Merge branch 'main' into instance_schema

* Update build_model_experimentation_lineage_with_prebuild_code.ipynb

* Update build_model_experimentation_lineage_with_prebuild_code.ipynb

* Update build_model_experimentation_lineage_with_prebuild_code.ipynb

* Update build_model_experimentation_lineage_with_prebuild_code.ipynb

* Update build_model_experimentation_lineage_with_prebuild_code.ipynb

Co-authored-by: Andrew Ferlitsch <aferlitsch@google.com>

.cloud-build/cleanup/cleanup-requirements.txt

@@ -1 +1,2 @@
+ratemate
 google-cloud-aiplatform

.cloud-build/cleanup/cleanup.py

@@ -1,45 +1,49 @@
 from typing import List
+from ratemate import RateLimit
 from resource_cleanup_manager import (
-	ResourceCleanupManager,
-	DatasetResourceCleanupManager,
-	EndpointResourceCleanupManager,
-	ModelResourceCleanupManager,
+    DatasetResourceCleanupManager,
+    ModelResourceCleanupManager,
+    EndpointResourceCleanupManager,
+    ResourceCleanupManager,
 )
 
+rate_limit = RateLimit(max_count=25, per=60, greedy=False)
+
 
 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:
+            try:
+                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:
+                    rate_limit.wait()  # wait before deleting
+                    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(),  # ModelResourceCleanupManager must follow EndpointResourceCleanupManager due to deployed models blocking model deletion.
 ]
 
 run_cleanup_managers(managers=managers, is_dry_run=is_dry_run)

.cloud-build/cleanup/resource_cleanup_manager.py

@@ -1,8 +1,9 @@
 import abc
+from typing import Any, Type
+
 from google.cloud import aiplatform
-from typing import Any
-from proto.datetime_helpers import DatetimeWithNanoseconds
 from google.cloud.aiplatform import base
+from proto.datetime_helpers import DatetimeWithNanoseconds
 
 # If a resource was updated within this number of seconds, do not delete.
 RESOURCE_UPDATE_BUFFER_IN_SECONDS = 60 * 60 * 8
@@ -40,7 +41,7 @@ class ResourceCleanupManager(abc.ABC):
         # Check that it wasn't created too recently, to prevent race conditions
         if time_difference <= RESOURCE_UPDATE_BUFFER_IN_SECONDS:
             print(
-                f"Skipping '{resource}' due update_time being '{time_difference}', which is less than '{RESOURCE_UPDATE_BUFFER_IN_SECONDS}'."
+                f"Skipping '{resource}' due to update_time being '{time_difference}', which is less than '{RESOURCE_UPDATE_BUFFER_IN_SECONDS}'."
             )
             return False
 
@@ -50,7 +51,7 @@ class ResourceCleanupManager(abc.ABC):
 class VertexAIResourceCleanupManager(ResourceCleanupManager):
     @property
     @abc.abstractmethod
-    def vertex_ai_resource(self) -> base.VertexAiResourceNounWithFutureManager:
+    def vertex_ai_resource(self) -> Type[base.VertexAiResourceNounWithFutureManager]:
         pass
 
     @property
@@ -60,7 +61,9 @@ class VertexAIResourceCleanupManager(ResourceCleanupManager):
     def list(self) -> Any:
         return self.vertex_ai_resource.list()
 
-    def resource_name(self, resource: Any) -> str:
+    def resource_name(
+        self, resource: Type[base.VertexAiResourceNounWithFutureManager]
+    ) -> str:
         return resource.display_name
 
     def delete(self, resource):
@@ -74,12 +77,33 @@ class VertexAIResourceCleanupManager(ResourceCleanupManager):
 
 class DatasetResourceCleanupManager(VertexAIResourceCleanupManager):
     vertex_ai_resource = aiplatform.datasets._Dataset
+    dataset_types = [
+        aiplatform.ImageDataset,
+        aiplatform.TabularDataset,
+        aiplatform.TextDataset,
+        aiplatform.TimeSeriesDataset,
+        aiplatform.VideoDataset,
+    ]
+
+    def list(self) -> Any:
+        return [
+            dataset
+            for dataset_type in self.dataset_types
+            for dataset in dataset_type.list()
+        ]
 
 
 class EndpointResourceCleanupManager(VertexAIResourceCleanupManager):
     vertex_ai_resource = aiplatform.Endpoint
 
     def delete(self, resource):
+        # TODO: Remove this once https://github.com/googleapis/python-aiplatform/issues/1441 is fixed
+        resource._sync_gca_resource()
+        for deployed_model_id in [
+            models.id for models in resource._gca_resource.deployed_models
+        ]:
+            resource._undeploy(deployed_model_id=deployed_model_id)
+
         resource.delete(force=True)
 
 

.cloud-build/execute_changed_notebooks_cli.py

@@ -17,6 +17,7 @@
 
 import argparse
 import pathlib
+
 import execute_changed_notebooks_helper
 
 
@@ -61,6 +62,18 @@ parser.add_argument(
     help="The GCP region. This is used to inject a variable value into the notebook before running.",
     required=True,
 )
+parser.add_argument(
+    "--variable_service_account",
+    type=str,
+    help="A service account. This is used to inject a variable value into the notebook before running. This is not the account that will run the notebook.",
+    required=True,
+)
+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(
     "--staging_bucket",
     type=str,
@@ -73,6 +86,13 @@ parser.add_argument(
     help="The GCP directory for storing executed notebooks.",
     required=True,
 )
+parser.add_argument(
+    "--timeout",
+    type=int,
+    help="Timeout in seconds",
+    default=86400,
+    required=False,
+)
 parser.add_argument(
     "--private_pool_id",
     type=str,
@@ -100,8 +120,11 @@ execute_changed_notebooks_helper.process_and_execute_notebooks(
     container_uri=args.container_uri,
     staging_bucket=args.staging_bucket,
     artifacts_bucket=args.artifacts_bucket,
+    should_parallelize=args.should_parallelize,
+    timeout=args.timeout,
     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,
+    variable_service_account=args.variable_service_account,
+    variable_vpc_network=args.variable_vpc_network,
+    private_pool_id=args.private_pool_id,
 )

.cloud-build/execute_changed_notebooks_helper.py

@@ -17,18 +17,28 @@ import concurrent
 import dataclasses
 import datetime
 import functools
+import json
+import git
+import operator
 import os
 import pathlib
-import nbformat
 import re
 import subprocess
+import utils
 from typing import List, Optional
-from tabulate import tabulate
-import operator
+from utils import util
 
+import execute_notebook_helper
 import execute_notebook_remote
-from utils import util, NotebookProcessors
+import nbformat
 from google.cloud.devtools.cloudbuild_v1.types import BuildOperationMetadata
+from ratemate import RateLimit
+from tabulate import tabulate
+from utils import NotebookProcessors, util
+
+# A buffer so that workers finish before the orchestrating job
+WORKER_TIMEOUT_BUFFER_IN_SECONDS: int = 60 * 60
+PYTHON_VERSION = "3.9" # Set default python version
 
 
 def format_timedelta(delta: datetime.timedelta) -> str:
@@ -60,13 +70,23 @@ class NotebookExecutionResult:
     log_url: str
     output_uri: str
     build_id: str
+    logs_bucket: str
     error_message: Optional[str]
 
+    @property
+    def output_uri_web(self) -> Optional[str]:
+        if self.output_uri.startswith("gs://"):
+            return f"https://storage.googleapis.com/{self.output_uri[5:]}"
+        else:
+            return None
+
 
 def _process_notebook(
     notebook_path: str,
     variable_project_id: str,
     variable_region: str,
+    variable_service_account: str,
+    variable_vpc_network: Optional[str],
 ):
     # Read notebook
     with open(notebook_path) as f:
@@ -78,6 +98,8 @@ def _process_notebook(
         replacement_map={
             "PROJECT_ID": variable_project_id,
             "REGION": variable_region,
+            "SERVICE_ACCOUNT": variable_service_account,
+            "VPC_NETWORK": variable_vpc_network,
         },
     )
 
@@ -93,6 +115,33 @@ def _process_notebook(
         nbformat.write(nb, new_file)
 
 
+def _get_notebook_python_version(notebook_path: str) -> str:
+    """
+    Get the python version for running the notebook if it is specified in
+    the notebook.
+    """
+    python_version = PYTHON_VERSION
+
+    # Load the notebook
+    file = open(notebook_path)
+    src = file.read()
+    nb_json = json.loads(src)
+
+    #Iterate over the cells in the ipynb
+    for cell in nb_json['cells']:
+        if cell['cell_type'] == 'markdown':
+            markdown = str.join('', cell['source'])
+
+            # Look for the python version specification pattern
+            re_match = re.search('python version = (\d\.\d)', markdown, flags=re.IGNORECASE)
+            if re_match:
+                # get the version number
+                python_version = re_match.group(1)
+                break
+
+    return python_version
+
+
 def _create_tag(filepath: str) -> str:
     tag = os.path.basename(os.path.normpath(filepath))
     tag = re.sub("[^0-9a-zA-Z_.-]+", "-", tag)
@@ -103,18 +152,33 @@ def _create_tag(filepath: str) -> str:
     return tag
 
 
+rate_limit = RateLimit(max_count=50, per=60, greedy=True)
+
+
 def process_and_execute_notebook(
     container_uri: str,
     staging_bucket: str,
     artifacts_bucket: str,
     variable_project_id: str,
     variable_region: str,
+    variable_service_account: str,
+    variable_vpc_network: Optional[str],
     private_pool_id: Optional[str],
+    deadline: datetime.datetime,
     notebook: str,
     should_get_tail_logs: bool = False,
 ) -> NotebookExecutionResult:
+    rate_limit.wait()  # wait before creating the task
+
     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])
 
@@ -128,6 +192,7 @@ def process_and_execute_notebook(
         output_uri=notebook_output_uri,
         log_url="",
         build_id="",
+        logs_bucket="",
         error_message=None,
     )
 
@@ -135,29 +200,43 @@ def process_and_execute_notebook(
     time_start = datetime.datetime.now()
     operation = None
     try:
+        # Get the python version for running the notebook if specified
+        notebook_exec_python_version = _get_notebook_python_version(notebook_path=notebook)
+        print(f"Running notebook with python {notebook_exec_python_version}")
+
         # Pre-process notebook by substituting variable names
         _process_notebook(
             notebook_path=notebook,
             variable_project_id=variable_project_id,
             variable_region=variable_region,
+            variable_service_account=variable_service_account,
+            variable_vpc_network=variable_vpc_network,
         )
 
         # Upload the pre-processed code to a GCS bucket
         code_archive_uri = util.archive_code_and_upload(staging_bucket=staging_bucket)
 
+        # Calculate timeout in seconds
+        timeout_in_seconds = max(
+            int((deadline - datetime.datetime.now()).total_seconds()), 1
+        )
+
         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,
-            region=variable_region,
             private_pool_id=private_pool_id,
+            private_pool_region=variable_region,
+            timeout_in_seconds=timeout_in_seconds,
+            python_version=notebook_exec_python_version
         )
 
         operation_metadata = BuildOperationMetadata(mapping=operation.metadata)
         result.build_id = operation_metadata.build.id
         result.log_url = operation_metadata.build.log_url
+        result.logs_bucket = operation_metadata.build.logs_bucket
 
         # Block and wait for the result
         operation_result = operation.result()
@@ -215,20 +294,40 @@ def get_changed_notebooks(
 
     # Find notebooks
     notebooks = []
+
+    # Instantiate GitPython objects
+    repo = git.Repo(os.getcwd())
+    index = repo.index
+
     if base_branch:
-        print(f"Looking for notebooks that changed from branch: {base_branch}")
-        notebooks = subprocess.check_output(
-            ["git", "diff", "--name-only", f"origin/{base_branch}..."] + test_paths
-        )
+        # Get the point at which this branch branches off from main
+        branching_commits = repo.merge_base("HEAD", f"origin/{base_branch}")
+
+        if len(branching_commits) > 0:
+            branching_commit = branching_commits[0]
+            print(f"Looking for notebooks that changed from branch: {branching_commit}")
+
+            notebooks = [
+                diff.b_path
+                for diff in index.diff(branching_commit, paths=test_paths)
+                if diff.b_path is not None
+            ]
+        else:
+            notebooks = []
     else:
         print(f"Looking for all notebooks.")
-        notebooks = subprocess.check_output(["git", "ls-files"] + test_paths)
+        notebooks_str = subprocess.check_output(["git", "ls-files"] + test_paths)
+        notebooks = notebooks_str.decode("utf-8").split("\n")
 
-    notebooks = notebooks.decode("utf-8").split("\n")
     notebooks = [notebook for notebook in notebooks if notebook.endswith(".ipynb")]
     notebooks = [notebook for notebook in notebooks if len(notebook) > 0]
     notebooks = [notebook for notebook in notebooks if pathlib.Path(notebook).exists()]
 
+    if len(notebooks) > 0:
+        print(f"Found {len(notebooks)} notebooks:")
+        for notebook in notebooks:
+            print(f"\t{notebook}")
+
     return notebooks
 
 
@@ -237,10 +336,13 @@ def process_and_execute_notebooks(
     container_uri: str,
     staging_bucket: str,
     artifacts_bucket: str,
+    should_parallelize: bool,
+    timeout: int,
     variable_project_id: str,
     variable_region: str,
-    private_pool_id: Optional[str],
-    should_parallelize: bool,
+    variable_service_account: str,
+    variable_vpc_network: Optional[str] = None,
+    private_pool_id: Optional[str] = None,
 ):
     """
     Run the notebooks that exist under the folders defined in the test_paths_file.
@@ -267,17 +369,27 @@ def process_and_execute_notebooks(
             Required. The value for REGION to inject into notebooks.
         should_parallelize (bool):
             Required. Should run notebooks in parallel using a thread pool as opposed to in sequence.
+        timeout (str):
+            Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
     """
-    notebook_execution_results: List[NotebookExecutionResult] = []
 
-    if len(notebooks) > 0:
+    # Calculate deadline
+    deadline = datetime.datetime.now() + datetime.timedelta(
+        seconds=max(timeout - WORKER_TIMEOUT_BUFFER_IN_SECONDS, 0)
+    )
+
+    if len(notebooks) >= 1:
+        notebook_execution_results: List[NotebookExecutionResult] = []
+
         print(f"Found {len(notebooks)} modified notebooks: {notebooks}")
 
         if should_parallelize and len(notebooks) > 1:
             print(
                 "Running notebooks in parallel, so no logs will be displayed. Please wait..."
             )
-            with concurrent.futures.ThreadPoolExecutor(max_workers=None) as executor:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
+                print(f"Max workers: {executor._max_workers}")
+
                 notebook_execution_results = list(
                     executor.map(
                         functools.partial(
@@ -287,7 +399,10 @@ def process_and_execute_notebooks(
                             artifacts_bucket,
                             variable_project_id,
                             variable_region,
+                            variable_service_account,
+                            variable_vpc_network,
                             private_pool_id,
+                            deadline,
                         ),
                         notebooks,
                     )
@@ -300,48 +415,88 @@ def process_and_execute_notebooks(
                     artifacts_bucket=artifacts_bucket,
                     variable_project_id=variable_project_id,
                     variable_region=variable_region,
+                    variable_service_account=variable_service_account,
+                    variable_vpc_network=variable_vpc_network,
                     private_pool_id=private_pool_id,
+                    deadline=deadline,
                     notebook=notebook,
                 )
                 for notebook in notebooks
             ]
+
+        print("\n=== RESULTS ===\n")
+
+        results_sorted = sorted(
+            notebook_execution_results,
+            key=lambda result: result.is_pass,
+            reverse=True,
+        )
+
+        # Print results
+        print(
+            tabulate(
+                [
+                    [
+                        result.name,
+                        "PASSED" if result.is_pass else "FAILED",
+                        format_timedelta(result.duration),
+                        result.log_url,
+                        result.output_uri,
+                        result.output_uri_web,
+                        result.logs_bucket
+                    ]
+                    for result in results_sorted
+                ],
+                headers=[
+                    "build_tag",
+                    "status",
+                    "duration",
+                    "log_url",
+                    "output_uri",
+                    "output_uri_web",
+                    "logs_bucket"
+                ],
+            )
+        )
+
+        if len(notebooks) == 1:
+          print("="*100)
+          print("The notebook execution build log:\n")
+          print("="*100)
+
+          build_id = results_sorted[0].build_id
+          logs_bucket_name = (results_sorted[0].logs_bucket).removeprefix("gs://")
+          log_file_name = f"log-{build_id}.txt"
+
+          log_contents = util.download_blob_into_memory(
+            bucket_name=logs_bucket_name,
+            blob_name=log_file_name,
+            download_as_text=True
+            )
+
+          # Remove extra steps from the log
+          match = re.search("starting Step #4", log_contents, flags=re.IGNORECASE)
+
+          if match is not None:
+            match_index = match.span()[0]
+            print(log_contents[match_index:])
+          else:
+            print(log_contents)
+
+        print("\n=== END RESULTS===\n")
+
+        total_notebook_duration = functools.reduce(
+            operator.add,
+            [datetime.timedelta(seconds=0)]
+            + [result.duration for result in results_sorted],
+        )
+
+        print(
+            f"Cumulative notebook duration: {format_timedelta(total_notebook_duration)}"
+        )
+
+        # 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")
     else:
         print("No notebooks modified in this pull request.")
-
-    print("\n=== RESULTS ===\n")
-
-    results_sorted = sorted(
-        notebook_execution_results,
-        key=lambda result: result.is_pass,
-        reverse=True,
-    )
-
-    # Print results
-    print(
-        tabulate(
-            [
-                [
-                    result.name,
-                    "PASSED" if result.is_pass else "FAILED",
-                    format_timedelta(result.duration),
-                    result.log_url,
-                ]
-                for result in results_sorted
-            ],
-            headers=["build_tag", "status", "duration", "log_url"],
-        )
-    )
-
-    print("\n=== END RESULTS===\n")
-
-    total_notebook_duration = functools.reduce(
-        operator.add,
-        [datetime.timedelta(seconds=0)]
-        + [result.duration for result in results_sorted],
-    )
-
-    print(f"Cumulative notebook duration: {format_timedelta(total_notebook_duration)}")
-
-    # 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")

.cloud-build/execute_notebook_cli.py

@@ -16,6 +16,7 @@
 """A CLI to download (optional) and run a single notebook locally"""
 
 import argparse
+
 import execute_notebook_helper
 
 parser = argparse.ArgumentParser(description="Run a single notebook locally.")

.cloud-build/execute_notebook_helper.py

@@ -15,17 +15,20 @@
 
 """Methods to run a notebook locally"""
 
-import sys
-import os
 import errno
-import papermill as pm
+import os
 import shutil
+import sys
 
-from utils import util
+import papermill as pm
 from google.cloud.aiplatform import utils
+from utils import util
 
 # This script is used to execute a notebook and write out the output notebook.
 
+# This is used to force papermill to use this kernel to run the notebook instead of any defined inside the notebook itself
+DEFAULT_KERNEL_NAME = "python3"
+
 
 def execute_notebook(
     notebook_source: str,
@@ -50,6 +53,17 @@ def execute_notebook(
 
     execution_exception = None
 
+    print("\n=== DOWNLOAD EXECUTED NOTEBOOK ===\n")
+    print(f"Please debug the executed notebook by downloading the executed notebook:")
+
+    print("Option 1. Using gsutil. Run the following command in your terminal.")
+    print(f'\tgsutil cp "{output_file_or_uri}" .')
+
+    print("Option 2. Using this link.")
+    print(f"\thttps://storage.googleapis.com/{output_file_or_uri[5:]}")
+
+    print("\n======\n")
+
     # Execute notebook
     try:
         # Execute notebook
@@ -58,6 +72,7 @@ def execute_notebook(
             output_path=notebook_source,
             progress_bar=should_log_output,
             request_save_on_cell_execute=should_log_output,
+            kernel_name=DEFAULT_KERNEL_NAME,
             log_output=should_log_output,
             stdout_file=sys.stdout if should_log_output else None,
             stderr_file=sys.stderr if should_log_output else None,
@@ -71,10 +86,6 @@ def execute_notebook(
             util.upload_file(notebook_source, remote_file_path=output_file_or_uri)
 
             print("\n=== EXECUTION FINISHED ===\n")
-            print(
-                f"Please debug the executed notebook by downloading: {output_file_or_uri}"
-            )
-            print("\n======\n")
         else:
             # Create directories if they don't exist
             if not os.path.exists(os.path.dirname(output_file_or_uri)):

.cloud-build/execute_notebook_remote.py

@@ -16,22 +16,18 @@
 """Methods to run a notebook on Google Cloud Build"""
 
 from re import sub
+from typing import Optional
+
+import google.auth
+import yaml
+from google.api_core import client_options, operation
+from google.cloud.aiplatform import utils
+from google.cloud.devtools import cloudbuild_v1
+from google.cloud.devtools.cloudbuild_v1.types import Source, StorageSource
 from google.protobuf import duration_pb2
 from yaml.loader import FullLoader
 
-import google.auth
-from google.cloud.devtools import cloudbuild_v1
-from google.cloud.devtools.cloudbuild_v1.types import Source, StorageSource
-
-from typing import Optional
-import yaml
-
-from google.cloud.aiplatform import utils
-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"
 
 
@@ -40,30 +36,38 @@ def execute_notebook_remote(
     notebook_uri: str,
     notebook_output_uri: str,
     container_uri: str,
-    region: str,
     private_pool_id: Optional[str],
+    private_pool_region: Optional[str],
     tag: Optional[str],
+    timeout_in_seconds: Optional[int] = None,
+    python_version: Optional[str] = None
 ) -> operation.Operation:
     """Create and execute a single notebook on Google Cloud Build"""
     # Load build steps from YAML
+
     cloudbuild_config = yaml.load(open(CLOUD_BUILD_FILEPATH), Loader=FullLoader)
 
     substitutions = {
         "_PYTHON_IMAGE": container_uri,
         "_NOTEBOOK_GCS_URI": notebook_uri,
         "_NOTEBOOK_OUTPUT_GCS_URI": notebook_output_uri,
+        "_PYTHON_VERSION" : f"python{python_version}"
     }
 
+    if python_version is not None:
+      substitutions["_PYTHON_VERSION"] = "python" + python_version
+
     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"]
+    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}
 
         # Switch to the regional endpoint of the pool
         options = client_options.ClientOptions(
-            api_endpoint=f"{region}-{SERVICE_BASE_PATH}"
+            api_endpoint=f"{private_pool_region}-{SERVICE_BASE_PATH}"
         )
 
     # Authorize the client with Google defaults
@@ -85,8 +89,8 @@ def execute_notebook_remote(
 
     build.steps = cloudbuild_config["steps"]
     build.substitutions = substitutions
-    build.timeout = duration_pb2.Duration(seconds=TIMEOUT_IN_SECONDS)
-    build.queue_ttl = duration_pb2.Duration(seconds=TIMEOUT_IN_SECONDS)
+    build.timeout = duration_pb2.Duration(seconds=timeout_in_seconds)
+    build.queue_ttl = duration_pb2.Duration(seconds=timeout_in_seconds)
 
     if tag:
         build.tags = [tag]

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

@@ -4,28 +4,35 @@ steps:
     entrypoint: /bin/sh
     args:
     - -c
-    - 'gcloud config list'
+    - 'gcloud config list --quiet'
   # Check the Python version
   - name: ${_PYTHON_IMAGE}
     entrypoint: /bin/sh
     args:
     - -c
-    - 'python3 .cloud-build/CheckPythonVersion.py'
+    - ${_PYTHON_VERSION} .cloud-build/CheckPythonVersion.py -q
+  # Create a virtual environment
+  - name: ${_PYTHON_IMAGE}
+    entrypoint: /bin/sh
+    args:
+    - -c
+    - ${_PYTHON_VERSION} -m venv workspace/env
   # Install Python dependencies
   - name: ${_PYTHON_IMAGE}
     entrypoint: /bin/sh
     args:
     - -c
-    - 'python3 -m pip install -U pip && python3 -m pip install -U --user -r .cloud-build/requirements.txt'
+    - . workspace/env/bin/activate &&
+      python -m pip -q install -U pip &&
+      python -m pip -q install -U -r .cloud-build/requirements.txt
   # Install Python dependencies and run testing script
   - name: ${_PYTHON_IMAGE}
     entrypoint: /bin/sh
     args:
     - -c
-    - 'python3 -m pip install -U pip && python3 -m pip freeze && python3 .cloud-build/execute_notebook_cli.py --notebook_source "${_NOTEBOOK_GCS_URI}" --output_file_or_uri "${_NOTEBOOK_OUTPUT_GCS_URI}"'
+    - |
+      . workspace/env/bin/activate &&
+      python .cloud-build/execute_notebook_cli.py --notebook_source "${_NOTEBOOK_GCS_URI}" --output_file_or_uri "${_NOTEBOOK_OUTPUT_GCS_URI}"
     env:
     - 'IS_TESTING=1'
 timeout: 86400s
-options:
-  pool:
-    name: ${_PRIVATE_POOL_NAME}

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

@@ -4,35 +4,42 @@ steps:
     entrypoint: /bin/sh
     args:
     - -c
-    - 'gcloud config list'
+    - gcloud config list --quiet
   # Check the Python version
   - name: ${_PYTHON_IMAGE}
     entrypoint: /bin/sh
     args:
     - -c
-    - 'python3 .cloud-build/CheckPythonVersion.py'
-  # Fetch base branch if required
+    - python3 .cloud-build/CheckPythonVersion.py -q
+  # Fetch full repo for diff purposes
+  - name: gcr.io/cloud-builders/git
+    args: [fetch, --unshallow, --quiet]
+  # Create a virtual environment
   - name: ${_PYTHON_IMAGE}
     entrypoint: /bin/sh
     args:
     - -c
-    - 'if [ -n "${_BASE_BRANCH}" ]; then git fetch origin "${_BASE_BRANCH}":refs/remotes/origin/"${_BASE_BRANCH}"; else echo "Skipping fetch."; fi'
+    - python3 -m venv workspace/env
   # Install Python dependencies
   - name: ${_PYTHON_IMAGE}
     entrypoint: /bin/sh
     args:
     - -c
-    - 'python3 -m pip install -U pip && python3 -m pip install -U --user -r .cloud-build/requirements.txt'
+    - . workspace/env/bin/activate &&
+      python3 -m pip -q install -U pip &&
+      python3 -m pip -q install -U -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} `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi`'
+    - |
+      . 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`
     env:
     - 'IS_TESTING=1'
 timeout: 86400s
 options:
   pool:
-    name: ${_PRIVATE_POOL_NAME}
+    name: ${_PRIVATE_POOL_NAME}

.cloud-build/requirements.txt

@@ -9,4 +9,5 @@ tabulate
 google-cloud-aiplatform
 google-cloud-storage
 google-cloud-build
-gcloud
+ratemate
+GitPython

.cloud-build/test_notebook_vm.txt

@@ -1,8 +1,6 @@
 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
+notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb
+notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb
+.cloud-build/tests/python_version_test.ipynb

.cloud-build/test_single_notebook_vm.txt

@@ -0,0 +1 @@
+notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb

.cloud-build/tests/python_version_test.ipynb

@@ -0,0 +1,61 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "57a3d44ed8a8"
+      },
+      "source": [
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.7\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c6516f90311b"
+      },
+      "outputs": [],
+      "source": [
+        "# test if the right python version is being used\n",
+        "import sys\n",
+        "\n",
+        "actual_python_version = f\"{sys.version_info.major}.{sys.version_info.minor}\"\n",
+        "print(f\"Runtime python version: {actual_python_version}\")\n",
+        "\n",
+        "assert actual_python_version == \"3.7\", \"Wrong python version!\""
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "python_version_test.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

.cloud-build/utils/NotebookProcessors.py

@@ -13,8 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nbconvert.preprocessors import Preprocessor
 from typing import Dict
+
+from nbconvert.preprocessors import Preprocessor
+
 from . import UpdateNotebookVariables as update_notebook_variables
 
 
@@ -60,4 +62,4 @@ class UpdateVariablesPreprocessor(Preprocessor):
 
             executable_cells.append(cell)
         notebook.cells = executable_cells
-        return notebook, resources
+        return notebook, resources

.cloud-build/utils/UpdateNotebookVariables.py

@@ -35,8 +35,8 @@ Variables in conditionals can also be replaced:
 
 def get_updated_value(content: str, variable_name: str, variable_value: str) -> str:
     return re.sub(
-        rf"({variable_name}.*?=.*?[\",\'])\[.+?\]([\",\'].*?)",
-        rf"\1{variable_value}\2",
+        rf"({variable_name}.*? = .*?[\",\'])\[.+?\]([\",\'].*?)",
+        rf"\g<1>{variable_value}\g<2>",
         content,
         flags=re.M,
     )
@@ -78,4 +78,27 @@ 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"}'
+
+
+def test_region_equal_equals_ignore():
+    # Tests that == is ignored
+    new_content = get_updated_value(
+        content='REGION == "[your-region]"  # @param {type:"string"}',
+        variable_name="REGION",
+        variable_value="us-central1",
+    )
+    assert new_content == 'REGION == "[your-region]"  # @param {type:"string"}'
+
+
+def test_service_account():
+    # Tests that == is ignored
+    new_content = get_updated_value(
+        content='SERVICE_ACCOUNT = "[your-service-account]"  # @param {type:"string"}',
+        variable_name="SERVICE_ACCOUNT",
+        variable_value="12345-compute@developer.gserviceaccount.com",
+    )
+    assert (
+        new_content
+        == 'SERVICE_ACCOUNT = "12345-compute@developer.gserviceaccount.com"  # @param {type:"string"}'
+    )

.cloud-build/utils/util.py

@@ -1,13 +1,13 @@
-from datetime import datetime
-from typing import Optional
-from google.cloud import storage
-from google.cloud.aiplatform import utils
-from google.auth import credentials as auth_credentials
 import os
-
 import subprocess
 import tarfile
 import uuid
+from datetime import datetime
+from typing import Optional, Union
+
+from google.auth import credentials as auth_credentials
+from google.cloud import storage
+from google.cloud.aiplatform import utils
 
 
 def download_file(bucket_name: str, blob_name: str, destination_file: str) -> str:
@@ -57,4 +57,35 @@ 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
+
+
+def download_blob_into_memory(
+  bucket_name: str,
+  blob_name: str,
+  download_as_text: Optional[bool]=False
+) -> Union[bytes, str]:
+    """
+    Downloads a blob into memory as byte or as text if
+    download_as_text is set to True.
+    """
+
+    storage_client = storage.Client()
+
+    bucket = storage_client.bucket(bucket_name)
+
+    # Construct a client side representation of a blob.
+    blob = bucket.blob(blob_name)
+
+    # Download the blob content
+    if download_as_text:
+      contents = blob.download_as_text()
+    else:
+      contents = blob.download_as_bytes()
+
+    print(
+        f"Downloaded storage object {blob_name} from bucket {bucket_name}."
+    )
+
+    return contents
+

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,18 +1,28 @@
-If you are opening a PR for `Official Notebooks` under the [notebooks/official](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/notebooks/official) folder, follow this mandatory checklist:
+**REQUIRED:** Add a summary of your PR here, typically including why the change is needed and what was changed. Include any design alternatives for discussion purposes.
+
+<br>
+--- YOUR PR SUMMARY GOES HERE ---
+<br><br><br>
+
+**REQUIRED:** Fill out the below checklists or remove if irrelevant
+1. If you are opening a PR for `Official Notebooks` under the [notebooks/official](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/notebooks/official) folder, follow this mandatory checklist:
 - [ ] Use the [notebook template](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/notebook_template.ipynb) as a starting point.
 - [ ] Follow the style and grammar rules outlined in the above notebook template.
 - [ ] Verify the notebook runs successfully in Colab since the automated tests cannot guarantee this even when it passes.
-- [ ] Passes all the required automated checks. You can locally test for formatting and linting with these [instructions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/docs/contributing.md#code-quality-checks).
+- [ ] Passes all the required automated checks. You can locally test for formatting and linting with these [instructions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/CONTRIBUTING.md#code-quality-checks).
 - [ ] You have consulted with a tech writer to see if tech writer review is necessary. If so, the notebook has been reviewed by a tech writer, and they have approved it.
-- [ ] This notebook has been added to the [CODEOWNERS](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/docs/CODEOWNERS) file under `# Official Notebooks` section, pointing to the author or the author's team.
+- [ ] This notebook has been added to the [CODEOWNERS](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/CODEOWNERS) file under the `Official Notebooks` section, pointing to the author or the author's team.
 - [ ] The Jupyter notebook cleans up any artifacts it has created (datasets, ML models, endpoints, etc) so as not to eat up unnecessary resources.
 
+<br>
 
-If you are opening a PR for `Community Notebooks` under the [notebooks/community](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/notebooks/community) folder:
-- [ ] This notebook has been added to the [CODEOWNERS](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/docs/CODEOWNERS) file under the `# Community Notebooks` section, pointing to the author or the author's team.
-- [ ] Passes all the required formatting and linting checks. You can locally test with these [instructions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/docs/contributing.md#code-quality-checks).
+2. If you are opening a PR for `Community Notebooks` under the [notebooks/community](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/notebooks/community) folder:
+- [ ] This notebook has been added to the [CODEOWNERS](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/CODEOWNERS) file under the `Community Notebooks` section, pointing to the author or the author's team.
+- [ ] Passes all the required formatting and linting checks. You can locally test with these [instructions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/CONTRIBUTING.md#code-quality-checks).
 
-If you are opening a PR for `Community Content` under the [community-content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/community-content) folder:
+<br>
+
+3. If you are opening a PR for `Community Content` under the [community-content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/community-content) folder:
 - [ ] Make sure your main `Content Directory Name` is descriptive, informative, and includes some of the key products and attributes of your content, so that it is differentiable from other content
-- [ ] The main content directory has been added to the [CODEOWNERS](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/docs/CODEOWNERS) file under the `# Community Content` section, pointing to the author or the author's team.
-- [ ] Passes all the required formatting and linting checks. You can locally test with these [instructions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/docs/contributing.md#code-quality-checks).
+- [ ] The main content directory has been added to the [CODEOWNERS](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/community-content/CODEOWNERS) file under the `Community Content` section, pointing to the author or the author's team.
+- [ ] Passes all the required formatting and linting checks. You can locally test with these [instructions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/CONTRIBUTING.md#code-quality-checks).

.github/workflows/ci.yaml

@@ -7,9 +7,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.x'      
     - name: Fetch pull request branch
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
       with:
         fetch-depth: 0      
     - name: Fetch base main branch

.github/workflows/linter/Dockerfile

@@ -0,0 +1,20 @@
+# To use this image, run this command with the desired notebook args from the top-level vertex-ai-samples directory:
+# 1. To lint all changed notebooks:
+# docker run -v ${PWD}:/setup/app gcr.io/python-docs-samples-tests/notebook_linter:latest
+# 2. To lint specific notebooks:
+# docker run -v ${PWD}:/setup/app gcr.io/python-docs-samples-tests/notebook_linter:latest notebooks/1.ipynb notebooks/2.ipynb
+
+FROM python:3.10
+
+WORKDIR setup
+
+COPY ./requirements.txt .
+COPY ./run_linter.sh .
+
+# Install dependencies.
+RUN pip install --upgrade pip
+RUN pip install -r requirements.txt
+
+WORKDIR app
+
+ENTRYPOINT ["/setup/run_linter.sh"]

.github/workflows/linter/requirements.txt

@@ -2,8 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==22.1.0
-pyupgrade==2.31.0
+black==22.6.0
+pyupgrade==2.34.0
 isort==5.10.1
 flake8==4.0.1
-nbqa==1.2.3
+nbqa==1.4.0
+

.github/workflows/linter/run_linter.sh

@@ -47,12 +47,22 @@ done
 
 echo "Test mode: $is_test"
 
+# Read in user-provided notebooks
+notebooks=()
+for arg in "$@"; do
+    if [[ $arg == *.ipynb ]]; then
+        notebooks+=("$arg")
+    fi
+done
+
 # Only check notebooks in test folders modified in this pull request.
 # Note: Use process substitution to persist the data in the array
-notebooks=()
-while read -r file || [ -n "$line" ]; do
-    notebooks+=("$file")
-done < <(git diff --name-only main... | grep '\.ipynb$')
+if [ ${#notebooks[@]} -eq 0 ]; then
+    echo "Checking for changed notebooked using git"
+    while read -r file || [ -n "$line" ]; do
+        notebooks+=("$file")
+    done < <(git diff --name-only main... | grep '\.ipynb$')
+fi
 
 problematic_notebooks=()
 if [ ${#notebooks[@]} -gt 0 ]; then
@@ -68,7 +78,7 @@ if [ ${#notebooks[@]} -gt 0 ]; then
 
             if [ "$is_test" = true ]; then
                 echo "Running nbfmt..."
-                python3 -m tensorflow_docs.tools.nbfmt --remove_outputs --test "$notebook"
+                python3 -m tensorflow_docs.tools.nbfmt --test "$notebook"
                 NBFMT_RTN=$?
                 # echo "Running black..."
                 # python3 -m nbqa black "$notebook" --check
@@ -93,7 +103,7 @@ if [ ${#notebooks[@]} -gt 0 ]; then
                 python3 -m nbqa isort "$notebook"
                 ISORT_RTN=$?
                 echo "Running nbfmt..."
-                python3 -m tensorflow_docs.tools.nbfmt --remove_outputs "$notebook"
+                python3 -m tensorflow_docs.tools.nbfmt "$notebook"
                 NBFMT_RTN=$?
                 echo "Running flake8..."
                 python3 -m nbqa flake8 "$notebook" --show-source --extend-ignore=W391,E501,F821,E402,F404,W503,E203,E722,W293,W291

CONTRIBUTING.md

@@ -48,8 +48,8 @@ then you will need to manually address them before submitting your PR.
 nbqa black "$notebook"
 nbqa pyupgrade "$notebook"
 nbqa isort "$notebook"
-python3 -m tensorflow_docs.tools.nbfmt --remove_outputs "$notebook"
 nbqa flake8 "$notebook" --extend-ignore=W391,E501,F821,E402,F404,W503,E203,E722,W293,W291
+python3 -m tensorflow_docs.tools.nbfmt --remove_outputs "$notebook"
 ```
 
 ## Code Reviews

README.md

@@ -6,7 +6,19 @@ Welcome to the Google Cloud [Vertex AI](https://cloud.google.com/vertex-ai/docs/
 
 ## Overview
 
-The repository contains [Notebooks](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/notebooks) and [Community Content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content) that demonstrate how to develop and manage ML workflows using Google Cloud Vertex AI.
+The repository contains [notebooks](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/notebooks) and [community content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content) that demonstrate how to develop and manage ML workflows using Google Cloud Vertex AI.
+
+## Repository structure
+
+```bash
+├── community-content - Sample code and tutorials contributed by the community
+├── notebooks
+│   ├── community - Notebooks contributed by the community
+│   ├── official - Notebooks demonstrating use of each Vertex AI service
+│   │   ├── automl
+│   │   ├── custom
+│   │   ├── ...
+```
 
 ## Contributing
 
@@ -19,3 +31,7 @@ Please use the [issues page](https://github.com/GoogleCloudPlatform/vertex-ai-sa
 ## Disclaimer
 
 This is not an officially supported Google product. The code in this repository is for demonstrative purposes only.
+
+## Feedback
+
+Please feel free to fill out our [survey](https://bit.ly/vertex-ai-samples-survey) to give us feedback on the repo and its content.

community-content/CODEOWNERS

@@ -1,5 +1,8 @@
 *  @vertex-ai-samples-contributors  @GoogleCloudPlatform/cloudml-samples-owners
 /tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk	@yinghsienwu
+/pytorch_pre_built_images_deployment @googleapis/vertex-prediction-team
 /pytorch_text_classification_using_vertex_sdk_and_gcloud  @RajeshThallam
 /pytorch_text_classification_using_vertex_sdk_and_gcloud		@RajeshThallam @ultrons
 /sklearn_text_classification_from_script_using_vertex_sdk       @maxhardt
+/pluto_on_workbench @wkharold
+/cpr-examples @samthrasher

community-content/cpr-examples/timm_serving/.dockerignore

@@ -0,0 +1,5 @@
+testdata/*
+build.py
+test.py
+state_dict.pth
+config.json

community-content/cpr-examples/timm_serving/.gitignore

@@ -0,0 +1,6 @@
+cpr_model_server.py
+entrypoint.py
+state_dict.pth
+config.json
+**/__pycache__
+!testdata/**

community-content/cpr-examples/timm_serving/README.md

@@ -0,0 +1,110 @@
+# CPR Example: PyTorch Image Models (timm)
+
+## About CPR
+
+CPR ([custom prediction routines](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/prediction/README.md)) is a framework designed by Google Cloud developers to make it easier to combine machine learning models with custom preprocessing and postprocessing logic in a real-time serving application. 
+
+## Using this example
+
+This code is a self-contained example of a custom model server project built using CPR.
+
+As is, you can use it to serve the ViT-Small image classification model from Ross Wightman's [`timm`](https://github.com/rwightman/pytorch-image-models) library of image model implementations in PyTorch. Both CPU and GPU are supported.
+
+You can also consider using the code here as a template for your own CPR project if you want to use a different model from `timm`, a different PyTorch model, or an entirely different framework.
+
+### Requirements
+
+In order to use this example, you'll need Docker and Python 3 installed on your system.
+
+To get started, first create a virtual environment in an empty directory:
+```sh
+mkdir cpr-example
+python3 -m venv cpr-example
+cd cpr-example && source bin/activate
+``` 
+
+Then, clone the [vertex-ai-samples repo](https://github.com/GoogleCloudPlatform/vertex-ai-samples) in that directory:
+```sh
+git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
+cd vertex-ai-samples/community-content/cpr-examples/timm_serving
+```
+
+Finally, install the Python modules required to build and run the model server:
+```sh
+pip install -r requirements.txt
+```
+
+### Auth
+
+This example uses Google Cloud Storage for hosting model artifacts and Artifact Registry to store the container image. 
+You'll need to authorize yourself before you can interact with these.
+
+First, log in to GCP with application default credentials:
+```sh
+gcloud auth application-default login
+```
+
+Next, if you haven't done so already, set up the [gcloud credential helper](https://cloud.google.com/artifact-registry/docs/docker/authentication)
+for the Artifact Registry region where you intend to host the image.  
+```
+gcloud auth configure-docker <region>-docker.pkg.dev
+```
+
+
+### Predictor
+
+The `TimmPredictor` class in `timm_serving/predictor.py` implements most of the important logic for the server.
+
+- `load(artifacts_dir)`: The predictor's `load` method is called when the server starts up in order to set up the predictor, usually by loading model weights and any artifacts needed for preprocessing and postprocessing. In this example, we initialize the saved model from the `state_dict.pth` file located inside the `artifacts_dir` folder and create the preprocessing transform from the model config.
+
+- `preprocess`, `predict`, `postprocess`: These methods are applied in sequence to the deserialized JSON data from each request. 
+    - `preprocess` decodes images from base64 and apply cropping, scaling and normalizing transforms. 
+    - `predict` runs the ViT-Small model on the preprocessed images and returns class scores.
+    - `postprocess` finds the top five classes and packs the class names, probabilities, and indices in a serializable result.
+
+### Building the container
+
+To build the model server locally, run the build command:
+```sh
+python build.py build
+```
+
+You can edit configuration values such as the model server's base image, the name and tag assigned to the image, and the path where model weights are stored locally.
+
+When you run the build command, model weights are downloaded and the model server container is built.
+
+### Running local tests
+
+`test.py` contains a suite of unit tests for the predictor as well as end-to-end tests for the model server. 
+
+To run the tests:
+```sh
+python test.py
+```
+
+All of the test images are public domain.
+- [Cat](https://commons.wikimedia.org/wiki/File:Stray_cat_on_wall.jpg)
+- [Airplane](https://commons.wikimedia.org/wiki/File:Airplanes_jets.jpg)
+- The infamous [mandrill](https://commons.wikimedia.org/wiki/File:Wikipedia-sipi-image-db-mandrill-4.2.03.png)
+
+### Deploying to Vertex AI
+
+Before uploading or deploying the container, you'll need to modify `config.py` to set appropriate values for: 
+- `project_id`: Your GCP project id.
+- `region`: Region where the model will be uploaded and deployed.
+- `repository`: [Artifact Registry repository](https://cloud.google.com/artifact-registry/docs/repositories/create-repos) in your project where the container image will be uploaded. 
+- `artifacts_gcs_dir`: Folder in a [Google Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) where the model weights will be uploaded.
+
+Once this is done, first upload the model:
+```sh
+python build.py upload
+```
+
+Then deploy it:
+```sh
+python build.py deploy
+```
+
+If you run the deploy command again, it will create a new endpoint. If you want to undeploy the model, you can do so using the Vertex AI dashboard on the Google Cloud console, or use `gcloud ai endpoints undeploy` from the command line.
+
+After deploying successfully, you can run `python build.py probe` to send a sample request to the deployed model. 

community-content/cpr-examples/timm_serving/build.py

@@ -0,0 +1,117 @@
+# 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.
+# You may obtain a copy of the License at
+
+#     https://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.
+
+"""Build the model server container."""
+import json
+import logging
+import os
+import pathlib
+from typing import Sequence
+
+from absl import app
+from absl import logging
+from config import CPRConfig
+from google.cloud import aiplatform
+from google.cloud.aiplatform import prediction as cpr
+import smart_open
+import timm
+from timm_serving import predictor
+import torch
+
+
+def build_container(config: CPRConfig, tag: str) -> cpr.LocalModel:
+    """Build the model server container.
+
+    Args:
+      tag: Output image tag.
+
+    Returns:
+      LocalModel exposing the built model server.
+    """
+    return cpr.LocalModel.build_cpr_model(
+        src_dir=os.path.join(os.getcwd()),
+        output_image_uri=tag,
+        base_image=config.base_image,
+        predictor=predictor.TimmPredictor,
+        requirements_path=os.path.join(os.getcwd(), "requirements.txt"),
+    )
+
+
+def save_model_artifact(destination: str) -> None:
+    """Save a copy of the model state dict."""
+    model = timm.create_model(predictor.TimmPredictor.TIMM_MODEL_NAME, pretrained=True)
+    dest_file = os.path.join(destination, predictor.TimmPredictor.WEIGHTS_FILE)
+    with smart_open.open(dest_file, "wb") as f:
+        torch.save(model, f)
+    logging.info("Saved model to %s", dest_file)
+    logging.info("%s parameters", sum(p.numel() for p in model.parameters()))
+
+
+def upload_model(config: CPRConfig) -> aiplatform.Model:
+    """Tag and upload the model server."""
+    ar_tag = (
+        f"{config.region}-docker.pkg.dev/{config.project_id}"
+        f"/{config.repository}/{config.image}"
+    )
+    local_model = build_container(config, tag=ar_tag)
+    aiplatform.init(project=config.project_id, location=config.region)
+    local_model.push_image()
+    aip_model = aiplatform.Model.upload(
+        local_model=local_model,
+        display_name=predictor.TimmPredictor.TIMM_MODEL_NAME,
+        artifact_uri=config.artifact_gcs_dir,
+    )
+    config.model_name = aip_model.resource_name
+    config.save()
+    return aip_model
+
+
+def deploy_model(config: CPRConfig) -> aiplatform.Endpoint:
+    """Deploy the model server to a Vertex Prediction endpoint."""
+    aiplatform.init(project=config.project_id, location=config.region)
+    aip_model = aiplatform.Model(model_name=config.model_name)
+    endpoint = aip_model.deploy(machine_type=config.machine_type)
+    config.endpoint_name = endpoint.resource_name
+    config.save()
+    return endpoint
+
+
+def probe_prediction(config: CPRConfig, request_path: str) -> None:
+    """Send a sample prediction request to the Vertex Prediction endpoint."""
+    aiplatform.init(project=config.project_id, location=config.region)
+    aip_endpoint = aiplatform.Endpoint(endpoint_name=config.endpoint_name)
+    with open(request_path) as f:
+        logging.info(aip_endpoint.predict(**json.load(f)))
+
+
+def main(argv: Sequence[str]):
+    config = CPRConfig()
+    if pathlib.Path(config.config_file).exists():
+        config.load()
+
+    actions = set(argv[1:])
+    if "build" in actions:
+        build_container(config, config.image)
+        save_model_artifact(config.artifact_local_dir)
+    if "upload" in actions:
+        save_model_artifact(config.artifact_gcs_dir)
+        upload_model(config)
+    if "deploy" in actions:
+        deploy_model(config)
+    if "probe" in actions:
+        probe_prediction(config, request_path="sample_request.json")
+
+
+if __name__ == "__main__":
+    app.run(main)

community-content/cpr-examples/timm_serving/config.py

@@ -0,0 +1,76 @@
+# 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.
+# You may obtain a copy of the License at
+
+#     https://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.
+
+import dataclasses
+import json
+
+
+@dataclasses.dataclass
+class CPRConfig(object):
+    """Configure the build process by editing the default values here.
+
+    config_file: File path used to save values in this config. (Some
+      values, such as the model name, are generated at build time and
+      depended on by future steps, so saving it allows this script to
+      deploy the model without re-uploading it, for example.)
+
+    base_image: Base Docker image on top of which the model server will
+      be built. By default, a Debian-based Python 3 image without GPU
+      support will be used.
+
+    image: Name and tag assigned to the built model server image.
+
+    artifact_local_dir: Local directory where a copy of the pretrained model weights
+      will be saved.
+
+    region: Google Cloud Region where the model will be uploaded during the
+      build process.
+
+    project_id: Google Cloud project ID.
+
+    repository: Name of the Artifact Registry repository where the container
+      will be uploaded.
+
+    artifact_gcs_dir: Location on GCS where a copy of the pretrained model
+      weights will be uploaded.
+
+    model_name: Full resource path of the uploaded model. This is a write-only
+      field, the value is generated by Vertex AI when the model is uploaded.
+
+    endpoint_name: Full resource path of the created endpoint. This is a
+      write-only field, the value is generated by Vertex AI when the model is
+      deployed to an endpoint.
+
+    machine_type: Machine type to use when deploying the model.
+    """
+
+    config_file: str = "config.json"
+    base_image: str = "python:3.10-bullseye"
+    image: str = "timm_predictor:latest"
+    artifact_local_dir: str = ""
+    region: str = "us-central1"
+    project_id: str = "<your project ID here>"
+    repository: str = "cpr-images"
+    artifact_gcs_dir: str = "gs://<your bucket ID here>/timm-vit224/"
+    model_name: str = ""
+    endpoint_name: str = ""
+    machine_type: str = "n1-standard-2"
+
+    def save(self):
+        with open(self.config_file, "w") as f:
+            json.dump(dataclasses.asdict(self), f, indent=2)
+
+    def load(self):
+        with open(self.config_file) as f:
+            self.__init__(**json.load(f))

community-content/cpr-examples/timm_serving/requirements.txt

@@ -0,0 +1,8 @@
+absl-py==1.1.0
+fastapi==0.75.2
+uvicorn==0.18.2
+timm==0.5.4
+smart_open==6.0.0
+
+google-cloud-storage>=1.26.0,<2.0.0dev
+google-cloud-aiplatform[prediction]>=1.16.0

community-content/cpr-examples/timm_serving/test.py

@@ -0,0 +1,255 @@
+"""Test the timm_serving predictor."""
+import base64
+import json
+import logging
+import os
+import pickle
+from typing import List, Dict
+
+from absl import flags
+from absl import logging
+from absl.testing import absltest
+from config import CPRConfig
+import fastapi
+from google.cloud import aiplatform
+from google.cloud.aiplatform import prediction as cpr
+import PIL
+from timm_serving import predictor
+import torch
+
+VIT_SMALL_PARAMS = 22878952
+
+
+def b64_encode_file(path: str) -> str:
+    """Encode a file's contents as base64.
+
+    Args:
+      path: Path to the file.
+
+    Returns:
+      Base64-encoded contents of the file.
+    """
+    with open(path, "rb") as f:
+        return str(base64.b64encode(f.read()), encoding="utf-8")
+
+
+def make_instance_dict(
+    image_paths: List[str], base64_encodings: List[str]
+) -> Dict[str, List[str]]:
+    """Generate a dictionary similar to a parsed prediction server request.
+
+    Args:
+      image_paths: Paths to image files to include.
+      base64_encodings: Pre-encoded base64 strings.
+
+    Returns:
+      Dictionary of instances in the format accepted by the preprocessor.
+    """
+    instances = [s for s in base64_encodings]
+    for path in image_paths:
+        instances.append(b64_encode_file(path))
+    return {"instances": instances}
+
+
+def count_parameters(model: torch.nn.Module):
+    """Count the parameters in a Pytorch model.
+
+    Args:
+      model: Pytorch model (nn.Module).
+
+    Returns:
+      Number of parameters in the model.
+
+    """
+    return sum(p.numel() for p in model.parameters())
+
+
+class PredictorUnitTests(absltest.TestCase):
+    """Unit tests for timm_serving.predictor."""
+
+    def setUp(self):
+        super().setUp()
+        self.config = CPRConfig()
+        try:
+            self.config.load()
+        except FileNotFoundError:
+            logging.info("No saved config file found, using default values.")
+        self.predictor = predictor.TimmPredictor()
+
+    def test_load_from_saved_state_dict_ok(self):
+        self.predictor.load(self.config.artifact_local_dir)
+        self.assertEqual(count_parameters(self.predictor._model), VIT_SMALL_PARAMS)
+
+    def test_load_bad_path(self):
+        with self.assertRaises(FileNotFoundError):
+            self.predictor.load("testdata/")
+        with self.assertRaisesRegex(ValueError, "not a directory"):
+            self.predictor.load("blah")
+
+    def test_load_bad_data(self):
+        with self.assertRaises(pickle.UnpicklingError):
+            self.predictor.load("testdata/bad_model_1")
+        with self.assertRaisesRegex(RuntimeError, "Invalid magic number"):
+            self.predictor.load("testdata/bad_model_2")
+
+    def test_preprocess_ok(self):
+        self.predictor.load(self.config.artifact_local_dir)
+        instance_dict = make_instance_dict(
+            base64_encodings=[],
+            image_paths=[
+                "testdata/airplane.jpg",
+                "testdata/mandrill.tiff",
+                "testdata/mandrill.tiff",
+                "testdata/cat_alpha.png",
+            ],
+        )
+        result = self.predictor.preprocess(instance_dict)
+        self.assertEqual(result.size(), torch.Size([4, 3, 224, 224]))
+        self.assertEqual(result.dtype, torch.float32)
+
+    def test_preprocess_no_instances(self):
+        self.predictor.load(self.config.artifact_local_dir)
+        with self.assertRaises(fastapi.HTTPException) as ctx:
+            self.predictor.preprocess({})
+        self.assertEqual(ctx.exception.status_code, 400)
+        self.assertRegex(ctx.exception.detail, 'must contain "instances"')
+
+    def test_preprocess_wrong_shape_instances(self):
+        self.predictor.load(self.config.artifact_local_dir)
+        instance_dict = {"instances": [[b64_encode_file("testdata/mandrill.tiff")]]}
+        with self.assertRaises(fastapi.HTTPException) as ctx:
+            self.predictor.preprocess(instance_dict)
+        self.assertEqual(ctx.exception.status_code, 400)
+        self.assertRegex(ctx.exception.detail, "not 'list'")
+
+    def test_preprocess_bad_base64(self):
+        self.predictor.load(self.config.artifact_local_dir)
+        instance_dict = make_instance_dict(base64_encodings=["!@#$"], image_paths=[])
+        with self.assertRaises(fastapi.HTTPException) as ctx:
+            self.predictor.preprocess(instance_dict)
+        self.assertEqual(ctx.exception.status_code, 400)
+        self.assertRegex(ctx.exception.detail, "[Bb]ase64")
+
+    def test_preprocess_not_image_data(self):
+        self.predictor.load(self.config.artifact_local_dir)
+        instance_dict = make_instance_dict(
+            base64_encodings=[], image_paths=["testdata/bad.jpg"]
+        )
+        with self.assertRaises(fastapi.HTTPException) as ctx:
+            self.predictor.preprocess(instance_dict)
+        self.assertEqual(ctx.exception.status_code, 400)
+        self.assertRegex(ctx.exception.detail, "image file")
+
+    def test_predict_ok(self):
+        self.predictor.load(self.config.artifact_local_dir)
+        inputs = torch.zeros(size=[2, 3, 224, 224], dtype=torch.float32)
+        if torch.cuda.device_count() > 0:
+            inputs = inputs.cuda()
+        result = self.predictor.predict(inputs)
+        self.assertEqual(result.size(), torch.Size([2, 1000]))
+        self.assertEqual(result.dtype, torch.float32)
+
+    def test_postprocess_ok(self):
+        class_probs = torch.zeros(size=[2, 1000])
+        class_probs[0, 0] = 1
+        class_probs[1, 123] = 1
+        result = self.predictor.postprocess(class_probs)
+        predictions = result["predictions"]
+        self.assertLen(predictions[0]["class_names"], 5)
+        self.assertLen(predictions[0]["indices"], 5)
+        self.assertLen(predictions[0]["probabilities"], 5)
+        self.assertLen(predictions[1]["class_names"], 5)
+        self.assertLen(predictions[1]["indices"], 5)
+        self.assertLen(predictions[1]["probabilities"], 5)
+        self.assertContainsSubsequence(predictions[0]["class_names"][0], "tench")
+        self.assertContainsSubsequence(
+            predictions[1]["class_names"][0], "spiny lobster"
+        )
+
+
+class ServerEndToEndTests(absltest.TestCase):
+    """End-to-end tests for the model server, using LocalEndpoint."""
+
+    def setUp(self):
+        super().setUp()
+        self.config = CPRConfig()
+        try:
+            self.config.load()
+        except FileNotFoundError:
+            logging.info("No saved config file found, using default values.")
+        self.local_model = cpr.LocalModel(
+            serving_container_spec=aiplatform.gapic.ModelContainerSpec(
+                image_uri=self.config.image
+            )
+        )
+
+        self.local_endpoint = self.local_model.deploy_to_local_endpoint(
+            artifact_uri=self.config.artifact_local_dir or os.getcwd()
+        )
+        self.local_endpoint.serve()
+
+    def tearDown(self):
+        self.local_endpoint.stop()
+        super().tearDown()
+
+    def test_e2e_healthcheck_ok(self):
+        health_check_response = self.local_endpoint.run_health_check()
+        self.assertEqual(health_check_response.status_code, 200)
+        self.assertEqual(health_check_response.content, b"{}")
+
+    def test_e2e_predict_ok(self):
+        predict_request = json.dumps(
+            make_instance_dict(
+                base64_encodings=[],
+                image_paths=[
+                    "testdata/mandrill.tiff",
+                ],
+            )
+        )
+        response = self.local_endpoint.predict(
+            request=predict_request, headers={"Content-Type": "application/json"}
+        )
+        logging.info(response.content)
+        self.assertEqual(response.status_code, 200)
+        predictions = response.json()["predictions"]
+        self.assertContainsSubsequence(predictions[0]["class_names"][0], "baboon")
+
+    def test_e2e_predict_bad_json_returns_400(self):
+        predict_request = "blah"
+        response = self.local_endpoint.predict(
+            request=predict_request, headers={"Content-Type": "application/json"}
+        )
+        logging.info(response.content)
+        self.assertEqual(response.status_code, 400)
+
+    def test_e2e_predict_no_instances_returns_400(self):
+        predict_request = json.dumps({})
+        response = self.local_endpoint.predict(
+            request=predict_request, headers={"Content-Type": "application/json"}
+        )
+        logging.info(response.content)
+        self.assertEqual(response.status_code, 400)
+
+    def test_e2e_predict_bad_base64_returns_400(self):
+        predict_request = json.dumps(
+            make_instance_dict(base64_encodings=["blah"], image_paths=[])
+        )
+        response = self.local_endpoint.predict(
+            request=predict_request, headers={"Content-Type": "application/json"}
+        )
+        logging.info(response.content)
+        self.assertEqual(response.status_code, 400)
+
+    def test_e2e_predict_bad_image_returns_400(self):
+        predict_request = json.dumps(
+            make_instance_dict(base64_encodings=[], image_paths=["testdata/bad.jpg"])
+        )
+        response = self.local_endpoint.predict(
+            request=predict_request, headers={"Content-Type": "application/json"}
+        )
+        logging.info(response.content)
+        self.assertEqual(response.status_code, 400)
+
+
+if __name__ == "__main__":
+    absltest.main()

community-content/cpr-examples/timm_serving/testdata/bad.jpg

@@ -0,0 +1 @@
+some non-image data

community-content/cpr-examples/timm_serving/testdata/bad.txt

@@ -0,0 +1 @@
+some non-image data

community-content/cpr-examples/timm_serving/testdata/bad_model_1/state_dict.pth

@@ -0,0 +1 @@
+blah

community-content/cpr-examples/timm_serving/timm_serving/predictor.py

@@ -0,0 +1,178 @@
+# 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.
+# You may obtain a copy of the License at
+
+#     https://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.
+
+"""Adapts a pretrained TIMM image classification model to the CPR framework.
+
+Documentation for the TIMM (Torch IMage Models) library is here:
+  https://rwightman.github.io/pytorch-image-models/
+
+Its source can also be found here:
+  https://github.com/rwightman/pytorch-image-models
+"""
+
+import base64
+import binascii
+import io
+import os
+from typing import Dict, List, Union
+
+from fastapi import HTTPException
+from google.cloud.aiplatform import prediction as cpr
+from pathlib import Path
+import PIL
+import smart_open
+import timm
+import torch
+import torch.nn.functional as F
+
+with open(Path(__file__).parent.absolute().joinpath("imagenet.txt")) as f:
+    IMAGENET_CLASSES = f.read().splitlines()
+
+
+class TimmPredictor(cpr.predictor.Predictor):
+    """Predictor class for image models based on TIMM."""
+
+    TIMM_MODEL_NAME = os.getenv("TIMM_MODEL_NAME", default="vit_small_patch32_224")
+    WEIGHTS_FILE = "state_dict.pth"
+    NUM_TOP_CLASSES_TO_RETURN = 5
+
+    def __init__(self):
+        self._cuda = torch.cuda.device_count() > 0
+
+    def load(self, artifacts_uri: str = ""):
+        """Initializes the model and preprocessing transforms.
+
+        Args:
+          artifacts_uri: Directory where state dict is stored. Can be a
+          GCS URI or local path.
+        """
+        if artifacts_uri:
+            artifact_path = os.path.join(artifacts_uri)
+            if not (os.path.isdir(artifact_path) or artifact_path.startswith("gs://")):
+                raise ValueError("Provided artifact_uri is not a directory.")
+        else:
+            artifact_path = os.getcwd()
+
+        artifact_path = os.path.join(artifact_path, self.WEIGHTS_FILE)
+        with smart_open.open(artifact_path, "rb") as f:
+            self._model = torch.load(f)
+
+        if self._cuda:
+            self._model.cuda()
+
+        config = timm.data.resolve_data_config(model=self.TIMM_MODEL_NAME, args=[])
+        self._transform = timm.data.create_transform(
+            is_training=False, use_prefetcher=False, **config
+        )
+
+    def preprocess(self, request_dict: Dict[str, List[str]]) -> torch.Tensor:
+        """Performs preprocessing.
+
+        By default, the server expects a request body consisting of a valid JSON
+        object. This will be parsed by the handler before it's evaluated by the
+        preprocess method.
+
+        Args:
+          request_dict: Parsed request body. We expect that the input consists of
+            a list of base64-encoded image files under the "instances" key. (Any
+            image format that PIL.image.open can handle is okay.)
+
+        Returns:
+          torch.Tensor containing the preprocessed images as a batch. If GPU is
+          available, the result tensor will be stored on GPU.
+        """
+
+        if "instances" not in request_dict:
+            raise HTTPException(
+                status_code=400,
+                detail='Request must contain "instances" as a top-level key.',
+            )
+
+        tensors = []
+
+        for (i, image) in enumerate(request_dict["instances"]):
+            # We use Base64 encoding to handle image data.
+            # This is probably the best we can do while still using JSON input.
+            # Overriding the input format requires building a custom Handler.
+            try:
+                image_bytes = base64.b64decode(image, validate=True)
+            except (binascii.Error, TypeError) as e:
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Base64 decoding of the input image at index {i} failed:"
+                    f" {str(e)}",
+                )
+
+            try:
+                pil_image = PIL.Image.open(io.BytesIO(image_bytes)).convert("RGB")
+            except PIL.UnidentifiedImageError:
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"The input image at index {i} could not be identified as an"
+                    " image file.",
+                )
+
+            tensors.append(self._transform(pil_image))
+
+        with torch.inference_mode():
+            result = torch.stack(tensors)
+            if self._cuda:
+                result = result.cuda()
+        return result
+
+    def predict(self, instances: torch.Tensor) -> torch.Tensor:
+        """Performs prediction.
+
+        Args:
+          instances: torch.Tensor with type torch.float32 and shape
+            [?, 3, 224, 224], containing the pre-processed input images.
+
+        Returns:
+          Vector of scores with type torch.float32 and shape [?, 1000],
+            representing the model's estimate of the likelihood that the
+            input belongs to the Imagenet class with that index.
+        """
+        with torch.inference_mode():
+            class_scores = self._model(instances)
+        return class_scores
+
+    def postprocess(
+        self, class_scores: torch.Tensor
+    ) -> Dict[str, List[Dict[str, Union[str, int, float]]]]:
+        """Translate the model output into a classification result.
+
+        Args:
+          class_scores: torch.Tensor with type torch.float32 and shape
+            [?, 1000], containing the scores assigned to each class by
+            the model.
+
+        Returns:
+          Dictionary containing the list of classification results. Each
+            classification result contains the probabilities, class names, and
+            class indices of the classes with the top class scores as reported by
+            the model.
+        """
+        class_probs = F.softmax(class_scores, dim=1)
+        top_k = class_probs.topk(self.NUM_TOP_CLASSES_TO_RETURN)
+        top_k_values = top_k.values.numpy().tolist()
+        top_k_indices = top_k.indices.numpy().tolist()
+        predictions = [
+            dict(
+                probabilities=values,
+                indices=indices,
+                class_names=[IMAGENET_CLASSES[int(class_num)] for class_num in indices],
+            )
+            for (values, indices) in zip(top_k_values, top_k_indices)
+        ]
+        return {"predictions": predictions}

community-content/pluto_on_workbench/README.md

@@ -0,0 +1,52 @@
+# Overview
+*Pluto* is a programming environment for Julia, designed to be interactive and helpful. It provides a familiar notebook interface but it is not a Jupyter notebook. The biggest difference is that Pluto notebooks are reactive, changing a variable or function in one cell causes the cells that depend on that variable or function to be reevaluated. Pluto also provides useful interaction mechanisms that allow users to dynamically interact with the notebooks computation state.
+
+The JuliaCon 2020 presentation: [Interactive notebooks ~ Pluto.jl]() provides a good introduction to Pluto. The source is at [fonsp/Pluto.jl]()
+
+# Install Pluto
+
+## Create a Vertex AI JupyterLab Instance
+
+1. From the [GCP console](https://console.cloud.google.com) "hamburger menu"
+
+   select Vertex AI > Workbench
+2. Click NEW NOTEBOOK
+   
+   * Choose Python 3 if you won't be using a GPU
+   * Choose Python 3 (CUDA Toolkit xx.y) if you do want use a GPU
+3. Give the notebook an appropriate name
+4. Edit Notebook properties if you have special requirements otherwise accept the defaults and click CREATE
+5. When the notebook instance is ready click OPEN JUPYTERLAB
+
+## Configure JupyterLab
+
+1. Open a terminal by clicking the Terminal icon.
+1. Install the plutoserver
+pip3 install git+https://github.com/fonsp/pluto-on-jupyterlab.git
+1. In a browser go to [julialang.org/downloads](https://julialang.org/downloads/)
+1. In the Current stable release right click on the `Generic Linux on x86 / 64-bit (glibc)` link
+Select copy link address
+1. Back in the terminal switch to root via
+sudo -i
+1. Download the release to /opt and install julia in /usr/local/bin
+    ```bash
+    cd /opt
+    wget <paste the release link address>
+    tar xf <name of the downloaded tar file>
+    ln -s /opt/<julia-x.y.z>/bin/julia /usr/local/bin
+    ^d
+    ```
+1. Add the Pluto package to Julia
+   ```bash
+    julia
+    julia> ]add Pluto
+    julia> bksp
+    julia> using Pluto
+    julia> ^d
+    ```
+1. From the JupyterLab menu bar select File > Shut Down
+
+# Start Pluto
+1. Click OPEN JUPYTERLAB in the Workbench
+1. In the Notebook section of the Launcher click Pluto.jl
+1. The welcome to Pluto.jl screen should appear

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb

@@ -1,474 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a6b56b1c7b76"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2021 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c414a395a19b"
-      },
-      "source": [
-        "# PyTorch Image Classification Multi-Node Distributed Data Parallel Training on CPU using Vertex Training with Custom Container"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "b98238e32cf7"
-      },
-      "source": [
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "03d216c7f7b1"
-      },
-      "source": [
-        "## Setup"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "c5ac73516218"
-      },
-      "outputs": [],
-      "source": [
-        "PROJECT_ID = \"YOUR PROJECT ID\"\n",
-        "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
-        "REGION = \"YOUR REGION\"\n",
-        "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "0b5ae674177e"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al $BUCKET_NAME"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "19a9b3bdd553"
-      },
-      "outputs": [],
-      "source": [
-        "content_name = \"pt-img-cls-multi-node-ddp-cust-cont\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "57bf6f8b4361"
-      },
-      "source": [
-        "## Local Training"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "e5d8a3443da0"
-      },
-      "outputs": [],
-      "source": [
-        "! ls trainer"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "07f79309472d"
-      },
-      "outputs": [],
-      "source": [
-        "! cat trainer/requirements.txt"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "e16cd8bb7483"
-      },
-      "outputs": [],
-      "source": [
-        "! pip install -r trainer/requirements.txt"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "0b8a210718c4"
-      },
-      "outputs": [],
-      "source": [
-        "! cat trainer/task.py"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "c0c6e7dfb3c6"
-      },
-      "outputs": [],
-      "source": [
-        "%run trainer/task.py --epochs 5 --no-cuda --local-mode"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "31dfdeede587"
-      },
-      "outputs": [],
-      "source": [
-        "! ls ./tmp"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "48d56ec621cc"
-      },
-      "outputs": [],
-      "source": [
-        "! rm -rf ./tmp"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8f3ea1210749"
-      },
-      "source": [
-        "## Vertex Training using Vertex SDK and Custom Container"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "93002a20a2a6"
-      },
-      "source": [
-        "### Build Custom Container"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "4130ce43fd08"
-      },
-      "outputs": [],
-      "source": [
-        "hostname = \"gcr.io\"\n",
-        "image_name = content_name\n",
-        "tag = \"latest\"\n",
-        "\n",
-        "custom_container_image_uri = f\"{hostname}/{PROJECT_ID}/{image_name}:{tag}\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2f1fc5b05240"
-      },
-      "outputs": [],
-      "source": [
-        "! cd trainer && docker build -t $custom_container_image_uri -f Dockerfile ."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b4f274f499ac"
-      },
-      "outputs": [],
-      "source": [
-        "! docker run --rm $custom_container_image_uri --epochs 5 --no-cuda --local-mode"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ee1a0a06d0b4"
-      },
-      "outputs": [],
-      "source": [
-        "! docker push $custom_container_image_uri"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "cb763be12fc9"
-      },
-      "outputs": [],
-      "source": [
-        "! gcloud container images list --repository $hostname/$PROJECT_ID"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "10c8cc6b3334"
-      },
-      "source": [
-        "### Initialize Vertex SDK"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1a12348169fa"
-      },
-      "outputs": [],
-      "source": [
-        "! pip install -r requirements.txt"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "42e981cefe41"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "aiplatform.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    staging_bucket=BUCKET_NAME,\n",
-        "    location=REGION,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "73c92c9298e9"
-      },
-      "source": [
-        "### Create a Vertex Tensorboard Instance"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "bde509558cd5"
-      },
-      "outputs": [],
-      "source": [
-        "content_name = content_name + \"-cpu\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "6d7908c0083c"
-      },
-      "outputs": [],
-      "source": [
-        "tensorboard = aiplatform.Tensorboard.create(\n",
-        "    display_name=content_name,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "a1f0a4f54037"
-      },
-      "source": [
-        "#### Option: Use a Previously Created Vertex Tensorboard Instance\n",
-        "\n",
-        "```\n",
-        "tensorboard_name = \"Your Tensorboard Resource Name or Tensorboard ID\"\n",
-        "tensorboard = aiplatform.Tensorboard(tensorboard_name=tensorboard_name)\n",
-        "```"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "a4cac84e04ac"
-      },
-      "source": [
-        "### Run a Vertex SDK CustomContainerTrainingJob"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f92e8fdd44ee"
-      },
-      "outputs": [],
-      "source": [
-        "display_name = content_name\n",
-        "gcs_output_uri_prefix = f\"{BUCKET_NAME}/{display_name}\"\n",
-        "\n",
-        "replica_count = 4\n",
-        "machine_type = \"n1-standard-4\"\n",
-        "\n",
-        "args = [\n",
-        "    \"--backend\",\n",
-        "    \"gloo\",\n",
-        "    \"--no-cuda\",\n",
-        "    \"--batch-size\",\n",
-        "    \"128\",\n",
-        "    \"--epochs\",\n",
-        "    \"25\",\n",
-        "]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ae4c57df7e07"
-      },
-      "outputs": [],
-      "source": [
-        "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
-        "    display_name=display_name,\n",
-        "    container_uri=custom_container_image_uri,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "35cf3ecdf0df"
-      },
-      "outputs": [],
-      "source": [
-        "custom_container_training_job.run(\n",
-        "    args=args,\n",
-        "    base_output_dir=gcs_output_uri_prefix,\n",
-        "    replica_count=replica_count,\n",
-        "    machine_type=machine_type,\n",
-        "    tensorboard=tensorboard.resource_name,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "49d10dded73b"
-      },
-      "outputs": [],
-      "source": [
-        "print(f\"Custom Training Job Name: {custom_container_training_job.resource_name}\")\n",
-        "print(f\"GCS Output URI Prefix: {gcs_output_uri_prefix}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "78398f52807b"
-      },
-      "source": [
-        "### Training Output Artifact"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fc74422de1d1"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls $gcs_output_uri_prefix"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5e99a6a05b10"
-      },
-      "source": [
-        "## Clean Up Artifact"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b0c1b3f7466b"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil rm -rf $gcs_output_uri_prefix"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "name": "multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb

@@ -1,347 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a6b56b1c7b76"
-      },
-      "outputs": [],
-      "source": [
-        "# Copyright 2021 Google LLC\n",
-        "#\n",
-        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-        "# you may not use this file except in compliance with the License.\n",
-        "# You may obtain a copy of the License at\n",
-        "#\n",
-        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
-        "#\n",
-        "# Unless required by applicable law or agreed to in writing, software\n",
-        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-        "# See the License for the specific language governing permissions and\n",
-        "# limitations under the License."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "20a5ea0081d0"
-      },
-      "source": [
-        "# PyTorch Image Classification Multi-Node Distributed Data Parallel Training on GPU using Vertex Training with Custom Container"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8752d4a255fb"
-      },
-      "source": [
-        "<table align=\"left\">\n",
-        "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
-        "    </a>\n",
-        "  </td>\n",
-        "</table>"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "03d216c7f7b1"
-      },
-      "source": [
-        "## Setup"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "c5ac73516218"
-      },
-      "outputs": [],
-      "source": [
-        "PROJECT_ID = \"YOUR PROJECT ID\"\n",
-        "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
-        "REGION = \"YOUR REGION\"\n",
-        "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "0b5ae674177e"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al $BUCKET_NAME"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "19a9b3bdd553"
-      },
-      "outputs": [],
-      "source": [
-        "content_name = \"pt-img-cls-multi-node-ddp-cust-cont\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5307fe28b633"
-      },
-      "source": [
-        "## Vertex Training using Vertex SDK and Custom Container"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "46cb58c7fbf9"
-      },
-      "source": [
-        "### Built Custom Container"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "97e66e9f9bab"
-      },
-      "outputs": [],
-      "source": [
-        "hostname = \"gcr.io\"\n",
-        "image_name = content_name\n",
-        "tag = \"latest\"\n",
-        "\n",
-        "custom_container_image_uri = f\"{hostname}/{PROJECT_ID}/{image_name}:{tag}\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ae9b29c4773f"
-      },
-      "source": [
-        "### Initialize Vertex SDK"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "dc1e84d5dec2"
-      },
-      "outputs": [],
-      "source": [
-        "! pip install -r requirements.txt"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "6964be27b98e"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "aiplatform.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    staging_bucket=BUCKET_NAME,\n",
-        "    location=REGION,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "594a91f438f2"
-      },
-      "source": [
-        "### Create a Vertex Tensorboard Instance"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "93134273261e"
-      },
-      "outputs": [],
-      "source": [
-        "content_name = content_name + \"-gpu\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "c2bd82dbcd9b"
-      },
-      "outputs": [],
-      "source": [
-        "tensorboard = aiplatform.Tensorboard.create(\n",
-        "    display_name=content_name,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ebc593c6472e"
-      },
-      "source": [
-        "#### Option: Use a Previously Created Vertex Tensorboard Instance\n",
-        "\n",
-        "```\n",
-        "tensorboard_name = \"Your Tensorboard Resource Name or Tensorboard ID\"\n",
-        "tensorboard = aiplatform.Tensorboard(tensorboard_name=tensorboard_name)\n",
-        "```"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0769e8e34c2f"
-      },
-      "source": [
-        "### Run a Vertex SDK CustomContainerTrainingJob"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "023f33ece826"
-      },
-      "outputs": [],
-      "source": [
-        "display_name = content_name\n",
-        "gcs_output_uri_prefix = f\"{BUCKET_NAME}/{display_name}\"\n",
-        "\n",
-        "replica_count = 1\n",
-        "machine_type = \"n1-standard-4\"\n",
-        "accelerator_count = 4\n",
-        "accelerator_type = \"NVIDIA_TESLA_K80\"\n",
-        "\n",
-        "args = [\n",
-        "    \"--backend\",\n",
-        "    \"nccl\",\n",
-        "    \"--batch-size\",\n",
-        "    \"128\",\n",
-        "    \"--epochs\",\n",
-        "    \"25\",\n",
-        "]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d4b599e726ef"
-      },
-      "outputs": [],
-      "source": [
-        "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
-        "    display_name=display_name,\n",
-        "    container_uri=custom_container_image_uri,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "81321e3bdf7f"
-      },
-      "outputs": [],
-      "source": [
-        "custom_container_training_job.run(\n",
-        "    args=args,\n",
-        "    base_output_dir=gcs_output_uri_prefix,\n",
-        "    replica_count=replica_count,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    tensorboard=tensorboard.resource_name,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "5100712c2c4c"
-      },
-      "outputs": [],
-      "source": [
-        "print(f\"Custom Training Job Name: {custom_container_training_job.resource_name}\")\n",
-        "print(f\"GCS Output URI Prefix: {gcs_output_uri_prefix}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f9b77676e5a6"
-      },
-      "source": [
-        "### Training Output Artifact"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "0e171ce95ace"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls $gcs_output_uri_prefix"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cf1b74a12b87"
-      },
-      "source": [
-        "## Clean Up Artifact"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a0b15089c341"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil rm -rf $gcs_output_uri_prefix"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "name": "multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}

community-content/pytorch_pre_built_images_deployment/README.md

@@ -0,0 +1,30 @@
+# PyTorch Deployment on Google Cloud: Text Classification
+
+**This is an Experimental release**, covered by the Pre-GA Offerings Terms of your Google Cloud Platform [Terms of Service](https://cloud.google.com/terms).
+
+Experiments are focused on validating a prototype and are not guaranteed to be released. They are not intended for production use or covered by any SLA, support obligation, or deprecation policy and might be subject to backward-incompatible changes.
+
+**Kindly drop us a note before you run any scale tests.**
+
+**Do not hesitate to contact vertexai-prediction-preview-feedback@google.com if you have any questions or run into any issues.**
+
+The projects need to be added to the allowlist in order to deploy PyTorch models using Vertex AI Prediction pre-built PyTorch images. If you are interested in the feature, please send an email to vertexai-prediction-preview-feedback@google.com to provide your project numbers OR project ids.
+
+## Overview
+
+In the PyTorch on Google Cloud series of blog posts, we aim to share how to deploy PyTorch models at scale on [Vertex AI](https://cloud.google.com/vertex-ai).
+
+This tutorial on text classification shows how to deploy a PyTorch based text classification model on [Vertex AI](https://cloud.google.com/vertex-ai/docs/start/client-libraries#python) using Vertex SDK and [`gcloud ai`](https://cloud.google.com/sdk/gcloud/reference/beta/ai).
+
+## Notebooks
+
+| <h4>Notebook</h4> | <h4>Description</h4>                       |
+| :-------- | :------- |
+| [pytorch-text-classification-vertex-ai-deploy.ipynb](./pytorch-text-classification-vertex-ai-deploy.ipynb) | Notebook to show deploying a PyTorch model on Vertex AI |
+
+## Folders
+
+
+| <h4>Folder Name</h4> | <h4>Description</h4>                       |
+| :-------- | :------- |
+| [`predictor`](./predictor) | Folder with custom prediction handler to deploy a PyTorch model to Vertex Prediction. In the [notebook](./pytorch-text-classification-vertex-ai-deploy.ipynb), this folder is used for deploying a PyTorch model on Vertex AI using Vertex Prediction pre-built PyTorch images |

community-content/pytorch_pre_built_images_deployment/predictor/custom_handler.py

@@ -0,0 +1,91 @@
+
+import os
+import json
+import logging
+
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from ts.torch_handler.base_handler import BaseHandler
+
+logger = logging.getLogger(__name__)
+
+
+class TransformersClassifierHandler(BaseHandler):
+    """
+    The handler takes an input string and returns the classification text 
+    based on the serialized transformers checkpoint.
+    """
+    def __init__(self):
+        super(TransformersClassifierHandler, self).__init__()
+        self.initialized = False
+
+    def initialize(self, ctx):
+        """ Loads the model.pt file and initialized the model object.
+        Instantiates Tokenizer for preprocessor to use
+        Loads labels to name mapping file for post-processing inference response
+        """
+        self.manifest = ctx.manifest
+
+        properties = ctx.system_properties
+        model_dir = properties.get("model_dir")
+        self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu")
+
+        # Read model serialize/pt file
+        serialized_file = self.manifest["model"]["serializedFile"]
+        model_pt_path = os.path.join(model_dir, serialized_file)
+        if not os.path.isfile(model_pt_path):
+            raise RuntimeError("Missing the model.pt or pytorch_model.bin file")
+        
+        # Load model
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+        self.model.to(self.device)
+        self.model.eval()
+        logger.debug('Transformer model from path {0} loaded successfully'.format(model_dir))
+        
+        # Ensure to use the same tokenizer used during training
+        self.tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
+
+        # Read the mapping file, index to object name
+        mapping_file_path = os.path.join(model_dir, "index_to_name.json")
+
+        if os.path.isfile(mapping_file_path):
+            with open(mapping_file_path) as f:
+                self.mapping = json.load(f)
+        else:
+            logger.warning('Missing the index_to_name.json file. Inference output will default.')
+            self.mapping = {"0": "Negative",  "1": "Positive"}
+
+        self.initialized = True
+
+    def preprocess(self, data):
+        """ Preprocessing input request by tokenizing
+            Extend with your own preprocessing steps as needed
+        """
+        text = data[0].get("data")
+        if text is None:
+            text = data[0].get("body")
+        sentences = text.decode('utf-8')
+        logger.info("Received text: '%s'", sentences)
+
+        # Tokenize the texts
+        tokenizer_args = ((sentences,))
+        inputs = self.tokenizer(*tokenizer_args,
+                                padding='max_length',
+                                max_length=128,
+                                truncation=True,
+                                return_tensors = "pt")
+        return inputs
+
+    def inference(self, inputs):
+        """ Predict the class of a text using a trained transformer model.
+        """
+        prediction = self.model(inputs['input_ids'].to(self.device))[0].argmax().item()
+
+        if self.mapping:
+            prediction = self.mapping[str(prediction)]
+
+        logger.info("Model predicted: '%s'", prediction)
+        return [prediction]
+
+    def postprocess(self, inference_output):
+        return inference_output

community-content/pytorch_pre_built_images_deployment/predictor/index_to_name.json

@@ -0,0 +1,5 @@
+
+{
+    "0": "Negative", 
+    "1": "Positive"
+}

community-content/pytorch_text_classification_using_vertex_sdk_and_gcloud/predictor/Dockerfile

@@ -2,10 +2,13 @@
 FROM pytorch/torchserve:latest-cpu
 
 # install dependencies
+RUN python3 -m pip install --upgrade pip
 RUN pip3 install transformers
 
+USER model-server
+
 # copy model artifacts, custom handler and other dependencies
-COPY ./custom_text_handler.py /home/model-server/
+COPY ./custom_handler.py /home/model-server/
 COPY ./index_to_name.json /home/model-server/
 COPY ./model/finetuned-bert-classifier/ /home/model-server/
 
@@ -21,7 +24,7 @@ EXPOSE 7080
 EXPOSE 7081
 
 # create model archive file packaging model artifacts and dependencies
-RUN torch-model-archiver -f   --model-name=finetuned-bert-classifier   --version=1.0   --serialized-file=/home/model-server/pytorch_model.bin   --handler=/home/model-server/custom_text_handler.py   --extra-files "/home/model-server/config.json,/home/model-server/tokenizer.json,/home/model-server/training_args.bin,/home/model-server/tokenizer_config.json,/home/model-server/special_tokens_map.json,/home/model-server/vocab.txt,/home/model-server/index_to_name.json"   --export-path=/home/model-server/model-store
+RUN torch-model-archiver -f   --model-name=finetuned-bert-classifier   --version=1.0   --serialized-file=/home/model-server/pytorch_model.bin   --handler=/home/model-server/custom_handler.py   --extra-files "/home/model-server/config.json,/home/model-server/tokenizer.json,/home/model-server/training_args.bin,/home/model-server/tokenizer_config.json,/home/model-server/special_tokens_map.json,/home/model-server/vocab.txt,/home/model-server/index_to_name.json"   --export-path=/home/model-server/model-store
 
 # run Torchserve HTTP serve to respond to prediction requests
-CMD ["torchserve",      "--start",      "--ts-config=/home/model-server/config.properties",      "--models",      "finetuned-bert-classifier=finetuned-bert-classifier.mar",      "--model-store",      "/home/model-server/model-store"]
+CMD ["torchserve",      "--start",      "--ts-config=/home/model-server/config.properties",      "--models",      "finetuned-bert-classifier=finetuned-bert-classifier.mar",      "--model-store",      "/home/model-server/model-store"]

community-content/pytorch_text_classification_using_vertex_sdk_and_gcloud/pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb

@@ -63,12 +63,12 @@
         "- [Training](#Training)\n",
         "    - [Run Training Locally in the Notebook](#Training-locally-in-the-notebook)\n",
         "    - [Run Training Job on Vertex AI](#Training-on-Vertex-AI)\n",
-        "        - [Training with pre-built container](#Run-Custom-Job-on-Vertex-Training-with-a-pre-built-container)\n",
-        "        - [Training with custom container](#Run-Custom-Job-on-Vertex-Training-with-custom-container)\n",
+        "        - [Training with pre-built container](#Run-Custom-Job-on-Vertex-AI-Training-with-a-pre-built-container)\n",
+        "        - [Training with custom container](#Run-Custom-Job-on-Vertex-AI-Training-with-custom-container)\n",
         "- [Tuning](#Hyperparameter-Tuning) \n",
         "    - [Run Hyperparameter Tuning job on Vertex AI](#Run-Hyperparameter-Tuning-Job-on-Vertex-AI)\n",
         "- [Deploying](#Deploying)\n",
-        "    - [Deploying model on Vertex Predictions with custom container](#Deploying-model-on-Vertex-Predictions-with-custom-container)\n",
+        "    - [Deploying model on Vertex AI Predictions with custom container](#Deploying-model-on-Vertex AI-Predictions-with-custom-container)\n",
         "\n",
         "### Costs \n",
         "\n",
@@ -202,9 +202,9 @@
         "id": "e0c1dcadc2c8"
       },
       "source": [
-        "We will be using [Vertex SDK for Python](https://cloud.google.com/vertex-ai/docs/start/client-libraries#python) to interact with Vertex AI services. The high-level `aiplatform` library is designed to simplify common data science workflows by using wrapper classes and opinionated defaults. \n",
+        "We will be using [Vertex AI SDK for Python](https://cloud.google.com/vertex-ai/docs/start/client-libraries#python) to interact with Vertex AI services. The high-level `aiplatform` library is designed to simplify common data science workflows by using wrapper classes and opinionated defaults. \n",
         "\n",
-        "#### Install Vertex SDK for Python"
+        "#### Install Vertex AI SDK for Python"
       ]
     },
     {
@@ -658,8 +658,8 @@
       },
       "outputs": [],
       "source": [
-        "datasets = load_dataset(\"imdb\")\n",
-        "datasets"
+        "dataset = load_dataset(\"imdb\")\n",
+        "dataset"
       ]
     },
     {
@@ -668,7 +668,7 @@
         "id": "RzfPtOMoIrIu"
       },
       "source": [
-        "The `datasets` object itself is [`DatasetDict`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasetdict), which contains one key for the training, validation and test set."
+        "The `dataset` object itself is [`DatasetDict`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasetdict), which contains one key for the training, validation and test set."
       ]
     },
     {
@@ -681,12 +681,12 @@
       "source": [
         "print(\n",
         "    \"Total # of rows in training dataset {} and size {:5.2f} MB\".format(\n",
-        "        datasets[\"train\"].shape[0], datasets[\"train\"].size_in_bytes / (1024 * 1024)\n",
+        "        dataset[\"train\"].shape[0], dataset[\"train\"].size_in_bytes / (1024 * 1024)\n",
         "    )\n",
         ")\n",
         "print(\n",
         "    \"Total # of rows in test dataset {} and size {:5.2f} MB\".format(\n",
-        "        datasets[\"test\"].shape[0], datasets[\"test\"].size_in_bytes / (1024 * 1024)\n",
+        "        dataset[\"test\"].shape[0], dataset[\"test\"].size_in_bytes / (1024 * 1024)\n",
         "    )\n",
         ")"
       ]
@@ -708,7 +708,7 @@
       },
       "outputs": [],
       "source": [
-        "datasets[\"train\"][0]"
+        "dataset[\"train\"][0]"
       ]
     },
     {
@@ -728,7 +728,7 @@
       },
       "outputs": [],
       "source": [
-        "label_list = datasets[\"train\"].unique(\"label\")\n",
+        "label_list = dataset[\"train\"].unique(\"label\")\n",
         "label_list"
       ]
     },
@@ -779,7 +779,7 @@
       },
       "outputs": [],
       "source": [
-        "show_random_elements(datasets[\"train\"])"
+        "show_random_elements(dataset[\"train\"])"
       ]
     },
     {
@@ -883,7 +883,7 @@
       },
       "outputs": [],
       "source": [
-        "example = datasets[\"train\"][4]\n",
+        "example = dataset[\"train\"][4]\n",
         "print(example)"
       ]
     },
@@ -920,7 +920,7 @@
       "source": [
         "# Dataset loading repeated here to make this cell idempotent\n",
         "# Since we are over-writing datasets variable\n",
-        "datasets = load_dataset(\"imdb\")\n",
+        "dataset = load_dataset(\"imdb\")\n",
         "\n",
         "# Mapping labels to ids\n",
         "# NOTE: We can extract this automatically but the `Unique` method of the datasets\n",
@@ -948,7 +948,7 @@
         "\n",
         "\n",
         "# apply preprocessing function to input examples\n",
-        "datasets = datasets.map(preprocess_function, batched=True, load_from_cache_file=True)"
+        "dataset = dataset.map(preprocess_function, batched=True, load_from_cache_file=True)"
       ]
     },
     {
@@ -1091,8 +1091,8 @@
         "trainer = Trainer(\n",
         "    model,\n",
         "    args,\n",
-        "    train_dataset=datasets[\"train\"],\n",
-        "    eval_dataset=datasets[\"test\"],\n",
+        "    train_dataset=dataset[\"train\"],\n",
+        "    eval_dataset=dataset[\"test\"],\n",
         "    data_collator=default_data_collator,\n",
         "    tokenizer=tokenizer,\n",
         "    compute_metrics=compute_metrics,\n",
@@ -1199,7 +1199,7 @@
       "source": [
         "### Run predictions locally with sample examples\n",
         "\n",
-        "Using the trained model, we can predict the sentiment label for an input text after applying the preprocessing function that was used during the training. We will run the predictions locally in the notebook and later show how you can deploy the model to an endpoint using [TorchServe](https://pytorch.org/serve/) on Vertex Predictions."
+        "Using the trained model, we can predict the sentiment label for an input text after applying the preprocessing function that was used during the training. We will run the predictions locally in the notebook and later show how you can deploy the model to an endpoint using [TorchServe](https://pytorch.org/serve/) on Vertex AI Predictions."
       ]
     },
     {
@@ -1382,7 +1382,7 @@
         "id": "f7466d414a0e"
       },
       "source": [
-        "### Run Custom Job on Vertex Training with a pre-built container"
+        "### Run Custom Job on Vertex AI Training with a pre-built container"
       ]
     },
     {
@@ -1395,7 +1395,7 @@
         "\n",
         "In this notebook, we are using Hugging Face Datasets and fine tuning a transformer model from Hugging Face Transformers Library for sentiment analysis task using PyTorch. We will use [pre-built container for PyTorch](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers#pytorch) and package the training application code by adding standard Python dependencies - `transformers`, `datasets` and `tqdm` - in the `setup.py` file. \n",
         "\n",
-        "![Training with Prebuilt Containers on Vertex Training](./images/training-with-prebuilt-containers-on-vertex-training.png)"
+        "![Training with Prebuilt Containers on Vertex AI Training](./images/training-with-prebuilt-containers-on-vertex-training.png)"
       ]
     },
     {
@@ -1569,7 +1569,7 @@
       "source": [
         "#### **Run custom training job on Vertex AI**\n",
         "\n",
-        "We use [Vertex SDK for Python](https://cloud.google.com/vertex-ai/docs/start/client-libraries#client_libraries) to create and submit training job to the Vertex training service."
+        "We use [Vertex AI SDK for Python](https://cloud.google.com/vertex-ai/docs/start/client-libraries#client_libraries) to create and submit training job to the Vertex AI training service."
       ]
     },
     {
@@ -1578,7 +1578,7 @@
         "id": "5d2957ef04fd"
       },
       "source": [
-        "##### **Initialize the Vertex SDK for Python**"
+        "##### **Initialize the Vertex AI SDK for Python**"
       ]
     },
     {
@@ -1598,7 +1598,7 @@
         "id": "6b0fed34b728"
       },
       "source": [
-        "##### **Configure and submit Custom Job to Vertex Training service**"
+        "##### **Configure and submit Custom Job to Vertex AI Training service**"
       ]
     },
     {
@@ -1609,7 +1609,7 @@
       "source": [
         "Configure a [Custom Job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) with the [pre-built container](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers) image for PyTorch and training code packaged as Python source distribution. \n",
         "\n",
-        "**NOTE:** When using Vertex SDK for Python for submitting a training job, it creates a [Training Pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the Custom Job on Vertex Training service."
+        "**NOTE:** When using Vertex AI SDK for Python for submitting a training job, it creates a [Training Pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the Custom Job on Vertex AI Training service."
       ]
     },
     {
@@ -1686,7 +1686,7 @@
         "\n",
         "You can monitor the custom job launched from Cloud Console following the link [here](https://console.cloud.google.com/vertex-ai/training/training-pipelines/) or use gcloud CLI command [`gcloud beta ai custom-jobs stream-logs`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs)\n",
         "\n",
-        "![Monitor custom job progress in Vertex Training](./images/vertex-training-monitor-custom-job.png)"
+        "![Monitor custom job progress in Vertex AI Training](./images/vertex-training-monitor-custom-job.png)"
       ]
     },
     {
@@ -1798,7 +1798,7 @@
         "id": "c170d386492b"
       },
       "source": [
-        "### Run Custom Job on Vertex Training with custom container"
+        "### Run Custom Job on Vertex AI Training with custom container"
       ]
     },
     {
@@ -1807,7 +1807,7 @@
         "id": "035227b6e581"
       },
       "source": [
-        "To create a [training job with custom container](https://cloud.google.com/vertex-ai/docs/training/create-custom-container?hl=hr), you define a `Dockerfile` to install or add the dependencies required for the training job. Then, you build and test your Docker image locally to verify, push the image to Container Registry and submit a Custom Job to Vertex Training service.\n",
+        "To create a [training job with custom container](https://cloud.google.com/vertex-ai/docs/training/create-custom-container?hl=hr), you define a `Dockerfile` to install or add the dependencies required for the training job. Then, you build and test your Docker image locally to verify, push the image to Container Registry and submit a Custom Job to Vertex AI Training service.\n",
         "\n",
         "![Training with custom containers on Vertex AI](./images/training-with-custom-containers-on-vertex-training.png)"
       ]
@@ -1834,7 +1834,7 @@
         "%%writefile ./custom_container/Dockerfile\n",
         "\n",
         "# Use pytorch GPU base image\n",
-        "FROM gcr.io/cloud-aiplatform/training/pytorch-gpu.1-7\n",
+        "FROM us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-10:latest\n",
         "\n",
         "# set working directory\n",
         "WORKDIR /app\n",
@@ -1968,7 +1968,7 @@
         "id": "a23e5e34bea9"
       },
       "source": [
-        "##### **Initialize the Vertex SDK for Python**"
+        "##### **Initialize the Vertex AI SDK for Python**"
       ]
     },
     {
@@ -1988,11 +1988,11 @@
         "id": "abf1fa4085cb"
       },
       "source": [
-        "##### **Configure and submit Custom Job to Vertex Training service**\n",
+        "##### **Configure and submit Custom Job to Vertex AI Training service**\n",
         "\n",
         "Configure a [Custom Job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) with the [custom container](https://cloud.google.com/vertex-ai/docs/training/create-custom-container) image with training code and other dependencies\n",
         "\n",
-        "**NOTE:** When using Vertex SDK for Python for submitting a training job, it creates a [Training Pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the Custom Job to train on Vertex Training."
+        "**NOTE:** When using Vertex AI SDK for Python for submitting a training job, it creates a [Training Pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the Custom Job to train on Vertex AI Training."
       ]
     },
     {
@@ -2044,7 +2044,7 @@
       },
       "outputs": [],
       "source": [
-        "# submit the custom job to Vertex training service\n",
+        "# submit the custom job to Vertex AI training service\n",
         "model = job.run(\n",
         "    replica_count=1,\n",
         "    machine_type=\"n1-standard-8\",\n",
@@ -2065,7 +2065,7 @@
         "\n",
         "You can monitor the custom job launched from Cloud Console following the link [here](https://console.cloud.google.com/vertex-ai/training/training-pipelines/) or use gcloud CLI command [`gcloud beta ai custom-jobs stream-logs`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs)\n",
         "\n",
-        "![Monitor custom job progress in Vertex Training](./images/vertex-training-monitor-custom-job-container.png)"
+        "![Monitor custom job progress in Vertex AI Training](./images/vertex-training-monitor-custom-job-container.png)"
       ]
     },
     {
@@ -2148,11 +2148,11 @@
         "id": "ba6122f929e3"
       },
       "source": [
-        "The training application code for fine-tuning a transformer model for sentiment analysis task uses hyperparameters such as learning rate and weight decay. These hyperparameters control the behavior of the training algorithm and can have a significant effect on the performance of the resulting model. This part of the notebook show how you can automate tuning these hyperparameters with Vertex Training service.\n",
+        "The training application code for fine-tuning a transformer model for sentiment analysis task uses hyperparameters such as learning rate and weight decay. These hyperparameters control the behavior of the training algorithm and can have a significant effect on the performance of the resulting model. This part of the notebook show how you can automate tuning these hyperparameters with Vertex AI Training service.\n",
         "\n",
-        "We submit a [Hyperparameter Tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to Vertex Training service by packaging the training application code and dependencies in a Docker container and push the container to Google Container Registry, similar to running a Custom Job on Vertex AI with Custom Container.\n",
+        "We submit a [Hyperparameter Tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to Vertex AI Training service by packaging the training application code and dependencies in a Docker container and push the container to Google Container Registry, similar to running a Custom Job on Vertex AI with Custom Container.\n",
         "\n",
-        "![Hyperparameter Tuning with Custom Containers on Vertex Training](./images/hp-tuning-with-custom-containers-on-vertex-training.png)"
+        "![Hyperparameter Tuning with Custom Containers on Vertex AI Training](./images/hp-tuning-with-custom-containers-on-vertex-training.png)"
       ]
     },
     {
@@ -2163,7 +2163,7 @@
       "source": [
         "### How hyperparameter tuning works in Vertex AI?\n",
         "\n",
-        "Following are the high level steps involved in running a Hyperparameter Tuning job on Vertex Training service:\n",
+        "Following are the high level steps involved in running a Hyperparameter Tuning job on Vertex AI Training service:\n",
         "\n",
         "- You define the hyperparameters to tune the model along with the metric (or goal) to optimize\n",
         "- Vertex AI runs multiple trials of your training application with the hyperparameters and limits you specified - maximum number of trials to run and number of parallel trials. \n",
@@ -2297,7 +2297,7 @@
       "source": [
         "### Run Hyperparameter Tuning Job on Vertex AI\n",
         "\n",
-        "Before submitting the hyperparameter tuning job to Vertex AI, push the custom container image with training application to Google Cloud Container Registry and then submit the job to Vertex AI. We will be using the same image used for running Custom Job on Vertex Training service."
+        "Before submitting the hyperparameter tuning job to Vertex AI, push the custom container image with training application to Google Cloud Container Registry and then submit the job to Vertex AI. We will be using the same image used for running Custom Job on Vertex AI Training service."
       ]
     },
     {
@@ -2326,7 +2326,7 @@
         "id": "f60fab07d67c"
       },
       "source": [
-        "##### **Initialize the Vertex SDK for Python**"
+        "##### **Initialize the Vertex AI SDK for Python**"
       ]
     },
     {
@@ -2346,7 +2346,7 @@
         "id": "6652aa63ddff"
       },
       "source": [
-        "##### **Configure and submit Hyperparameter Tuning Job to Vertex Training service**\n",
+        "##### **Configure and submit Hyperparameter Tuning Job to Vertex AI Training service**\n",
         "\n",
         "Configure a [Hyperparameter Tuning Job](https://cloud.google.com/vertex-ai/docs/training/using-hyperparameter-tuning) with the [custom container](https://cloud.google.com/vertex-ai/docs/training/create-custom-container) image with training code and other dependencies.\n",
         "\n",
@@ -2374,7 +2374,7 @@
         "id": "9d46db3a8b23"
       },
       "source": [
-        "Define the training arguments with `hp-tune` argument set to `y` so that training application code can report metrics to Vertex"
+        "Define the training arguments with `hp-tune` argument set to `y` so that training application code can report metrics to Vertex AI"
       ]
     },
     {
@@ -2548,7 +2548,7 @@
         "\n",
         "You can monitor the hyperparameter tuning job launched from Cloud Console following the link [here](https://console.cloud.google.com/vertex-ai/training/hyperparameter-tuning-jobs/) or use gcloud CLI command [`gcloud beta ai custom-jobs stream-logs`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs)\n",
         "\n",
-        "![Monitor hyperparameter tuning job progress in Vertex Training](./images/vertex-training-monitor-hptuning-job-container.png)"
+        "![Monitor hyperparameter tuning job progress in Vertex AI Training](./images/vertex-training-monitor-hptuning-job-container.png)"
       ]
     },
     {
@@ -2557,7 +2557,7 @@
         "id": "ba934b434f03"
       },
       "source": [
-        "After the job is finished, you can view and format the results of the hyperparameter tuning Trials (run by Vertex Training service) as a Pandas dataframe"
+        "After the job is finished, you can view and format the results of the hyperparameter tuning Trials (run by Vertex AI Training service) as a Pandas dataframe"
       ]
     },
     {
@@ -2612,7 +2612,7 @@
         "id": "5dbccb2b7d32"
       },
       "source": [
-        "Now from the results of Trials, you can pick the best performing Trial to deploy to Vertex Predictions"
+        "Now from the results of Trials, you can pick the best performing Trial to deploy to Vertex AI Predictions"
       ]
     },
     {
@@ -2701,8 +2701,8 @@
         "JOB_NAME=${JOB_PREFIX}-pytorch-hptune-$(date +%Y%m%d%H%M%S)\n",
         "echo \"Launching hyperparameter tuning job with display name as \"$JOB_NAME\n",
         "\n",
-        "# BUCKET_NAME: Change to your bucket name\n",
-        "BUCKET_NAME=$1 # <-- CHANGE TO YOUR BUCKET NAME\n",
+        "# BUCKET_NAME is a required parameter to run the cell.\n",
+        "BUCKET_NAME=$1\n",
         "\n",
         "# APP_NAME: get application name\n",
         "APP_NAME=$2\n",
@@ -2711,7 +2711,7 @@
         "JOB_DIR=${BUCKET_NAME}/${JOB_PREFIX}/model/${JOB_NAME}\n",
         "\n",
         "# custom container image URI\n",
-        "CUSTOM_TRAIN_IMAGE_URI=f'gcr.io/'${PROJECT_ID}'/pytorch_gpu_train_'${APP_NAME}\n",
+        "CUSTOM_TRAIN_IMAGE_URI='gcr.io/'${PROJECT_ID}'/pytorch_gpu_train_'${APP_NAME}\n",
         "\n",
         "# ========================================================\n",
         "# create hyperparameter tuning configuration file\n",
@@ -2772,20 +2772,20 @@
       "source": [
         "## Deploying\n",
         "\n",
-        "Deploying a PyTorch model on [Vertex Predictions](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions) requires to use a custom container that serves online predictions. You will deploy a container running [PyTorch's TorchServe](https://pytorch.org/serve/) tool in order to serve predictions from a fine-tuned transformer model from Hugging Face Transformers for sentiment analysis task. You can then use Vertex Predictions to classify sentiment of input texts. \n",
+        "Deploying a PyTorch model on [Vertex AI Predictions](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions) requires to use a custom container that serves online predictions. You will deploy a container running [PyTorch's TorchServe](https://pytorch.org/serve/) tool in order to serve predictions from a fine-tuned transformer model from Hugging Face Transformers for sentiment analysis task. You can then use Vertex AI Predictions to classify sentiment of input texts. \n",
         "\n",
-        "### Deploying model on Vertex Predictions with custom container\n",
+        "### Deploying model on Vertex AI Predictions with custom container\n",
         "\n",
-        "To use a custom container to serve predictions from a PyTorch model, you must provide Vertex AI with a Docker container image that runs an HTTP server, such as TorchServe in this case. Please refer to [documentation](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements) that describes the container image requirements to be compatible with Vertex Predictions.\n",
+        "To use a custom container to serve predictions from a PyTorch model, you must provide Vertex AI with a Docker container image that runs an HTTP server, such as TorchServe in this case. Please refer to [documentation](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements) that describes the container image requirements to be compatible with Vertex AI Predictions.\n",
         "\n",
-        "![Serving with Custom Containers on Vertex Predictions](./images/serve-pytorch-model-on-vertex-predictions-with-custom-containers.png)\n",
+        "![Serving with Custom Containers on Vertex AI Predictions](./images/serve-pytorch-model-on-vertex-predictions-with-custom-containers.png)\n",
         "\n",
-        "Essentially, to deploy a PyTorch model on Vertex Predictions following are the steps:\n",
+        "Essentially, to deploy a PyTorch model on Vertex AI Predictions following are the steps:\n",
         "\n",
         "1. Package the trained model artifacts including [default](https://pytorch.org/serve/#default-handlers) or [custom](https://pytorch.org/serve/custom_service.html) handlers by creating an archive file using [Torch model archiver](https://github.com/pytorch/serve/tree/master/model-archiver)\n",
-        "2. Build a [custom container](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements) compatible with Vertex Predictions to serve the model using Torchserve\n",
-        "3. Upload the model with custom container image to serve predictions as a Vertex Model resource\n",
-        "4. Create a Vertex Endpoint and [deploy the model](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api) resource"
+        "2. Build a [custom container](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements) compatible with Vertex AI Predictions to serve the model using Torchserve\n",
+        "3. Upload the model with custom container image to serve predictions as a Vertex AI Model resource\n",
+        "4. Create a Vertex AI Endpoint and [deploy the model](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api) resource"
       ]
     },
     {
@@ -3048,10 +3048,13 @@
         "FROM pytorch/torchserve:latest-cpu\n",
         "\n",
         "# install dependencies\n",
+        "RUN python3 -m pip install --upgrade pip\n",
         "RUN pip3 install transformers\n",
         "\n",
+        "USER model-server\n",
+        "\n",
         "# copy model artifacts, custom handler and other dependencies\n",
-        "COPY ./custom_text_handler.py /home/model-server/\n",
+        "COPY ./custom_handler.py /home/model-server/\n",
         "COPY ./index_to_name.json /home/model-server/\n",
         "COPY ./model/$APP_NAME/ /home/model-server/\n",
         "\n",
@@ -3071,7 +3074,7 @@
         "  --model-name=$APP_NAME \\\n",
         "  --version=1.0 \\\n",
         "  --serialized-file=/home/model-server/pytorch_model.bin \\\n",
-        "  --handler=/home/model-server/custom_text_handler.py \\\n",
+        "  --handler=/home/model-server/custom_handler.py \\\n",
         "  --extra-files \"/home/model-server/config.json,/home/model-server/tokenizer.json,/home/model-server/training_args.bin,/home/model-server/tokenizer_config.json,/home/model-server/special_tokens_map.json,/home/model-server/vocab.txt,/home/model-server/index_to_name.json\" \\\n",
         "  --export-path=/home/model-server/model-store\n",
         "\n",
@@ -3130,7 +3133,7 @@
       "source": [
         "#### **Run the container locally** ***[Optional]***\n",
         "\n",
-        "Before push the container image to Container Registry to use it with Vertex Predictions, you can run it as a container in your local environment to verify that the server works as expected"
+        "Before push the container image to Container Registry to use it with Vertex AI Predictions, you can run it as a container in your local environment to verify that the server works as expected"
       ]
     },
     {
@@ -3268,9 +3271,9 @@
         "id": "69477b3a00c0"
       },
       "source": [
-        "#### **Deploying the serving container to Vertex Predictions**\n",
+        "#### **Deploying the serving container to Vertex AI Predictions**\n",
         "\n",
-        "We create a model resource on Vertex AI and deploy the model to a Vertex Endpoints. You must deploy a model to an endpoint before using the model. The deployed model runs the custom container image to serve predictions. "
+        "We create a model resource on Vertex AI and deploy the model to a Vertex AI Endpoints. You must deploy a model to an endpoint before using the model. The deployed model runs the custom container image to serve predictions. "
       ]
     },
     {
@@ -3301,7 +3304,7 @@
         "id": "a3da91e19af4"
       },
       "source": [
-        "##### **Initialize the Vertex SDK for Python**"
+        "##### **Initialize the Vertex AI SDK for Python**"
       ]
     },
     {
@@ -3438,7 +3441,7 @@
         "id": "bc4673478269"
       },
       "source": [
-        "#### **Invoking the Endpoint with deployed Model using Vertex SDK to make predictions**"
+        "#### **Invoking the Endpoint with deployed Model using Vertex AI SDK to make predictions**"
       ]
     },
     {
@@ -3488,7 +3491,7 @@
       "source": [
         "##### **Formatting input for online prediction**\n",
         "\n",
-        "For online prediction requests, the prediction input instances must be formatted as JSON with base64 encoding as shown here:\n",
+        "This notebook uses [Torchserve's KServe based inference API](https://pytorch.org/serve/inference_api.html#kserve-inference-api) which is also [Vertex AI Predictions compatible format](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#prediction). For online prediction requests, format the prediction input instances as JSON with base64 encoding as shown here:\n",
         "\n",
         "```\n",
         "[\n",
@@ -3561,9 +3564,9 @@
       },
       "source": [
         "##### ***[Optional]*** **Make prediction requests using gcloud CLI**\n",
-        "You can also call the Vertex Endpoint to make predictions using [`gcloud beta ai endpoints predict`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/endpoints/predict). \n",
+        "You can also call the Vertex AI Endpoint to make predictions using [`gcloud beta ai endpoints predict`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/endpoints/predict). \n",
         "\n",
-        "The following cell shows how to make a prediction request to Vertex Endpoints using `gcloud` CLI: "
+        "The following cell shows how to make a prediction request to Vertex AI Endpoints using `gcloud` CLI: "
       ]
     },
     {
@@ -3654,12 +3657,12 @@
       },
       "outputs": [],
       "source": [
-        "delete_custom_job = True\n",
-        "delete_hp_tuning_job = True\n",
+        "delete_custom_job = False\n",
+        "delete_hp_tuning_job = False\n",
         "delete_endpoint = True\n",
-        "delete_model = True\n",
-        "delete_bucket = True\n",
-        "delete_image = True"
+        "delete_model = False\n",
+        "delete_bucket = False\n",
+        "delete_image = False"
       ]
     },
     {
@@ -3687,7 +3690,7 @@
         "\n",
         "client_options = {\"api_endpoint\": API_ENDPOINT}\n",
         "\n",
-        "# Initialize Vertex SDK\n",
+        "# Initialize Vertex AI SDK\n",
         "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
       ]
     },

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb

@@ -188,11 +188,14 @@
       },
       "outputs": [],
       "source": [
-        "! pip3 install {USER_FLAG} google-cloud-aiplatform==1.0.1\n",
-        "! pip3 install {USER_FLAG} google-cloud-pipeline-components==0.1.3\n",
+        "! pip3 install {USER_FLAG} google-cloud-aiplatform\n",
+        "! pip3 install {USER_FLAG} google-cloud-pipeline-components\n",
         "! pip3 install {USER_FLAG} --upgrade kfp\n",
-        "! pip3 install {USER_FLAG} numpy==1.20.3\n",
-        "! pip3 install {USER_FLAG} --upgrade tensorflow"
+        "! pip3 install {USER_FLAG} numpy\n",
+        "! pip3 install {USER_FLAG} --upgrade tensorflow\n",
+        "! pip3 install {USER_FLAG} --upgrade pillow\n",
+        "! pip3 install {USER_FLAG} --upgrade tf-agents\n",
+        "! pip3 install {USER_FLAG} --upgrade fastapi"
       ]
     },
     {
@@ -287,7 +290,7 @@
         "\n",
         "# Get your Google Cloud project ID from gcloud\n",
         "if not os.getenv(\"IS_TESTING\"):\n",
-        "    shell_output=!gcloud config list --format 'value(core.project)' 2>/dev/null\n",
+        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
         "    PROJECT_ID = shell_output[0]\n",
         "    print(\"Project ID: \", PROJECT_ID)"
       ]
@@ -518,6 +521,7 @@
         "import os\n",
         "import sys\n",
         "\n",
+        "from google.cloud import aiplatform\n",
         "from google_cloud_pipeline_components import aiplatform as gcc_aip\n",
         "from kfp.v2 import compiler, dsl\n",
         "from kfp.v2.google.client import AIPlatformClient"
@@ -561,13 +565,34 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "H3530hdGGilo"
+        "id": "895ac243c125"
       },
       "outputs": [],
       "source": [
         "# Dataset parameters\n",
-        "RAW_DATA_PATH = \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\"  # Location of the MovieLens 100K dataset's \"u.data\" file.\n",
-        "\n",
+        "RAW_DATA_PATH = \"gs://[your-bucket-name]/raw_data/u.data\"   # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "62bfb9a820f6"
+      },
+      "outputs": [],
+      "source": [
+        "# Download the sample data into your RAW_DATA_PATH\n",
+        "! gsutil cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\" $RAW_DATA_PATH"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "H3530hdGGilo"
+      },
+      "outputs": [],
+      "source": [
         "# Pipeline parameters\n",
         "PIPELINE_NAME = \"movielens-pipeline\"  # Pipeline display name.\n",
         "ENABLE_CACHING = False  # Whether to enable execution caching for the pipeline.\n",
@@ -635,7 +660,7 @@
       "source": [
         "#### Run unit tests on the Generator component\n",
         "\n",
-        "Before running the command, fill in `RAW_DATA_PATH` in [`src/generator/test_generator_component.py`](src/generator/test_generator_component.py)."
+        "Before running the command, you should update the `RAW_DATA_PATH` in [`src/generator/test_generator_component.py`](src/generator/test_generator_component.py)."
       ]
     },
     {
@@ -713,12 +738,12 @@
         "TRAINING_ARTIFACTS_DIR = (\n",
         "    f\"{BUCKET_NAME}/artifacts\"  # Root directory for training artifacts.\n",
         ")\n",
-        "TRAINING_REPLICA_COUNT = \"1\"  # Number of replica to run the custom training job.\n",
+        "TRAINING_REPLICA_COUNT = 1  # Number of replica to run the custom training job.\n",
         "TRAINING_MACHINE_TYPE = (\n",
         "    \"n1-standard-4\"  # Type of machine to run the custom training job.\n",
         ")\n",
         "TRAINING_ACCELERATOR_TYPE = \"ACCELERATOR_TYPE_UNSPECIFIED\"  # Type of accelerators to run the custom training job.\n",
-        "TRAINING_ACCELERATOR_COUNT = \"0\"  # Number of accelerators for the custom training job."
+        "TRAINING_ACCELERATOR_COUNT = 0  # Number of accelerators for the custom training job."
       ]
     },
     {
@@ -769,8 +794,12 @@
         "TRAINED_POLICY_DISPLAY_NAME = (\n",
         "    \"movielens-trained-policy\"  # Display name of the uploaded and deployed policy.\n",
         ")\n",
+        "TRAFFIC_SPLIT = {\"0\": 100}\n",
         "ENDPOINT_DISPLAY_NAME = \"movielens-endpoint\"  # Display name of the prediction endpoint.\n",
-        "ENDPOINT_MACHINE_TYPE = \"n1-standard-4\"  # Type of machine of the prediction endpoint."
+        "ENDPOINT_MACHINE_TYPE = \"n1-standard-4\"  # Type of machine of the prediction endpoint.\n",
+        "ENDPOINT_REPLICA_COUNT = 1  # Number of replicas of the prediction endpoint.\n",
+        "ENDPOINT_ACCELERATOR_TYPE = \"ACCELERATOR_TYPE_UNSPECIFIED\"  # Type of accelerators to run the custom training job.\n",
+        "ENDPOINT_ACCELERATOR_COUNT = 0  # Number of accelerators for the custom training job."
       ]
     },
     {
@@ -900,16 +929,17 @@
       },
       "outputs": [],
       "source": [
+        "from google_cloud_pipeline_components.experimental.custom_job import utils\n",
         "from kfp.components import load_component_from_url\n",
         "\n",
         "generate_op = load_component_from_url(\n",
-        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/component.yaml\"\n",
+        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/62a2a7611499490b4b04d731d48a7ba87c2d636f/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/component.yaml\"\n",
         ")\n",
         "ingest_op = load_component_from_url(\n",
-        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml\"\n",
+        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/62a2a7611499490b4b04d731d48a7ba87c2d636f/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml\"\n",
         ")\n",
         "train_op = load_component_from_url(\n",
-        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml\"\n",
+        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/62a2a7611499490b4b04d731d48a7ba87c2d636f/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml\"\n",
         ")\n",
         "\n",
         "\n",
@@ -978,7 +1008,7 @@
         "        bigquery_location=bigquery_location,\n",
         "        bigquery_table_id=bigquery_table_id,\n",
         "    )\n",
-        "\n",
+        "    \n",
         "    # Run the Ingester component.\n",
         "    ingest_task = ingest_op(\n",
         "        project_id=project_id,\n",
@@ -988,7 +1018,16 @@
         "    )\n",
         "\n",
         "    # Run the Trainer component and submit custom job to Vertex AI.\n",
-        "    train_task = train_op(\n",
+        "    # Convert the train_op component into a Vertex AI Custom Job pre-built component\n",
+        "    custom_job_training_op = utils.create_custom_training_job_op_from_component(\n",
+        "        component_spec=train_op,\n",
+        "        replica_count=TRAINING_REPLICA_COUNT,\n",
+        "        machine_type=TRAINING_MACHINE_TYPE,\n",
+        "        accelerator_type=TRAINING_ACCELERATOR_TYPE,\n",
+        "        accelerator_count=TRAINING_ACCELERATOR_COUNT,\n",
+        "    )\n",
+        "\n",
+        "    train_task = custom_job_training_op(\n",
         "        training_artifacts_dir=training_artifacts_dir,\n",
         "        tfrecord_file=ingest_task.outputs[\"tfrecord_file\"],\n",
         "        num_epochs=num_epochs,\n",
@@ -996,28 +1035,10 @@
         "        num_actions=num_actions,\n",
         "        tikhonov_weight=tikhonov_weight,\n",
         "        agent_alpha=agent_alpha,\n",
+        "        project=PROJECT_ID,\n",
+        "        location=REGION,\n",
         "    )\n",
         "\n",
-        "    worker_pool_specs = [\n",
-        "        {\n",
-        "            \"containerSpec\": {\n",
-        "                \"imageUri\": train_task.container.image,\n",
-        "            },\n",
-        "            \"replicaCount\": TRAINING_REPLICA_COUNT,\n",
-        "            \"machineSpec\": {\n",
-        "                \"machineType\": TRAINING_MACHINE_TYPE,\n",
-        "                \"acceleratorType\": TRAINING_ACCELERATOR_TYPE,\n",
-        "                \"acceleratorCount\": TRAINING_ACCELERATOR_COUNT,\n",
-        "            },\n",
-        "        },\n",
-        "    ]\n",
-        "    train_task.custom_job_spec = {\n",
-        "        \"displayName\": train_task.name,\n",
-        "        \"jobSpec\": {\n",
-        "            \"workerPoolSpecs\": worker_pool_specs,\n",
-        "        },\n",
-        "    }\n",
-        "\n",
         "    # Run the Deployer components.\n",
         "    # Upload the trained policy as a model.\n",
         "    model_upload_op = gcc_aip.ModelUploadOp(\n",
@@ -1034,11 +1055,14 @@
         "    # Deploy the uploaded, trained policy to the created endpoint. (This operation\n",
         "    # has to occur after both model uploading and endpoint creation complete.)\n",
         "    gcc_aip.ModelDeployOp(\n",
-        "        project=project_id,\n",
         "        endpoint=endpoint_create_op.outputs[\"endpoint\"],\n",
         "        model=model_upload_op.outputs[\"model\"],\n",
         "        deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,\n",
-        "        machine_type=ENDPOINT_MACHINE_TYPE,\n",
+        "        traffic_split=TRAFFIC_SPLIT,\n",
+        "        dedicated_resources_machine_type=ENDPOINT_MACHINE_TYPE,\n",
+        "        dedicated_resources_accelerator_type=ENDPOINT_ACCELERATOR_TYPE,\n",
+        "        dedicated_resources_accelerator_count=ENDPOINT_ACCELERATOR_COUNT,\n",
+        "        dedicated_resources_min_replica_count=ENDPOINT_REPLICA_COUNT,\n",
         "    )"
       ]
     },
@@ -1053,12 +1077,11 @@
         "# Compile the authored pipeline.\n",
         "compiler.Compiler().compile(pipeline_func=pipeline, package_path=PIPELINE_SPEC_PATH)\n",
         "\n",
-        "# Createa Vertex AI client.\n",
-        "api_client = AIPlatformClient(project_id=PROJECT_ID, region=REGION)\n",
-        "\n",
         "# Create a pipeline run job.\n",
-        "response = api_client.create_run_from_job_spec(\n",
-        "    job_spec_path=PIPELINE_SPEC_PATH,\n",
+        "job = aiplatform.PipelineJob(\n",
+        "    display_name=f\"{PIPELINE_NAME}-startup\",\n",
+        "    template_path=PIPELINE_SPEC_PATH,\n",
+        "    pipeline_root=PIPELINE_ROOT,\n",
         "    parameter_values={\n",
         "        # Pipeline configs\n",
         "        \"project_id\": PROJECT_ID,\n",
@@ -1070,7 +1093,9 @@
         "        \"bigquery_table_id\": BIGQUERY_TABLE_ID,\n",
         "    },\n",
         "    enable_caching=ENABLE_CACHING,\n",
-        ")"
+        ")\n",
+        "\n",
+        "job.run()"
       ]
     },
     {
@@ -1111,7 +1136,11 @@
         "SIMULATOR_SCHEDULE = \"*/5 * * * *\"  # Cloud Scheduler cron job schedule for the Simulator. Eg. \"*/5 * * * *\" means every 5 mins.\n",
         "SIMULATOR_SCHEDULER_MESSAGE = (\n",
         "    \"simulator-message\"  # Cloud Scheduler message for the Simulator.\n",
-        ")"
+        ")\n",
+        "# TF-Agents RL configs\n",
+        "BATCH_SIZE = 8\n",
+        "RANK_K = 20\n",
+        "NUM_ACTIONS = 20"
       ]
     },
     {
@@ -1221,7 +1250,7 @@
       },
       "outputs": [],
       "source": [
-        "endpoints = ! gcloud beta ai endpoints list \\\n",
+        "endpoints = ! gcloud ai endpoints list \\\n",
         "    --region=$REGION \\\n",
         "    --filter=display_name=$ENDPOINT_DISPLAY_NAME\n",
         "print(\"\\n\".join(endpoints), \"\\n\")\n",
@@ -1424,13 +1453,11 @@
       },
       "outputs": [],
       "source": [
-        "from kfp.components import load_component_from_url\n",
-        "\n",
         "ingest_op = load_component_from_url(\n",
-        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml\"\n",
+        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/62a2a7611499490b4b04d731d48a7ba87c2d636f/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml\"\n",
         ")\n",
         "train_op = load_component_from_url(\n",
-        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml\"\n",
+        "    \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/62a2a7611499490b4b04d731d48a7ba87c2d636f/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml\"\n",
         ")\n",
         "\n",
         "\n",
@@ -1481,7 +1508,16 @@
         "    )\n",
         "\n",
         "    # Run the Trainer component and submit custom job to Vertex AI.\n",
-        "    train_task = train_op(\n",
+        "    # Convert the train_op component into a Vertex AI Custom Job pre-built component\n",
+        "    custom_job_training_op = utils.create_custom_training_job_op_from_component(\n",
+        "        component_spec=train_op,\n",
+        "        replica_count=TRAINING_REPLICA_COUNT,\n",
+        "        machine_type=TRAINING_MACHINE_TYPE,\n",
+        "        accelerator_type=TRAINING_ACCELERATOR_TYPE,\n",
+        "        accelerator_count=TRAINING_ACCELERATOR_COUNT,\n",
+        "    )\n",
+        "\n",
+        "    train_task = custom_job_training_op(\n",
         "        training_artifacts_dir=training_artifacts_dir,\n",
         "        tfrecord_file=ingest_task.outputs[\"tfrecord_file\"],\n",
         "        num_epochs=num_epochs,\n",
@@ -1489,28 +1525,10 @@
         "        num_actions=num_actions,\n",
         "        tikhonov_weight=tikhonov_weight,\n",
         "        agent_alpha=agent_alpha,\n",
+        "        project=PROJECT_ID,\n",
+        "        location=REGION,\n",
         "    )\n",
         "\n",
-        "    worker_pool_specs = [\n",
-        "        {\n",
-        "            \"containerSpec\": {\n",
-        "                \"imageUri\": train_task.container.image,\n",
-        "            },\n",
-        "            \"replicaCount\": TRAINING_REPLICA_COUNT,\n",
-        "            \"machineSpec\": {\n",
-        "                \"machineType\": TRAINING_MACHINE_TYPE,\n",
-        "                \"acceleratorType\": TRAINING_ACCELERATOR_TYPE,\n",
-        "                \"acceleratorCount\": TRAINING_ACCELERATOR_COUNT,\n",
-        "            },\n",
-        "        },\n",
-        "    ]\n",
-        "    train_task.custom_job_spec = {\n",
-        "        \"displayName\": train_task.name,\n",
-        "        \"jobSpec\": {\n",
-        "            \"workerPoolSpecs\": worker_pool_specs,\n",
-        "        },\n",
-        "    }\n",
-        "\n",
         "    # Run the Deployer components.\n",
         "    # Upload the trained policy as a model.\n",
         "    model_upload_op = gcc_aip.ModelUploadOp(\n",
@@ -1527,11 +1545,13 @@
         "    # Deploy the uploaded, trained policy to the created endpoint. (This operation\n",
         "    # has to occur after both model uploading and endpoint creation complete.)\n",
         "    gcc_aip.ModelDeployOp(\n",
-        "        project=project_id,\n",
         "        endpoint=endpoint_create_op.outputs[\"endpoint\"],\n",
         "        model=model_upload_op.outputs[\"model\"],\n",
         "        deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,\n",
-        "        machine_type=ENDPOINT_MACHINE_TYPE,\n",
+        "        dedicated_resources_machine_type=ENDPOINT_MACHINE_TYPE,\n",
+        "        dedicated_resources_accelerator_type=ENDPOINT_ACCELERATOR_TYPE,\n",
+        "        dedicated_resources_accelerator_count=ENDPOINT_ACCELERATOR_COUNT,\n",
+        "        dedicated_resources_min_replica_count=ENDPOINT_REPLICA_COUNT,\n",
         "    )"
       ]
     },

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/component.yaml

@@ -39,14 +39,15 @@ outputs:
 - {name: bigquery_table_id, type: String}
 implementation:
   container:
-    image: tensorflow/tensorflow:2.5.0
+    image: python:3.7
     command:
     - sh
     - -c
     - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location
-      'google-cloud-bigquery==2.20.0' 'tensorflow==2.5.0' 'tf-agents==0.8.0' || PIP_DISABLE_PIP_VERSION_CHECK=1
-      python3 -m pip install --quiet --no-warn-script-location 'google-cloud-bigquery==2.20.0'
-      'tensorflow==2.5.0' 'tf-agents==0.8.0' --user) && "$0" "$@"
+      'google-cloud-bigquery==2.20.0' 'pillow' 'tensorflow==2.5.0' 'tf-agents==0.8.0'
+      || PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location
+      'google-cloud-bigquery==2.20.0' 'pillow' 'tensorflow==2.5.0' 'tf-agents==0.8.0'
+      --user) && "$0" "$@"
     - sh
     - -ec
     - |
@@ -296,7 +297,8 @@ implementation:
 
       def _serialize_str(str_value: str) -> str:
           if not isinstance(str_value, str):
-              raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value))))
+              raise TypeError('Value "{}" has type "{}" instead of str.'.format(
+                  str(str_value), str(type(str_value))))
           return str_value
 
       import argparse

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml

@@ -20,7 +20,7 @@ outputs:
 - {name: tfrecord_file, type: String}
 implementation:
   container:
-    image: tensorflow/tensorflow:2.5.0
+    image: python:3.7
     command:
     - sh
     - -c
@@ -187,7 +187,8 @@ implementation:
 
       def _serialize_str(str_value: str) -> str:
           if not isinstance(str_value, str):
-              raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value))))
+              raise TypeError('Value "{}" has type "{}" instead of str.'.format(
+                  str(str_value), str(type(str_value))))
           return str_value
 
       import argparse

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/logger/requirements.txt

@@ -1,3 +1,4 @@
 google-cloud-bigquery==2.20.0
-tensorflow==2.5.2
+tensorflow==2.7.2
+pillow==9.0.1
 tf-agents==0.8.0

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/prediction_container/requirements.txt

@@ -1,2 +1,4 @@
+google-cloud-pubsub==2.5.0
+pillow==9.0.1
 tf-agents==0.8.0
-tensorflow==2.5.2
+tensorflow==2.7.2

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/simulator/requirements.txt

@@ -0,0 +1,5 @@
+dataclasses==0.6
+google-cloud-aiplatform==1.8.1
+tensorflow==2.7.2
+pillow==9.0.1
+tf-agents==0.8.0

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml

@@ -27,14 +27,14 @@ outputs:
 - {name: training_artifacts_dir, type: String}
 implementation:
   container:
-    image: tensorflow/tensorflow:2.5.0
+    image: python:3.7
     command:
     - sh
     - -c
     - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location
-      'tensorflow==2.5.0' 'tf-agents==0.8.0' || PIP_DISABLE_PIP_VERSION_CHECK=1 python3
-      -m pip install --quiet --no-warn-script-location 'tensorflow==2.5.0' 'tf-agents==0.8.0'
-      --user) && "$0" "$@"
+      'tensorflow==2.5.0' 'tf-agents==0.8.0' 'Pillow' || PIP_DISABLE_PIP_VERSION_CHECK=1
+      python3 -m pip install --quiet --no-warn-script-location 'tensorflow==2.5.0'
+      'tf-agents==0.8.0' 'Pillow' --user) && "$0" "$@"
     - sh
     - -ec
     - |
@@ -270,7 +270,8 @@ implementation:
 
       def _serialize_str(str_value: str) -> str:
           if not isinstance(str_value, str):
-              raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value))))
+              raise TypeError('Value "{}" has type "{}" instead of str.'.format(
+                  str(str_value), str(type(str_value))))
           return str_value
 
       import argparse

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/step_by_step_sdk_tf_agents_bandits_movie_recommendation/src/tests/test_task.py

@@ -22,13 +22,13 @@ from src.training import task
 
 
 # Paths and configurations
-DATA_PATH = "gs://[your-bucket-name]/[your-dataset-dir]/u.data"  # FILL IN
+DATA_PATH = "gs://[your-bucket-name]/artifacts/u.data"  # FILL IN
 ROOT_DIR = "gs://[your-bucket-name]/artifacts"  # FILL IN
 ARTIFACTS_DIR = "gs://[your-bucket-name]/artifacts"  # FILL IN
 PROFILER_DIR = "gs://[your-bucket-name]/profiler"  # FILL IN
 HPTUNING_RESULT_DIR = "[your-hptuning-result-dir]/"  # FILL IN
 HPTUNING_RESULT_PATH = os.path.join(HPTUNING_RESULT_DIR,
-                                    "[your-file-name].json")  # FILL IN
+                                    "result.json")  # FILL IN
 RAW_BUCKET_NAME = "[your-hptuning-result-bucket-name]"  # FILL IN
 
 # Hyperparameters

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/requirements.txt

@@ -1 +1 @@
-tensorflow==2.5.2
+tensorflow==2.7.2

community-content/tf_keras_text_classification_distributed_single_worker_gpus_with_gcloud_local_run_and_vertex_sdk/trainer/requirements.txt

@@ -1 +1 @@
-tensorflow==2.5.2
+tensorflow==2.7.2

community-content/tf_keras_text_classification_distributed_single_worker_gpus_with_gcloud_local_run_and_vertex_sdk/vertex_training_with_local_mode_container.ipynb

@@ -8,7 +8,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2021 Google LLC\n",
+        "# Copyright 2022 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -113,8 +113,8 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud beta ai custom-jobs local-run \\\n",
-        "  --base-image=$BASE_IMAGE_URI \\\n",
+        "! gcloud ai custom-jobs local-run \\\n",
+        "  --executor-image-uri=$BASE_IMAGE_URI \\\n",
         "  --script=$SCRIPT_PATH \\\n",
         "  --output-image-uri=$OUTPUT_IMAGE_NAME \\\n",
         "  -- \\\n",

notebooks/README.md

@@ -0,0 +1,5 @@
+The [official](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/notebooks/official) folder contains notebooks organized by Google Cloud product. These are tested weekly and maintained by Google.
+
+The [community](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/notebooks/community) folder contains notebooks that may be created by Google or external contributors. They are not necessary maintained.
+
+Contributions to the repo should use the [notebook template](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/notebook_template.ipynb) as a starting point.

notebooks/community/CODEOWNERS

@@ -3,16 +3,34 @@
 # @global-owner1 and @global-owner2 will be requested for
 # review when someone opens a pull request.
 
-/sdk/sdk_* @aferlitsch
-/gapic @aferlitsch
-/ml_ops @aferlitsch
-/model_monitoring/* @mco
+/sdk/sdk_* @andrewferlitsch
+/gapic @andrewferlitsch
+/gapic/custom/showcase_custom_image_classification_online_explain_example_based_api.ipynb @inardini
+/ml_ops @andrewferlitsch
+/model_monitoring/* @andrewferlitsch
 /structured_data/rapid_prototyping_* @rafael-carvalho
 
-/managed_notebooks/ @notebooks-team
-/sdk/SDK_FBProphet_Forecasting_Online.ipynb			@brianchunkang
+/managed_notebooks/
+/bigquery_ml/ @polong
+/sdk/SDK_FBProphet_Forecasting_Online.ipynb @brianchunkang
 /pipelines/google_cloud_pipeline_components_TPU_model_train_upload_deploy.ipynb @brianchunkang
-/sdk/SDK_AutoML_Forecasting_Model_Training_Example.ipynb @thehardikv
-/sdk/sdk_automl_forecasting_evaluating_a_model.ipynb @thehardikv
-/matching_engine @yinghsienwu
-/neo4j @benofben @htappen
+/explainable_ai/SDK_Custom_Container_XAI.ipynb @brianchunkang
+/matching_engine/sdk_matching_engine_for_indexing.ipynb @ivanmkc
+/matching_engine/matching_engine_for_indexing.ipynb @yinghsienwu
+/matching_engine/stream_update_for_matching_engine.ipynb @peterping666
+/sdk/pytorch_lightning_custom_container_training.ipynb @brianchunkang
+/tensorboard @yfang1
+/feature_store @nayaknishant @morgandu
+/prediction @googleapis/vertex-prediction-team
+/vertex_endpoints/tf_hub_obj_detection/deploy_tfhub_object_detection_on_vertex_endpoints.ipynb @entrpn
+/vertex_endpoints/nvidia-triton/nvidia-triton-custom-container-prediction.ipynb @RajeshThallam
+/vertex_endpoints/optimized_tensorflow_runtime @vlasenkoalexey
+/notebooks/community/ml_ops/stage2/get_started_with_visionapi_and_automl.ipynb @mansari
+/notebooks/community/neo4j/graph_paysim.ipynb @benofben @laeg
+/notebooks/community/ml_ops/stage1/get_started_with_visionapi_and_vertex_datasets.ipynb @mansari
+/notebooks/community/pipelines/google_cloud_pipeline_components_bqml_pipeline_demand_forecasting.ipynb @inardini
+/notebooks/community/ml_ops/stage2/get_started_vertex_hpt_r_kernel.ipynb @fhirschmann
+/notebooks/community/ml_ops/stage2/get_started_vertex_training_r_using_r_kernel.ipynb @fhirschmann
+/notebooks/community/vertex-ai-samples/notebooks/community/model_registry/vertex_ai_model_registry_bqml_custom_model_versioning.ipynb @inardini
+/notebooks/community/vertex-ai-samples/notebooks/community/model_registry/vertex_ai_model_registry_automl_model_versioning.ipynb @inardini
+/notebooks/community/vizier/conversions_vertex_vizier_and_open_source_vizier.ipynb @halio-g

notebooks/community/feature_store/gapic-feature-store.ipynb

@@ -171,7 +171,7 @@
       },
       "outputs": [],
       "source": [
-        "! pip install {USER_FLAG} --upgrade git+https://github.com/googleapis/python-aiplatform.git@v1.6.0"
+        "! pip install {USER_FLAG} --upgrade google-cloud-aiplatform"
       ]
     },
     {
@@ -266,7 +266,7 @@
         "\n",
         "# Get your Google Cloud project ID from gcloud\n",
         "if not os.getenv(\"IS_TESTING\"):\n",
-        "    shell_output=!gcloud config list --format 'value(core.project)' 2>/dev/null\n",
+        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
         "    PROJECT_ID = shell_output[0]\n",
         "    print(\"Project ID: \", PROJECT_ID)"
       ]
@@ -292,6 +292,37 @@
         "    PROJECT_ID = \"python-docs-samples-tests\"  # @param {type:\"string\"}"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d9f118b92c74"
+      },
+      "source": [
+        "#### UUID\n",
+        "\n",
+        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3ee72715c0fd"
+      },
+      "outputs": [],
+      "source": [
+        "import random\n",
+        "import string\n",
+        "\n",
+        "\n",
+        "# Generate a uuid of a specifed length(default=8)\n",
+        "def generate_uuid(length: int = 8) -> str:\n",
+        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
+        "\n",
+        "\n",
+        "UUID = generate_uuid()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -478,7 +509,6 @@
         "from google.cloud.aiplatform_v1.types import \\\n",
         "    featurestore_service as featurestore_service_pb2\n",
         "from google.cloud.aiplatform_v1.types import io as io_pb2\n",
-        "from google.protobuf.duration_pb2 import Duration\n",
         "\n",
         "# Create admin_client for CRUD and data_client for reading feature values.\n",
         "admin_client = FeaturestoreServiceClient(client_options={\"api_endpoint\": API_ENDPOINT})\n",
@@ -542,30 +572,23 @@
       },
       "outputs": [],
       "source": [
-        "FEATURESTORE_ID = \"movie_prediction\"\n",
-        "create_lro = admin_client.create_featurestore(\n",
-        "    featurestore_service_pb2.CreateFeaturestoreRequest(\n",
-        "        parent=BASE_RESOURCE_PATH,\n",
-        "        featurestore_id=FEATURESTORE_ID,\n",
-        "        featurestore=featurestore_pb2.Featurestore(\n",
-        "            online_serving_config=featurestore_pb2.Featurestore.OnlineServingConfig(\n",
-        "                fixed_node_count=1\n",
+        "FEATURESTORE_ID = f\"movie_prediction_{UUID}\"\n",
+        "try:\n",
+        "    create_lro = admin_client.create_featurestore(\n",
+        "        featurestore_service_pb2.CreateFeaturestoreRequest(\n",
+        "            parent=BASE_RESOURCE_PATH,\n",
+        "            featurestore_id=FEATURESTORE_ID,\n",
+        "            featurestore=featurestore_pb2.Featurestore(\n",
+        "                online_serving_config=featurestore_pb2.Featurestore.OnlineServingConfig(\n",
+        "                    fixed_node_count=1\n",
+        "                ),\n",
         "            ),\n",
-        "        ),\n",
+        "        )\n",
         "    )\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "57V8eVcB5VFZ"
-      },
-      "outputs": [],
-      "source": [
-        "# Wait for LRO to finish and get the LRO result.\n",
-        "print(create_lro.result())"
+        "    # Wait for LRO to finish and get the LRO result.\n",
+        "    print(create_lro.result())\n",
+        "except Exception as e:\n",
+        "    print(e)"
       ]
     },
     {
@@ -574,7 +597,7 @@
         "id": "ag8pCQ7rNjVf"
       },
       "source": [
-        "You can use [GetFeaturestore](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1beta1#google.cloud.aiplatform.v1beta1.FeaturestoreService.GetFeaturestore) or [ListFeaturestores](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1beta1#google.cloud.aiplatform.v1beta1.FeaturestoreService.ListFeaturestores) to check if the Featurestore was successfully created. The following example gets the details of the Featurestore.\n"
+        "You can use [GetFeaturestore](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1#google.cloud.aiplatform.v1.FeaturestoreService.GetFeaturestore) or [ListFeaturestores](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1#google.cloud.aiplatform.v1.FeaturestoreService.ListFeaturestores) to check if the Featurestore was successfully created. The following example gets the details of the Featurestore.\n"
       ]
     },
     {
@@ -590,6 +613,41 @@
         ")"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "018ab19d934f"
+      },
+      "source": [
+        "Auto scaling is available in v1 since v1.11. Below is the example for the `CreateFeaturestoreRequest` with auto-scaling, use it with `aiplatform_v1.FeaturestoreServiceClient` to create Featurestore:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aea39718b5d3"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud.aiplatform_v1.types import \\\n",
+        "    featurestore as v1_featurestore_pb2\n",
+        "from google.cloud.aiplatform_v1.types import \\\n",
+        "    featurestore_service as v1_featurestore_service_pb2\n",
+        "\n",
+        "create_featurestore_request = v1_featurestore_service_pb2.CreateFeaturestoreRequest(\n",
+        "    parent=BASE_RESOURCE_PATH,\n",
+        "    featurestore_id=FEATURESTORE_ID,\n",
+        "    featurestore=v1_featurestore_pb2.Featurestore(\n",
+        "        online_serving_config=v1_featurestore_pb2.Featurestore.OnlineServingConfig(\n",
+        "            scaling=v1_featurestore_pb2.Featurestore.OnlineServingConfig.Scaling(\n",
+        "                min_node_count=1, max_node_count=5\n",
+        "            )\n",
+        "        ),\n",
+        "    ),\n",
+        ")"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -607,18 +665,20 @@
       },
       "outputs": [],
       "source": [
-        "users_entity_type_lro = admin_client.create_entity_type(\n",
-        "    featurestore_service_pb2.CreateEntityTypeRequest(\n",
-        "        parent=admin_client.featurestore_path(PROJECT_ID, REGION, FEATURESTORE_ID),\n",
-        "        entity_type_id=\"users\",\n",
-        "        entity_type=entity_type_pb2.EntityType(\n",
-        "            description=\"Users entity\",\n",
-        "        ),\n",
+        "try:\n",
+        "    users_entity_type_lro = admin_client.create_entity_type(\n",
+        "        featurestore_service_pb2.CreateEntityTypeRequest(\n",
+        "            parent=admin_client.featurestore_path(PROJECT_ID, REGION, FEATURESTORE_ID),\n",
+        "            entity_type_id=\"users\",\n",
+        "            entity_type=entity_type_pb2.EntityType(\n",
+        "                description=\"Users entity\",\n",
+        "            ),\n",
+        "        )\n",
         "    )\n",
-        ")\n",
-        "\n",
-        "# Similarly, wait for EntityType creation operation.\n",
-        "print(users_entity_type_lro.result())"
+        "    # Similarly, wait for EntityType creation operation.\n",
+        "    print(users_entity_type_lro.result())\n",
+        "except Exception as e:\n",
+        "    print(e)"
       ]
     },
     {
@@ -630,16 +690,73 @@
       "outputs": [],
       "source": [
         "# Create movies entity type without a monitoring configuration.\n",
-        "movies_entity_type_lro = admin_client.create_entity_type(\n",
-        "    featurestore_service_pb2.CreateEntityTypeRequest(\n",
-        "        parent=admin_client.featurestore_path(PROJECT_ID, REGION, FEATURESTORE_ID),\n",
-        "        entity_type_id=\"movies\",\n",
-        "        entity_type=entity_type_pb2.EntityType(description=\"Movies entity\"),\n",
+        "try:\n",
+        "    movies_entity_type_lro = admin_client.create_entity_type(\n",
+        "        featurestore_service_pb2.CreateEntityTypeRequest(\n",
+        "            parent=admin_client.featurestore_path(PROJECT_ID, REGION, FEATURESTORE_ID),\n",
+        "            entity_type_id=\"movies\",\n",
+        "            entity_type=entity_type_pb2.EntityType(description=\"Movies entity\"),\n",
+        "        )\n",
         "    )\n",
+        "\n",
+        "    # Similarly, wait for EntityType creation operation.\n",
+        "    print(movies_entity_type_lro.result())\n",
+        "except Exception as e:\n",
+        "    print(e)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dPkT7KDuEvWv"
+      },
+      "source": [
+        "Feature [monitoring](https://cloud.google.com/vertex-ai/docs/featurestore/monitoring) is in preview, so you need to use v1 Python. Import feature analysis is only available through SDK for now."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9kiqrBN6E28r"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud.aiplatform_v1 import \\\n",
+        "    FeaturestoreServiceClient as v1_FeaturestoreServiceClient\n",
+        "from google.cloud.aiplatform_v1.types import entity_type as v1_entity_type_pb2\n",
+        "from google.cloud.aiplatform_v1.types import \\\n",
+        "    featurestore_monitoring as v1_featurestore_monitoring_pb2\n",
+        "from google.cloud.aiplatform_v1.types import \\\n",
+        "    featurestore_service as v1_featurestore_service_pb2\n",
+        "\n",
+        "v1_admin_client = v1_FeaturestoreServiceClient(\n",
+        "    client_options={\"api_endpoint\": API_ENDPOINT}\n",
         ")\n",
         "\n",
-        "# Similarly, wait for EntityType creation operation.\n",
-        "print(movies_entity_type_lro.result())"
+        "# Enable import feature analysis for users entity type.\n",
+        "# All Features belonging to this EntityType will by default inherit the monitoring config.\n",
+        "v1_admin_client.update_entity_type(\n",
+        "    v1_featurestore_service_pb2.UpdateEntityTypeRequest(\n",
+        "        entity_type=v1_entity_type_pb2.EntityType(\n",
+        "            name=admin_client.entity_type_path(\n",
+        "                PROJECT_ID, REGION, FEATURESTORE_ID, \"users\"\n",
+        "            ),\n",
+        "            monitoring_config=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig(\n",
+        "                import_features_analysis=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.ImportFeaturesAnalysis(\n",
+        "                    anomaly_detection_baseline=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.ImportFeaturesAnalysis.Baseline.LATEST_STATS,\n",
+        "                    state=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.ImportFeaturesAnalysis.State.ENABLED,\n",
+        "                ),\n",
+        "                numerical_threshold_config=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.ThresholdConfig(\n",
+        "                    value=0.001,\n",
+        "                ),\n",
+        "                categorical_threshold_config=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.ThresholdConfig(\n",
+        "                    value=0.001,\n",
+        "                ),\n",
+        "            ),\n",
+        "        ),\n",
+        "    )\n",
+        ")"
       ]
     },
     {
@@ -648,7 +765,9 @@
         "id": "85b1f59fbf6d"
       },
       "source": [
-        "Feature [monitoring](https://cloud.google.com/vertex-ai/docs/featurestore/monitoring) is in preview, so you need to use v1beta1 Python. The easiest way to set this for now is using [console UI](https://console.cloud.google.com/vertex-ai/features). For completeness, below is example to do this using v1beta1 SDK\n"
+        "The easiest way to set up snapshot analysis for now is using [console UI](https://console.cloud.google.com/vertex-ai/features). For completeness, below is example to do this using v1 SDK.\n",
+        "\n",
+        "You can view monitoring statistics on [console UI](https://console.cloud.google.com/vertex-ai/features)."
       ]
     },
     {
@@ -659,30 +778,36 @@
       },
       "outputs": [],
       "source": [
-        "from google.cloud.aiplatform_v1beta1 import \\\n",
-        "    FeaturestoreServiceClient as v1beta1_FeaturestoreServiceClient\n",
-        "from google.cloud.aiplatform_v1beta1.types import \\\n",
-        "    entity_type as v1beta1_entity_type_pb2\n",
-        "from google.cloud.aiplatform_v1beta1.types import \\\n",
-        "    featurestore_monitoring as v1beta1_featurestore_monitoring_pb2\n",
-        "from google.cloud.aiplatform_v1beta1.types import \\\n",
-        "    featurestore_service as v1beta1_featurestore_service_pb2\n",
+        "from google.cloud.aiplatform_v1 import \\\n",
+        "    FeaturestoreServiceClient as v1_FeaturestoreServiceClient\n",
+        "from google.cloud.aiplatform_v1.types import entity_type as v1_entity_type_pb2\n",
+        "from google.cloud.aiplatform_v1.types import \\\n",
+        "    featurestore_monitoring as v1_featurestore_monitoring_pb2\n",
+        "from google.cloud.aiplatform_v1.types import \\\n",
+        "    featurestore_service as v1_featurestore_service_pb2\n",
         "\n",
-        "v1beta1_admin_client = v1beta1_FeaturestoreServiceClient(\n",
+        "v1_admin_client = v1_FeaturestoreServiceClient(\n",
         "    client_options={\"api_endpoint\": API_ENDPOINT}\n",
         ")\n",
         "\n",
-        "# Enable monitoring for users entity type.\n",
+        "# Enable snapshot analysis for users entity type.\n",
         "# All Features belonging to this EntityType will by default inherit the monitoring config.\n",
-        "v1beta1_admin_client.update_entity_type(\n",
-        "    v1beta1_featurestore_service_pb2.UpdateEntityTypeRequest(\n",
-        "        entity_type=v1beta1_entity_type_pb2.EntityType(\n",
+        "v1_admin_client.update_entity_type(\n",
+        "    v1_featurestore_service_pb2.UpdateEntityTypeRequest(\n",
+        "        entity_type=v1_entity_type_pb2.EntityType(\n",
         "            name=admin_client.entity_type_path(\n",
         "                PROJECT_ID, REGION, FEATURESTORE_ID, \"users\"\n",
         "            ),\n",
-        "            monitoring_config=v1beta1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig(\n",
-        "                snapshot_analysis=v1beta1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.SnapshotAnalysis(\n",
-        "                    monitoring_interval=Duration(seconds=86400),  # 1 day\n",
+        "            monitoring_config=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig(\n",
+        "                snapshot_analysis=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.SnapshotAnalysis(\n",
+        "                    monitoring_interval_days=1,  # 1 day\n",
+        "                    staleness_days=30,\n",
+        "                ),\n",
+        "                numerical_threshold_config=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.ThresholdConfig(\n",
+        "                    value=0.001,\n",
+        "                ),\n",
+        "                categorical_threshold_config=v1_featurestore_monitoring_pb2.FeaturestoreMonitoringConfig.ThresholdConfig(\n",
+        "                    value=0.001,\n",
         "                ),\n",
         "            ),\n",
         "        ),\n",
@@ -708,32 +833,40 @@
       "outputs": [],
       "source": [
         "# Create features for the 'users' entity.\n",
-        "admin_client.batch_create_features(\n",
-        "    parent=admin_client.entity_type_path(PROJECT_ID, REGION, FEATURESTORE_ID, \"users\"),\n",
-        "    requests=[\n",
-        "        featurestore_service_pb2.CreateFeatureRequest(\n",
-        "            feature=feature_pb2.Feature(\n",
-        "                value_type=feature_pb2.Feature.ValueType.INT64,\n",
-        "                description=\"User age\",\n",
-        "            ),\n",
-        "            feature_id=\"age\",\n",
+        "try:\n",
+        "    admin_client.batch_create_features(\n",
+        "        parent=admin_client.entity_type_path(\n",
+        "            PROJECT_ID, REGION, FEATURESTORE_ID, \"users\"\n",
         "        ),\n",
-        "        featurestore_service_pb2.CreateFeatureRequest(\n",
-        "            feature=feature_pb2.Feature(\n",
-        "                value_type=feature_pb2.Feature.ValueType.STRING,\n",
-        "                description=\"User gender\",\n",
+        "        requests=[\n",
+        "            featurestore_service_pb2.CreateFeatureRequest(\n",
+        "                feature=feature_pb2.Feature(\n",
+        "                    value_type=feature_pb2.Feature.ValueType.INT64,\n",
+        "                    description=\"User age\",\n",
+        "                    disable_monitoring=False,\n",
+        "                ),\n",
+        "                feature_id=\"age\",\n",
         "            ),\n",
-        "            feature_id=\"gender\",\n",
-        "        ),\n",
-        "        featurestore_service_pb2.CreateFeatureRequest(\n",
-        "            feature=feature_pb2.Feature(\n",
-        "                value_type=feature_pb2.Feature.ValueType.STRING_ARRAY,\n",
-        "                description=\"An array of genres that this user liked\",\n",
+        "            featurestore_service_pb2.CreateFeatureRequest(\n",
+        "                feature=feature_pb2.Feature(\n",
+        "                    value_type=feature_pb2.Feature.ValueType.STRING,\n",
+        "                    description=\"User gender\",\n",
+        "                    # Default is False. If True, Feature 'gender' monitoring analysis is disabled.\n",
+        "                    disable_monitoring=True,\n",
+        "                ),\n",
+        "                feature_id=\"gender\",\n",
         "            ),\n",
-        "            feature_id=\"liked_genres\",\n",
-        "        ),\n",
-        "    ],\n",
-        ").result()"
+        "            featurestore_service_pb2.CreateFeatureRequest(\n",
+        "                feature=feature_pb2.Feature(\n",
+        "                    value_type=feature_pb2.Feature.ValueType.STRING_ARRAY,\n",
+        "                    description=\"An array of genres that this user liked\",\n",
+        "                ),\n",
+        "                feature_id=\"liked_genres\",\n",
+        "            ),\n",
+        "        ],\n",
+        "    ).result()\n",
+        "except Exception as e:\n",
+        "    print(e)"
       ]
     },
     {
@@ -745,33 +878,37 @@
       "outputs": [],
       "source": [
         "# Create features for movies type.\n",
-        "# 'title' Feature enables monitoring.\n",
-        "admin_client.batch_create_features(\n",
-        "    parent=admin_client.entity_type_path(PROJECT_ID, REGION, FEATURESTORE_ID, \"movies\"),\n",
-        "    requests=[\n",
-        "        featurestore_service_pb2.CreateFeatureRequest(\n",
-        "            feature=feature_pb2.Feature(\n",
-        "                value_type=feature_pb2.Feature.ValueType.STRING,\n",
-        "                description=\"The title of the movie\",\n",
-        "            ),\n",
-        "            feature_id=\"title\",\n",
+        "try:\n",
+        "    admin_client.batch_create_features(\n",
+        "        parent=admin_client.entity_type_path(\n",
+        "            PROJECT_ID, REGION, FEATURESTORE_ID, \"movies\"\n",
         "        ),\n",
-        "        featurestore_service_pb2.CreateFeatureRequest(\n",
-        "            feature=feature_pb2.Feature(\n",
-        "                value_type=feature_pb2.Feature.ValueType.STRING,\n",
-        "                description=\"The genres of the movie\",\n",
+        "        requests=[\n",
+        "            featurestore_service_pb2.CreateFeatureRequest(\n",
+        "                feature=feature_pb2.Feature(\n",
+        "                    value_type=feature_pb2.Feature.ValueType.STRING,\n",
+        "                    description=\"The title of the movie\",\n",
+        "                ),\n",
+        "                feature_id=\"title\",\n",
         "            ),\n",
-        "            feature_id=\"genres\",\n",
-        "        ),\n",
-        "        featurestore_service_pb2.CreateFeatureRequest(\n",
-        "            feature=feature_pb2.Feature(\n",
-        "                value_type=feature_pb2.Feature.ValueType.DOUBLE,\n",
-        "                description=\"The average rating for the movie, range is [1.0-5.0]\",\n",
+        "            featurestore_service_pb2.CreateFeatureRequest(\n",
+        "                feature=feature_pb2.Feature(\n",
+        "                    value_type=feature_pb2.Feature.ValueType.STRING,\n",
+        "                    description=\"The genres of the movie\",\n",
+        "                ),\n",
+        "                feature_id=\"genres\",\n",
         "            ),\n",
-        "            feature_id=\"average_rating\",\n",
-        "        ),\n",
-        "    ],\n",
-        ").result()"
+        "            featurestore_service_pb2.CreateFeatureRequest(\n",
+        "                feature=feature_pb2.Feature(\n",
+        "                    value_type=feature_pb2.Feature.ValueType.DOUBLE,\n",
+        "                    description=\"The average rating for the movie, range is [1.0-5.0]\",\n",
+        "                ),\n",
+        "                feature_id=\"average_rating\",\n",
+        "            ),\n",
+        "        ],\n",
+        "    ).result()\n",
+        "except Exception as e:\n",
+        "    print(e)"
       ]
     },
     {
@@ -782,8 +919,8 @@
       "source": [
         "## Search created features\n",
         "\n",
-        "While the [ListFeatures](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1beta1#google.cloud.aiplatform.v1beta1.FeaturestoreService.ListFeatures) method allows you to easily view all features of a single\n",
-        "entity type, the [SearchFeatures](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1beta1#google.cloud.aiplatform.v1beta1.FeaturestoreService.SearchFeatures) method searches across all featurestores\n",
+        "While the [ListFeatures](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1#google.cloud.aiplatform.v1.FeaturestoreService.ListFeatures) method allows you to easily view all features of a single\n",
+        "entity type, the [SearchFeatures](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1#google.cloud.aiplatform.v1.FeaturestoreService.SearchFeatures) method searches across all featurestores\n",
         "and entity types in a given location (such as `us-central1`). This can help you discover features that were created by someone else.\n",
         "\n",
         "You can query based on feature properties including feature ID, entity type ID,\n",
@@ -985,6 +1122,8 @@
         "    ],\n",
         "    feature_time_field=\"update_time\",\n",
         "    worker_count=1,\n",
+        "    # Default is False. If True, the import feature analysis won't happen for this specific operation.\n",
+        "    disable_ingestion_analysis=False,\n",
         ")"
       ]
     },
@@ -1095,7 +1234,7 @@
       },
       "source": [
         "The\n",
-        "[Online Serving APIs](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1beta1#featurestoreonlineservingservice)\n",
+        "[Online Serving APIs](https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1#featurestoreonlineservingservice)\n",
         "lets you serve feature values for small batches of entities. It's designed for latency-sensitive service, such as online model prediction. For example, for a movie service, you might want to quickly shows movies that the current user would most likely watch by using online predictions."
       ]
     },
@@ -1393,9 +1532,7 @@
   ],
   "metadata": {
     "colab": {
-      "collapsed_sections": [
-        "ze4-nDLfK4pw"
-      ],
+      "collapsed_sections": [],
       "name": "gapic-feature-store.ipynb",
       "toc_visible": true
     },

notebooks/community/managed_notebooks/pricing_optimization/pricing-optimization.ipynb

@@ -1,819 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d1cc1c1fa076"
-      },
-      "source": [
-        "# Pricing Optimization \n",
-        "## Table of contents\n",
-        "* [Overview](#section-1)\n",
-        "* [Dataset](#section-2)\n",
-        "* [Objective](#section-3)\n",
-        "* [Costs](#section-4)\n",
-        "* [Create a BigQuery dataset](#section-5)\n",
-        "* [Load the dataset from Cloud Storage](#section-6)\n",
-        "* [Data Analysis](#section-7)\n",
-        "* [Preprocess the data for training](#section-8)\n",
-        "* [Train the model using BigQuery ML](#section-9)\n",
-        "* [Generate forecasts from the model](#section-10)\n",
-        "* [Interpret the results to choose the best price](#section-11)\n",
-        "* [Clean Up](#section-12)\n",
-        "\n",
-        "\n",
-        "## Overview\n",
-        "<a name=\"section-1\"></a>\n",
-        "This notebook demonstrates analysis of pricing optimization on [CDM Pricing Data](https://github.com/trifacta/trifacta-google-cloud/tree/main/design-pattern-pricing-optimization) and automating the workflow using Vertex AI Workbench's managed notebooks.\n",
-        "\n",
-        "<b>Note</b>: This notebook is designed to run on managed notebooks instance of Vertex AI Workbench. Some components of this notebook may not work in other notebook environments.\n",
-        "\n",
-        "## Dataset\n",
-        "<a name=\"section-2\"></a>\n",
-        "The dataset used in this notebook is a part of the [CDM Pricing Data](https://github.com/trifacta/trifacta-google-cloud/blob/main/design-pattern-pricing-optimization/CDM_Pricing_large_table.csv) which consists of products sales information on specified dates.\n",
-        "\n",
-        "## Objective\n",
-        "<a name=\"section-3\"></a>\n",
-        "The objective of this notebook is to build a pricing optimization model using Vertex AI on GCP. The following steps have been followed in this usecase :  \n",
-        "\n",
-        "- Load the required dataset from a Cloud Storage bucket.\n",
-        "- Analyze the fields present in the dataset.\n",
-        "- Process the data to build a model.\n",
-        "- Build a BigQuery ML forecast model on the processed data.\n",
-        "- Get forecasted values from the BigQuery ML model.\n",
-        "- Interpret the forecasts to identify best prices.\n",
-        "- Clean up.\n",
-        "\n",
-        "\n",
-        "## Costs\n",
-        "<a name=\"section-4\"></a>\n",
-        "This tutorial uses the following billable components of Google Cloud:\n",
-        "\n",
-        "- Vertex AI\n",
-        "- Bigquery\n",
-        "- Cloud Storage\n",
-        "\n",
-        "\n",
-        "Learn about [Vertex AI\n",
-        "pricing](https://cloud.google.com/vertex-ai/pricing), [Bigquery pricing](https://cloud.google.com/bigquery/pricing) and [Cloud Storage\n",
-        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
-        "Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage.\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5ed1f5e85640"
-      },
-      "source": [
-        "#### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "c3f30148b66d"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "PROJECT_ID = \"\"\n",
-        "\n",
-        "# Get your Google Cloud project ID from gcloud\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID: \", PROJECT_ID)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "750bf2883c2d"
-      },
-      "source": [
-        "Otherwise, set your project ID here."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "3c6db1ca88b9"
-      },
-      "outputs": [],
-      "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
-        "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2a1c270c7d34"
-      },
-      "source": [
-        "### Import the required libraries and define constants\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "acc6fac1fa55"
-      },
-      "outputs": [],
-      "source": [
-        "import matplotlib.pyplot as plt\n",
-        "import pandas as pd\n",
-        "import seaborn as sns\n",
-        "from google.cloud import bigquery\n",
-        "from google.cloud.bigquery import Client"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a06006dff8f9"
-      },
-      "outputs": [],
-      "source": [
-        "DATASET = \"[your-bigquery-dataset-id]\"  # set the Bigquery dataset-id\n",
-        "TRAINING_DATA_TABLE = \"[your-bigquery-table-id-to-store-the-training-data]\"  # set the Bigquery table-id to store the training data"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "016c3d47cc69"
-      },
-      "source": [
-        "## Create a BigQuery dataset\n",
-        "<a name=\"section-5\"></a>\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "12ccd8d7956e"
-      },
-      "source": [
-        "#@bigquery\n",
-        "-- create a dataset in BigQuery\n",
-        "\n",
-        "CREATE SCHEMA pricing_optimization\n",
-        "OPTIONS(\n",
-        "  location=\"us\"\n",
-        "  )"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c106b978a79b"
-      },
-      "source": [
-        "## Load the dataset from Cloud Storage\n",
-        "<a name=\"section-6\"></a>\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "8aeae9da9796"
-      },
-      "outputs": [],
-      "source": [
-        "DATA_LOCATION = \"gs://cloud-samples-data/ai-platform-unified/datasets/tabular/cdm_pricing_large_table.csv\"\n",
-        "df = pd.read_csv(DATA_LOCATION)\n",
-        "print(df.shape)\n",
-        "df.head()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "7b98d5f09842"
-      },
-      "source": [
-        "We will build a forecast model on this data and thus determine the best price for a product. For this type of model, we may not be using many fields but just the sales and price related ones. For the current execrcise, we will just focus on the following fields :\n",
-        "- `Product_ID`\n",
-        "- `Customer_Hierarchy`\n",
-        "- `Fiscal_Date`\n",
-        "- `List_Price_Converged`\n",
-        "- `Invoiced_quantity_in_Pieces`\n",
-        "- `Net_Sales`\n",
-        "\n",
-        "\n",
-        "\n",
-        "## Data Analysis\n",
-        "<a name=\"section-7\"></a>\n",
-        "\n",
-        "First, we will explore the data and distributions.\n",
-        "\n",
-        "Select the required columns from the dataframe."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "af4b41c5eb1f"
-      },
-      "outputs": [],
-      "source": [
-        "id_col = \"Product_ID\"\n",
-        "date_col = \"Fiscal_Date\"\n",
-        "categ_cols = [\"Customer_Hierarchy\"]\n",
-        "num_cols = [\"List_Price_Converged\", \"Invoiced_quantity_in_Pieces\", \"Net_Sales\"]\n",
-        "\n",
-        "df = df[[id_col, date_col] + categ_cols + num_cols].copy()\n",
-        "df.head()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "3d780043ee5b"
-      },
-      "source": [
-        "Check the column types and null values in the dataframe."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f54c445a1288"
-      },
-      "outputs": [],
-      "source": [
-        "df.info()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cd817b414c4d"
-      },
-      "source": [
-        "This data description reveals that there are no null values in the data. Also, the field `Fiscal_Date` which is a date field is loaded as an object type. \n",
-        "\n",
-        "Change the type of the date field to datetime."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b160fac085c8"
-      },
-      "outputs": [],
-      "source": [
-        "df[\"Fiscal_Date\"] = pd.to_datetime(df[\"Fiscal_Date\"], infer_datetime_format=True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fb4778578064"
-      },
-      "source": [
-        "Plot the distributions for the categorical fields."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "dd0467cd57c3"
-      },
-      "outputs": [],
-      "source": [
-        "for i in categ_cols:\n",
-        "    df[i].value_counts(normalize=True).plot(kind=\"bar\")\n",
-        "    plt.title(i)\n",
-        "    plt.show()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "145deed255e0"
-      },
-      "source": [
-        "Plot the distributions for the numerical fields."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f934137c6d82"
-      },
-      "outputs": [],
-      "source": [
-        "for i in num_cols:\n",
-        "    _, ax = plt.subplots(1, 2, figsize=(10, 4))\n",
-        "    df[i].plot(kind=\"box\", ax=ax[0])\n",
-        "    df[i].plot(kind=\"hist\", ax=ax[1])\n",
-        "    ax[0].set_title(i + \"-Boxplot\")\n",
-        "    ax[1].set_title(i + \"-Histogram\")\n",
-        "    plt.show()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f9b9c2e58380"
-      },
-      "source": [
-        "Check maximum date and minimum date in Fiscal_Date column."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2a10aa689f9d"
-      },
-      "outputs": [],
-      "source": [
-        "print(df[\"Fiscal_Date\"].max())\n",
-        "print(df[\"Fiscal_Date\"].min())"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "4834f63e2e59"
-      },
-      "source": [
-        "Check the product distribution across each category."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "4664877f5304"
-      },
-      "outputs": [],
-      "source": [
-        "grp_cols = [\"Customer_Hierarchy\", \"Product_ID\"]\n",
-        "grp_df = df[grp_cols].groupby(by=grp_cols).count().reset_index()\n",
-        "grp_df.groupby(\"Customer_Hierarchy\").nunique()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "01ed02b9c8fd"
-      },
-      "source": [
-        "Check the percentage changes in the orders based on the percentage changes in the price."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "0b2c428cb135"
-      },
-      "outputs": [],
-      "source": [
-        "# aggregate the data\n",
-        "df_aggr = (\n",
-        "    df.groupby([\"Product_ID\", \"List_Price_Converged\"])\n",
-        "    .agg({\"Fiscal_Date\": min, \"Invoiced_quantity_in_Pieces\": sum, \"Net_Sales\": sum})\n",
-        "    .reset_index()\n",
-        ")\n",
-        "# rename the aggregated columns\n",
-        "df_aggr.rename(\n",
-        "    columns={\n",
-        "        \"Fiscal_Date\": \"First_price_date\",\n",
-        "        \"Invoiced_quantity_in_Pieces\": \"Total_ordered_pieces\",\n",
-        "        \"Net_Sales\": \"Total_net_sales\",\n",
-        "    },\n",
-        "    inplace=True,\n",
-        ")\n",
-        "\n",
-        "# sort values chronologically\n",
-        "df_aggr.sort_values(by=[\"Product_ID\", \"First_price_date\"], inplace=True)\n",
-        "df_aggr.reset_index(drop=True, inplace=True)\n",
-        "\n",
-        "# add columns for previous values\n",
-        "df_aggr[\"Previous_List\"] = df_aggr.groupby([\"Product_ID\"])[\n",
-        "    \"List_Price_Converged\"\n",
-        "].shift()\n",
-        "df_aggr[\"Previous_Total_ordered_pieces\"] = df_aggr.groupby([\"Product_ID\"])[\n",
-        "    \"Total_ordered_pieces\"\n",
-        "].shift()\n",
-        "\n",
-        "# average price change across sku's\n",
-        "df_aggr[\"price_change_perc\"] = (\n",
-        "    (df_aggr[\"List_Price_Converged\"] - df_aggr[\"Previous_List\"])\n",
-        "    / df_aggr[\"Previous_List\"].fillna(0)\n",
-        "    * 100\n",
-        ")\n",
-        "df_aggr[\"order_change_perc\"] = (\n",
-        "    (df_aggr[\"Total_ordered_pieces\"] - df_aggr[\"Previous_Total_ordered_pieces\"])\n",
-        "    / df_aggr[\"Previous_Total_ordered_pieces\"].fillna(0)\n",
-        "    * 100\n",
-        ")\n",
-        "\n",
-        "# plot a scatterplot to visualize the changes\n",
-        "sns.scatterplot(\n",
-        "    x=\"price_change_perc\",\n",
-        "    y=\"order_change_perc\",\n",
-        "    data=df_aggr,\n",
-        "    hue=\"Product_ID\",\n",
-        "    legend=False,\n",
-        ")\n",
-        "plt.title(\"Percentage of change in price vs order\")\n",
-        "plt.show()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8259e916fe25"
-      },
-      "source": [
-        "For most of the products, we see that the percentage change in orders are high where the percentage changes in the prices are low. This suggests that too much change in the prices can affect the number of orders. \n",
-        "\n",
-        "**Note**: There seem to be some outliers in the data as percentage changes greater than 800 are found and as evident from the box plots made earlier. In the current exercise, we will not take any manual measures to deal with outliers as we will create a BigQuery ML timeseries model that already deals with outliers.\n",
-        "\n",
-        "## Preprocess the data for training\n",
-        "<a name=\"section-8\"></a>\n",
-        "\n",
-        "Check which `Product_ID`s that have the maximum orders."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f5cbc7709c6a"
-      },
-      "outputs": [],
-      "source": [
-        "df_orders = df.groupby([\"Product_ID\", \"Customer_Hierarchy\"], as_index=False)[\n",
-        "    \"Invoiced_quantity_in_Pieces\"\n",
-        "].sum()\n",
-        "df_orders.loc[\n",
-        "    df_orders.groupby(\"Customer_Hierarchy\")[\"Invoiced_quantity_in_Pieces\"].idxmax()\n",
-        "]"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fd6d227e513e"
-      },
-      "source": [
-        "From the above result, we can infer the following :\n",
-        "\n",
-        "- Under **Food** category, **SKU 62** has maximum orders.\n",
-        "- Under **Manufacturing** category, **SKU 17** has maximum orders.\n",
-        "- Under **Paper** category, **SKU 107** has maximum orders.\n",
-        "- Under **Publishing** category, **SKU 8** has maximum orders.\n",
-        "- Under **Utilities** category, **SKU 140** has maximum orders.\n",
-        "\n",
-        "Given there are too many ids and only a few records for most of them, we will consider only the above `Product_ID`s for which there are maximum number of orders. \n",
-        "\n",
-        "**Note**: The `Invoiced_quantity_in_Pieces` field seem to be a *float* type rather than an *int* type as it should be. This could be probably because of the data itself might be averaged in the first place."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2dbc0d64d157"
-      },
-      "source": [
-        "Check the various prices available for these `Product_ID`s."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "acc1dbd2d838"
-      },
-      "outputs": [],
-      "source": [
-        "df_type_food = df[(df[\"Product_ID\"] == \"SKU 62\") & (df[\"Customer_Hierarchy\"] == \"Food\")]\n",
-        "print(\"Food :\")\n",
-        "print(df_type_food[\"List_Price_Converged\"].value_counts())\n",
-        "df_type_manuf = df[\n",
-        "    (df[\"Product_ID\"] == \"SKU 17\") & (df[\"Customer_Hierarchy\"] == \"Manufacturing\")\n",
-        "]\n",
-        "print(\"Manufacturing :\")\n",
-        "print(df_type_manuf[\"List_Price_Converged\"].value_counts())\n",
-        "df_type_paper = df[\n",
-        "    (df[\"Product_ID\"] == \"SKU 107\") & (df[\"Customer_Hierarchy\"] == \"Paper\")\n",
-        "]\n",
-        "print(\"Paper :\")\n",
-        "print(df_type_paper[\"List_Price_Converged\"].value_counts())\n",
-        "df_type_pub = df[\n",
-        "    (df[\"Product_ID\"] == \"SKU 8\") & (df[\"Customer_Hierarchy\"] == \"Publishing\")\n",
-        "]\n",
-        "print(\"Publishing :\")\n",
-        "print(df_type_pub[\"List_Price_Converged\"].value_counts())\n",
-        "df_type_util = df[\n",
-        "    (df[\"Product_ID\"] == \"SKU 140\") & (df[\"Customer_Hierarchy\"] == \"Utilities\")\n",
-        "]\n",
-        "print(\"Utilities :\")\n",
-        "print(df_type_util[\"List_Price_Converged\"].value_counts())"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "f023af578c0f"
-      },
-      "source": [
-        "In the publishing category, `Product_ID` `SKU 8` and `SKU 17` has less than or equal to two different prices in the entire data and so we will exclude them and consider the rest for building the forecast model. The idea here is to train a forecast model on the timeseries data for products with different prices.\n",
-        "\n",
-        "Join the data for all the `Product_ID`s into one dataframe and remove duplicate records."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "a44771cc4c20"
-      },
-      "outputs": [],
-      "source": [
-        "df_final = pd.concat([df_type_food, df_type_paper, df_type_util])\n",
-        "df_final = (\n",
-        "    df_final[\n",
-        "        [\n",
-        "            \"Product_ID\",\n",
-        "            \"Fiscal_Date\",\n",
-        "            \"Customer_Hierarchy\",\n",
-        "            \"List_Price_Converged\",\n",
-        "            \"Invoiced_quantity_in_Pieces\",\n",
-        "        ]\n",
-        "    ]\n",
-        "    .drop_duplicates()\n",
-        "    .reset_index(drop=True)\n",
-        ")\n",
-        "df_final.head()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "add5063df368"
-      },
-      "source": [
-        "Save the data to a BigQuery table."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fd82ba56571f"
-      },
-      "outputs": [],
-      "source": [
-        "bq_client = bigquery.Client(project=PROJECT_ID)\n",
-        "\n",
-        "job_config = bigquery.LoadJobConfig(\n",
-        "    # Specify a (partial) schema. All columns are always written to the\n",
-        "    # table. The schema is used to assist in data type definitions.\n",
-        "    schema=[\n",
-        "        bigquery.SchemaField(\"Product_ID\", bigquery.enums.SqlTypeNames.STRING),\n",
-        "        bigquery.SchemaField(\"Fiscal_Date\", bigquery.enums.SqlTypeNames.DATE),\n",
-        "        bigquery.SchemaField(\"List_Price_Converged\", bigquery.enums.SqlTypeNames.FLOAT),\n",
-        "        bigquery.SchemaField(\n",
-        "            \"Invoiced_quantity_in_Pieces\", bigquery.enums.SqlTypeNames.FLOAT\n",
-        "        ),\n",
-        "    ],\n",
-        "    # Optionally, set the write disposition. BigQuery appends loaded rows\n",
-        "    # to an existing table by default, but with WRITE_TRUNCATE write\n",
-        "    # disposition it replaces the table with the loaded data.\n",
-        "    write_disposition=\"WRITE_TRUNCATE\",\n",
-        ")\n",
-        "\n",
-        "# save the dataframe to a table in the created dataset\n",
-        "job = bq_client.load_table_from_dataframe(\n",
-        "    df_final,\n",
-        "    \"{}.{}.{}\".format(PROJECT_ID, DATASET, TRAINING_DATA_TABLE),\n",
-        "    job_config=job_config,\n",
-        ")  # Make an API request.\n",
-        "job.result()  # Wait for the job to complete."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fca77641b03b"
-      },
-      "source": [
-        "# Train the model using BigQuery ML\n",
-        "<a name=\"section-9\"></a>\n",
-        "\n",
-        "Train an [Arima-Plus](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-create-time-series) model on the data using BigQuery ML."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cded27507891"
-      },
-      "source": [
-        "#@bigquery\n",
-        "create or replace model pricing_optimization.bqml_arima\n",
-        "options\n",
-        " (model_type = 'ARIMA_PLUS',\n",
-        "  time_series_timestamp_col = 'Fiscal_Date',\n",
-        "  time_series_data_col = 'Invoiced_quantity_in_Pieces',\n",
-        "  time_series_id_col = 'ID'\n",
-        " ) as\n",
-        "select\n",
-        " Fiscal_Date,\n",
-        " Concat(Product_ID,\"_\" ,Cast(List_Price_Converged as string)) as ID,\n",
-        " Invoiced_quantity_in_Pieces\n",
-        "from\n",
-        " pricing_optimization.TRAINING_DATA\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "332fd11ff32b"
-      },
-      "source": [
-        "## Generate forecasts from the model\n",
-        "<a name=\"section-10\"></a>\n",
-        "\n",
-        "Predict the sales for the next 30 days for each id and save to a dataframe."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ef926cdbf28e"
-      },
-      "outputs": [],
-      "source": [
-        "client = Client()\n",
-        "\n",
-        "query = '''\n",
-        "DECLARE HORIZON STRING DEFAULT \"30\"; #number of values to forecast\n",
-        "DECLARE CONFIDENCE_LEVEL STRING DEFAULT \"0.90\"; ## required confidence level\n",
-        "\n",
-        "EXECUTE IMMEDIATE format(\"\"\"\n",
-        "    SELECT\n",
-        "      *\n",
-        "    FROM \n",
-        "      ML.FORECAST(MODEL pricing_optimization.bqml_arima, \n",
-        "                  STRUCT(%s AS horizon, \n",
-        "                         %s AS confidence_level)\n",
-        "                 )\n",
-        "    \"\"\",HORIZON,CONFIDENCE_LEVEL)'''\n",
-        "job = client.query(query)\n",
-        "dfforecast = job.to_dataframe()\n",
-        "dfforecast.head()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "608c7de72dae"
-      },
-      "source": [
-        "## Interpret the results to choose the best price\n",
-        "<a name=\"section-11\"></a>\n",
-        "\n",
-        "Calculate average forecast values for the forecast duration."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "e1e193680400"
-      },
-      "outputs": [],
-      "source": [
-        "dfforecast_avg = (\n",
-        "    dfforecast[[\"ID\", \"forecast_value\"]].groupby(\"ID\", as_index=False).mean()\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5ce395d652a3"
-      },
-      "source": [
-        "Extract the ID and Price fields from the ID field."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "452c56fa58ed"
-      },
-      "outputs": [],
-      "source": [
-        "dfforecast_avg[\"Product_ID\"] = dfforecast_avg[\"ID\"].apply(lambda x: x.split(\"_\")[0])\n",
-        "dfforecast_avg[\"Price\"] = dfforecast_avg[\"ID\"].apply(lambda x: x.split(\"_\")[1])"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "3cee67f4028f"
-      },
-      "source": [
-        "Plot the average forecasted sales vs. the price of the product."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "fb351c8f383d"
-      },
-      "outputs": [],
-      "source": [
-        "for i in dfforecast_avg[\"Product_ID\"].unique():\n",
-        "    dfforecast_avg[dfforecast_avg[\"Product_ID\"] == i].set_index(\"Price\").sort_values(\n",
-        "        \"forecast_value\"\n",
-        "    ).plot(kind=\"bar\")\n",
-        "    plt.title(\"Price vs. Average Sales for \" + i)\n",
-        "    plt.show()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "67ff3acc74a5"
-      },
-      "source": [
-        "Based on the plots for price vs. the average forecasted orders, it can be said that to avail the maximum orders, each of the considered `Product_ID`s can follow the below prices :\n",
-        "- SKU 107's price range can be from 4.44 - 4.73 units\n",
-        "- SKU 140's price can be 1.95 units\n",
-        "- SKU 62's price can be 4.23 units\n",
-        "\n",
-        "\n",
-        "## Clean Up\n",
-        "<a name=\"section-12\"></a>\n",
-        "\n",
-        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial. The following code deletes the entire dataset."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d78908b8134d"
-      },
-      "outputs": [],
-      "source": [
-        "# Construct a BigQuery client object.\n",
-        "client = bigquery.Client()\n",
-        "\n",
-        "# TODO(developer): Set model_id to the ID of the model to fetch.\n",
-        "dataset_id = \"{PROJECT}.{DATASET}\".format(PROJECT=PROJECT_ID, DATASET=DATASET)\n",
-        "\n",
-        "# Use the delete_contents parameter to delete a dataset and its contents.\n",
-        "# Use the not_found_ok parameter to not receive an error if the dataset has already been deleted.\n",
-        "client.delete_dataset(\n",
-        "    dataset_id, delete_contents=True, not_found_ok=True\n",
-        ")  # Make an API request.\n",
-        "\n",
-        "print(\"Deleted dataset '{}'.\".format(dataset_id))"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "name": "pricing-optimization.ipynb",
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}

notebooks/community/matching_engine/matching_engine_for_indexing.ipynb

@@ -459,7 +459,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil cp gs://cloud-samples-data/ai-platform-unified/matching_engine/glove-100-angular.hdf5 ."
+        "! gsutil cp gs://cloud-samples-data/vertex-ai/matching_engine/glove-100-angular.hdf5 ."
       ]
     },
     {

notebooks/community/ml_ops/README.md

@@ -12,7 +12,7 @@ The purpose of this set of notebooks and markdown files is to demonstrate Google
 2. [Experimentation](stage2)
 3. [Formalization](stage3)
 4. [Evaluation](stage4)
-5. Deployment
-6. Serving
-7. Monitoring
+5. [Deployment](stage5)
+6. [Serving](stage6)
+7. Monitoring(stage7)
 8. Continuous Training

notebooks/community/ml_ops/setup.md

@@ -0,0 +1,43 @@
+## Before you begin
+
+### Set up your Google Cloud project
+
+**The following steps are required, regardless of your notebook environment.**
+
+1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.
+
+1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).
+
+1. [Enable the Vertex AI, BigQuery, Compute Engine and Cloud Storage APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,bigquery,compute_component,storage_component).
+
+1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).
+
+1. Enter your project ID in the cell below. Then run the cell to make sure the
+Cloud SDK uses the right project for all the commands in this notebook.
+
+**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.
+
+### Set up your local development environment
+
+**If you are using Colab or Vertex AI Workbench Notebooks**, your environment already meets all the requirements to run this notebook. You can skip this step.
+
+**Otherwise**, make sure your environment meets this notebook's requirements. You need the following:
+
+- The Cloud Storage SDK
+- Python 3
+- virtualenv
+- Jupyter notebook running in a virtual environment with Python 3
+
+The Cloud Storage guide to [Setting up a Python development environment](https://cloud.google.com/python/setup) and the [Jupyter installation guide](https://jupyter.org/install) provide detailed instructions for meeting these requirements. The following steps provide a condensed set of instructions:
+
+1. [Install and initialize the SDK](https://cloud.google.com/sdk/docs/).
+
+2. [Install Python 3](https://cloud.google.com/python/setup#installing_python).
+
+3. [Install virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv) and create a virtual environment that uses Python 3.  Activate the virtual environment.
+
+4. To install Jupyter, run `pip3 install jupyter` on the command-line in a terminal shell.
+
+5. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.
+
+6. Open this notebook in the Jupyter Notebook Dashboard.

notebooks/community/ml_ops/setup.py

@@ -0,0 +1,112 @@
+import os
+import sys
+import argparse
+import subprocess
+import random
+import string
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--bucket', dest='bucket_required', action='store_true', 
+                    default=False, help='Bucket required')
+parser.add_argument('--email', dest='email_required', action='store_true', 
+                    default=False, help='Email required')
+parser.add_argument('--sa', dest='sa_required', action='store_true', 
+                    default=False, help='Service account required')
+parser.add_argument('--packages', dest='extra_packages',
+                    default='', type=str, help='additional required packages')
+args = parser.parse_args()
+
+extra_pkgs = args.extra_packages
+
+
+# Installation
+
+
+# The Vertex AI Workbench Notebook product has specific requirements
+IS_WORKBENCH_NOTEBOOK = os.getenv("DL_ANACONDA_HOME")
+IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(
+    "/opt/deeplearning/metadata/env_version"
+)
+IS_COLAB = "google.colab" in sys.modules
+
+# Vertex AI Notebook requires dependencies to be installed with '--user'
+USER_FLAG = ""
+if IS_WORKBENCH_NOTEBOOK:
+    USER_FLAG = "--user"
+
+# not used
+'''
+print("Installing packages")
+os.system(f"pip3 install --upgrade --quiet {USER_FLAG} google-cloud-aiplatform {args.extra_packages}")
+print("Done installation")
+'''
+
+# Authenticate
+if IS_COLAB:
+    from google.colab import auth as google_auth
+
+    google_auth.authenticate_user()
+
+
+# project ID
+if IS_WORKBENCH_NOTEBOOK:
+    shell_output = subprocess.check_output("gcloud config list --format 'value(core.project)' 2>/dev/null", shell=True)
+    PROJECT_ID = shell_output[0:-1].decode('utf-8')
+    print("PROJECT ID: ", PROJECT_ID)
+else:
+    PROJECT_ID = input("Enter PROJECT_ID: ")
+    os.system(f"gcloud config set project {PROJECT_ID}")
+
+# email
+if args.email_required:
+    shell_output = subprocess.check_output("gcloud config list --format 'value(core.account)' 2>/dev/null", shell=True)
+    EMAIL_ADDR = shell_output[0:-1].decode('utf-8')
+    if EMAIL_ADDR == '':
+            EMAIL_ADDR = input("Enter Email Address: ")
+    print("EMAIL_ADDR: ", EMAIL_ADDR)
+
+# region
+shell_output = subprocess.check_output("gcloud config list --format 'value(ai.region)'", shell=True)
+REGION = shell_output[0:-1].decode('utf-8')
+if REGION == '':
+    REGION = input("Enter REGION: ")
+print("REGION: ", REGION)
+
+# multi-region
+MULTI_REGION = REGION.split('-')[0]
+
+    
+# UUID
+# Generate a uuid of a specifed length(default=8)
+def generate_uuid(length: int = 8) -> str:
+    return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))
+
+
+UUID = generate_uuid()
+print("UUID", UUID)
+    
+# Bucket
+if args.bucket_required:
+    BUCKET_NAME = PROJECT_ID + "aip-" + UUID
+    BUCKET_URI = f"gs://{BUCKET_NAME}"
+    os.system(f"gsutil mb -l {REGION} {BUCKET_URI}")
+    print("BUCKET_URI", BUCKET_URI)
+    
+    
+# Project Number
+if args.sa_required:
+    if IS_WORKBENCH_NOTEBOOK:
+        shell_output = subprocess.check_output("gcloud auth list 2>/dev/null", shell=True)
+        SERVICE_ACCOUNT = shell_output[:-1].decode('utf-8').split('\n')[2].strip()
+        PROJECT_NUMBER = SERVICE_ACCOUNT.split('-')[0]
+    else:
+        shell_output = subprocess.check_output(f"gcloud projects describe {PROJECT_ID}", shell=True)
+        try:
+            PROJECT_NUMBER = shell_output[:-1].decode('utf-8').split('\n')[7].split(':')[-1].strip().replace("'", "")
+            SERVICE_ACCOUNT = f"{PROJECT_NUMBER}-compute@developer.gserviceaccount.com"
+        except:
+            PROJECT_NUMBER = input("Enter project number: ")
+            SERVICE_ACCOUNT = f"{PROJECT_NUMBER}-compute@developer.gserviceaccount.com"
+    
+    print("SERVICE_ACCOUNT", SERVICE_ACCOUNT)
+    print("PROJECT_NUMBER", PROJECT_NUMBER)

notebooks/community/ml_ops/stage1/README.md

@@ -22,31 +22,31 @@ The first stage in MLOps is the collection and preparation for the purpose of de
 - Data is preprocessed for training and evaluation using Dataflow.
 - Data augmentation is performed on-the-fly and is coupled with model feeding.
 
-<img src='stage1.jpg'>
+<img src='stage1v2.png'>
 
 ## Notebooks
 
 ### Get Started
 
-[Get Started with BQ datasets](get_started_bq_datasets.ipynb)
+[Get started with Dataflow](community/ml_ops/stage1/get_started_dataflow.ipynb)
+
+In this tutorial, you learn how to use `Dataflow` for training with `Vertex AI`.
 
-```
 The steps performed include:
 
-- Create a Vertex AI `Dataset` resource from `BigQuery` table -- compatible for `AutoML` training.
-- Extract a copy of the dataset from `BigQuery` to a CSV file in Cloud Storage -- compatible for `AutoML` or custom training.
-- Select rows from a `BigQuery` dataset into a `pandas` dataframe -- compatible for custom training.
-- Select rows from a `BigQuery` dataset into a `tf.data.Dataset` -- compatible for custom training `TensorFlow` models.
-- Select rows from extracted CSV files into a `tf.data.Dataset` -- compatible for custom training `TensorFlow` models.
-- Create a `BigQuery` dataset from CSV files.
-- Extract data from `BigQuery` table into a `DMatrix` -- compatible for custom training `XGBoost` models.
-```
+- Offline preprocessing of data:
+    - Serially - w/o dataflow
+    - Parallel - with dataflow
+- Upstream preprocessing of data:
+    - tabular data
+    - image data
 
-[Get Started with Vertex datasets](get_started_vertex_datasets.ipynb)
+[Get started with Vertex AI datasets](community/ml_ops/stage1/get_started_vertex_datasets.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Dataset` for training with `Vertex AI`.
 
-```
 The steps performed include:
-  
+
 - Create a Vertex AI `Dataset` resource for:
     - image data
     - text data
@@ -61,28 +61,54 @@ The steps performed include:
 - Detect anomalies in new data using TensorFlow Data Validation.
 - Generate a TFRecord feature specification using TensorFlow Transform from the data schema.
 - Export a dataset and convert to TFRecords.
-```
 
-[Get Started with Dataflow](get_started_dataflow.ipynb)
+[Get started with BigQuery datasets](community/ml_ops/stage1/get_started_bq_datasets.ipynb)
+
+In this tutorial, you learn how to use `BigQuery` as a dataset for training with `Vertex AI`.
 
-```
 The steps performed include:
 
-- Offline preprocessing of data:
-    - Serially - w/o dataflow
-    - Parallel - with dataflow
-- Upstream preprocessing of data:
-    - tabular data
-    - image data
-```
+- Create a Vertex AI `Dataset` resource from `BigQuery` table -- compatible for `AutoML` training.
+- Extract a copy of the dataset from `BigQuery` to a CSV file in Cloud Storage -- compatible for `AutoML` or custom training.
+- Select rows from a `BigQuery` dataset into a `pandas` dataframe -- compatible for custom training.
+- Select rows from a `BigQuery` dataset into a `tf.data.Dataset` -- compatible for custom training `TensorFlow` models.
+- Select rows from extracted CSV files into a `tf.data.Dataset` -- compatible for custom training `TensorFlow` models.
+- Create a `BigQuery` dataset from CSV files.
+- Extract data from `BigQuery` table into a `DMatrix` -- compatible for custom training `XGBoost` models.
+
+[Get started with Vertex AI Data Labeling](community/ml_ops/stage1/get_started_with_data_labeling.ipynb)
+
+In this tutorial, you learn how to use the `Vertex AI Data Labeling` service.
+
+The steps performed include:
+
+- Create a Specialist Pool for data labelers.
+- Create a data labeling job.
+- Submit the data labeling job.
+- List data labeling jobs.
+- Cancel a data labeling job.
+
+
+
+[Create an unlabelled Vertex AI AutoML text entity extraction dataset from PDFs using Vision API](community/ml_ops/stage1/get_started_with_visionapi_and_vertex_datasets.ipynb)
+
+In this tutorial, you learn to use `Vision API` to extract text from PDF files stored on a Cloud Storage bucket. You then process the results and create an unlabelled `Vertex AI Dataset`, compatible with `AutoML`, for text entity extraction.
+
+The steps performed include:
+
+1. Using `Vision API` to perform Optical Character Recognition (OCR) to extract text from PDF files.
+2. Processing the results and saving them to text files.
+3. Generating a `Vertex AI Dataset` import file.
+4. Creating a new unlabelled text entity extraction `Vertex AI Dataset` resource in `Vertex AI`.
+
 
 ### E2E Stage Example
 
 [Stage 1: Data Management](mlops_data_management.ipynb)
-  
+
+
 ```
 The steps performed include:
-
 - Explore and visualize the data.
 - Create a Vertex AI `Dataset` resource from `BigQuery` table -- for AutoML training.
 - Extract a copy of the dataset to a CSV file in Cloud Storage.
@@ -92,4 +118,3 @@ The steps performed include:
 - Generate a TFRecord feature specification using TensorFlow Data Validation from the data schema.
 - Preprocess a portion of the BigQuery data using `Dataflow` -- for custom training.
 ```
-

notebooks/community/ml_ops/stage1/get_started_bq_datasets.ipynb

@@ -8,7 +8,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2021 Google LLC\n",
+        "# Copyright 2022 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -34,13 +34,18 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_bq_datasets.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_bq_datasets.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
         "  </td>\n",
         "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/ai/platform/notebooks/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_bq_datasets.ipynb\">\n",
-        "      Open in Google Cloud Notebooks\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage1/get_started_bq_datasets.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>\n",
         "</table>\n",
@@ -59,17 +64,6 @@
         "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 1 : data management: get started with BigQuery datasets."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dataset:gsod,lrg"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "The dataset used for this tutorial is the GSOD dataset from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). The version of the dataset you use only the fields year, month and day to predict the value of mean daily temperature (mean_temp)."
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -104,7 +98,7 @@
       "source": [
         "### Recommendations\n",
         "\n",
-        "When doing E2E MLOps on Google Cloud, the following best practices with structured (tabular) data in BigQuery:\n",
+        "When doing E2E MLOps on Google Cloud, following are the best practices when dealing with structured (tabular) data in BigQuery:\n",
         "\n",
         "- For AutoML training:\n",
         "  - Create a managed dataset with Vertex AI `TabularDataset`.\n",
@@ -124,20 +118,47 @@
         "        - Within the generator (upstream)\n",
         "        - Within the model (downstream)\n",
         "    - XGBoost model training:\n",
-        "      - Use BigQuery ML builtin XGBoost training.\n",
+        "      - Use BigQuery ML built-in XGBoost training.\n",
         "      - Alternatively, create a DMatrix generator from CSV files extracted from BigQuery table.\n",
-        "    - Pytorch model training:\n",
+        "    - PyTorch model training:\n",
         "        - Extract the BigQuery to a pandas dataframe.\n",
         "        - Preprocess the data in the dataframe.\n",
         "        - Create a DataLoader generator from the pandas dataframe.\n",
         "\n",
         "\n",
-        "- Alternately:\n",
+        "- Alternatively:\n",
         "    - Extract the BigQuery table to CSV files.\n",
         "    - Preprocess the CSV files.\n",
         "    - Create a tf.data.Dataset generator from the CSV files."
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dataset:gsod,lrg"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "The dataset used for this tutorial is the GSOD dataset from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). In this version of the dataset you consider the fields year, month and day to predict the value of mean daily temperature (mean_temp)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9e483012a752"
+      },
+      "source": [
+        "### Costs\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "- Vertex AI\n",
+        "- Cloud Storage\n",
+        "- BigQuery\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing) and [BigQuery pricing](https://cloud.google.com/bigquery/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -146,7 +167,7 @@
       "source": [
         "## Installations\n",
         "\n",
-        "Install *one time* the packages for executing the MLOps notebooks."
+        "Install the following packages to execute this notebook."
       ]
     },
     {
@@ -157,40 +178,21 @@
       },
       "outputs": [],
       "source": [
-        "ONCE_ONLY = False\n",
-        "if ONCE_ONLY:\n",
-        "    ! pip3 install -U tensorflow==2.5 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-data-validation==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-transform==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-io==0.18 $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-bigquery $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-logging $USER_FLAG\n",
-        "    ! pip3 install --upgrade apache-beam[gcp] $USER_FLAG\n",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
-        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG\n",
-        "    ! pip3 install --upgrade kfp $USER_FLAG"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "install_xgboost"
-      },
-      "source": [
-        "Install the latest GA version of *XGBoost* library as well."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "install_xgboost"
-      },
-      "outputs": [],
-      "source": [
-        "! pip3 install -U xgboost $USER_FLAG"
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\") and not os.getenv(\"VIRTUAL_ENV\")\n",
+        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
+        "    \"/opt/deeplearning/metadata/env_version\"\n",
+        ")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
+        "extra_pkgs = \"tensorflow tensorflow-io==0.18 pyarrow xgboost google-cloud-bigquery\"\n",
+        "! pip3 install --upgrade --quiet {USER_FLAG} google-cloud-aiplatform $extra_pkgs"
       ]
     },
     {
@@ -212,9 +214,9 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "import sys\n",
         "\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
+        "if \"google.colab\" in sys.modules:\n",
         "    # Automatically restart kernel after installs\n",
         "    import IPython\n",
         "\n",
@@ -225,187 +227,42 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "project_id"
+        "id": "fc8fb52b5cca"
       },
       "source": [
-        "#### Set your project ID\n",
+        "### Common setup\n",
         "\n",
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+        "Now, execute the common setup for the notebook tutorials."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "set_project_id"
+        "id": "001a0fcd5d78"
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
+        "# Common code setup for notebook tutorials\n",
+        "\n",
+        "! wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/setup.py -O setup.py\n",
+        "\n",
+        "%run setup.py --bucket"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "autoset_project_id"
+        "id": "d809f07a8935"
       },
       "outputs": [],
       "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
-        "    # Get your GCP project id from gcloud\n",
-        "    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID:\", PROJECT_ID)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "set_gcloud_project_id"
-      },
-      "outputs": [],
-      "source": [
-        "! gcloud config set project $PROJECT_ID"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "region"
-      },
-      "source": [
-        "#### Region\n",
+        "# Other Common setup instructions for notebook tutorials\n",
         "\n",
-        "You can also change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
+        "! wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/setup.md -O setup.md\n",
         "\n",
-        "- Americas: `us-central1`\n",
-        "- Europe: `europe-west4`\n",
-        "- Asia Pacific: `asia-east1`\n",
-        "\n",
-        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
-        "\n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "timestamp"
-      },
-      "source": [
-        "#### Timestamp\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append the timestamp onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "timestamp"
-      },
-      "outputs": [],
-      "source": [
-        "from datetime import datetime\n",
-        "\n",
-        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bucket:custom"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
-        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "When you submit a custom training job using the Vertex SDK, you upload a Python package\n",
-        "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
-        "the code from this package. In this tutorial, Vertex AI also saves the\n",
-        "trained model that results from your job in the same bucket. You can then\n",
-        "create an `Endpoint` resource based on this output in order to serve\n",
-        "online predictions.\n",
-        "\n",
-        "Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "bucket"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "autoset_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "create_bucket"
-      },
-      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "create_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "validate_bucket"
-      },
-      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "validate_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al $BUCKET_NAME"
+        "%load setup.md"
       ]
     },
     {
@@ -414,9 +271,6 @@
         "id": "setup_vars"
       },
       "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
         "### Import libraries and define constants"
       ]
     },
@@ -428,75 +282,12 @@
       },
       "outputs": [],
       "source": [
-        "import google.cloud.aiplatform as aip"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_bq"
-      },
-      "source": [
-        "#### Import BigQuery\n",
-        "\n",
-        "Import the BigQuery package into your Python environment."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_bq"
-      },
-      "outputs": [],
-      "source": [
+        "import google.cloud.aiplatform as aiplatform\n",
+        "import pandas as pd\n",
+        "import xgboost as xgb\n",
         "from google.cloud import bigquery"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_xgboost"
-      },
-      "source": [
-        "#### Import XGBoost\n",
-        "\n",
-        "Import the XGBoost package into your Python environment."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_xgboost"
-      },
-      "outputs": [],
-      "source": [
-        "import xgboost as xgb"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_pandas"
-      },
-      "source": [
-        "#### Import pandas\n",
-        "\n",
-        "Import the pandas package into your Python environment."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_pandas"
-      },
-      "outputs": [],
-      "source": [
-        "import pandas as pd"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -516,7 +307,7 @@
       },
       "outputs": [],
       "source": [
-        "aip.init(project=PROJECT_ID, location=REGION)"
+        "aiplatform.init(project=PROJECT_ID, location=REGION)"
       ]
     },
     {
@@ -538,7 +329,7 @@
       },
       "outputs": [],
       "source": [
-        "bqclient = bigquery.Client()"
+        "bqclient = bigquery.Client(project=PROJECT_ID)"
       ]
     },
     {
@@ -549,7 +340,7 @@
       "source": [
         "#### Location of BigQuery training data.\n",
         "\n",
-        "Now set the variable `IMPORT_FILE` to the location of the data table in BigQuery."
+        "Now, set the variable `IMPORT_FILE` to the location of the data table in BigQuery and `BQ_TABLE` with the table id."
       ]
     },
     {
@@ -591,10 +382,10 @@
       },
       "outputs": [],
       "source": [
-        "dataset = aip.TabularDataset.create(\n",
-        "    display_name=\"NOAA historical weather data\" + \"_\" + TIMESTAMP,\n",
+        "dataset = aiplatform.TabularDataset.create(\n",
+        "    display_name=\"NOAA historical weather data\" + \"_\" + UUID,\n",
         "    bq_source=[IMPORT_FILE],\n",
-        "    labels={\"user_metadata\": BUCKET_NAME[5:]},\n",
+        "    labels={\"user_metadata\": BUCKET_URI[5:]},\n",
         ")\n",
         "\n",
         "label_column = \"mean_temp\"\n",
@@ -610,7 +401,7 @@
       "source": [
         "### Copy the dataset to Cloud Storage\n",
         "\n",
-        "Next, you make a copy of the BigQuery dataset, as a CSV file, to Cloud Storage using the BigQuery extract command.\n",
+        "Next, you make a copy of the BigQuery table as a CSV file, to Cloud Storage using the BigQuery extract command.\n",
         "\n",
         "Learn more about [BigQuery command line interface](https://cloud.google.com/bigquery/docs/reference/bq-cli-reference)."
       ]
@@ -626,9 +417,9 @@
         "comps = BQ_TABLE.split(\".\")\n",
         "BQ_PROJECT_DATASET_TABLE = comps[0] + \":\" + comps[1] + \".\" + comps[2]\n",
         "\n",
-        "! bq --location=us extract --destination_format CSV $BQ_PROJECT_DATASET_TABLE $BUCKET_NAME/mydata*.csv\n",
+        "! bq --location=us extract --destination_format CSV $BQ_PROJECT_DATASET_TABLE $BUCKET_URI/mydata*.csv\n",
         "\n",
-        "IMPORT_FILES = ! gsutil ls $BUCKET_NAME/mydata*.csv\n",
+        "IMPORT_FILES = ! gsutil ls $BUCKET_URI/mydata*.csv\n",
         "\n",
         "print(IMPORT_FILES)\n",
         "\n",
@@ -664,15 +455,12 @@
       },
       "outputs": [],
       "source": [
-        "if \"IMPORT_FILES\" in globals():\n",
-        "    gcs_source = IMPORT_FILES\n",
-        "else:\n",
-        "    gcs_source = [IMPORT_FILE]\n",
+        "gcs_source = IMPORT_FILES\n",
         "\n",
-        "dataset = aip.TabularDataset.create(\n",
-        "    display_name=\"NOAA historical weather data\" + \"_\" + TIMESTAMP,\n",
+        "dataset = aiplatform.TabularDataset.create(\n",
+        "    display_name=\"NOAA historical weather data\" + \"_\" + UUID,\n",
         "    gcs_source=gcs_source,\n",
-        "    labels={\"user_metadata\": BUCKET_NAME[5:]},\n",
+        "    labels={\"user_metadata\": BUCKET_URI[5:]},\n",
         ")\n",
         "\n",
         "\n",
@@ -694,6 +482,30 @@
         "Learn more about [Creating BigQuery views](https://cloud.google.com/bigquery/docs/views)."
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7dc142433e50"
+      },
+      "outputs": [],
+      "source": [
+        "# Set dataset name and view name in BigQuery\n",
+        "BQ_MY_DATASET = \"[your-dataset-name]\"\n",
+        "BQ_MY_TABLE = \"[your-view-name]\"\n",
+        "\n",
+        "# Otherwise, use the default names\n",
+        "if (\n",
+        "    BQ_MY_DATASET == \"\"\n",
+        "    or BQ_MY_DATASET is None\n",
+        "    or BQ_MY_DATASET == \"[your-dataset-name]\"\n",
+        "):\n",
+        "    BQ_MY_DATASET = \"mlops_dataset_\" + UUID\n",
+        "\n",
+        "if BQ_MY_TABLE == \"\" or BQ_MY_TABLE is None or BQ_MY_TABLE == \"[your-view-name]\":\n",
+        "    BQ_MY_TABLE = \"mlops_view_\" + UUID"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -702,8 +514,7 @@
       },
       "outputs": [],
       "source": [
-        "BQ_MY_DATASET = 'mydataset'\n",
-        "BQ_MY_TABLE = 'myview'\n",
+        "# Create the resources\n",
         "! bq --location=US mk -d \\\n",
         "$PROJECT_ID:$BQ_MY_DATASET\n",
         "\n",
@@ -744,8 +555,8 @@
       },
       "outputs": [],
       "source": [
-        "# Download a table.\n",
-        "table = bigquery.TableReference.from_string(\"bigquery-public-data.samples.gsod\")\n",
+        "# Download the table.\n",
+        "table = bigquery.TableReference.from_string(BQ_TABLE)\n",
         "\n",
         "rows = bqclient.list_rows(\n",
         "    table,\n",
@@ -1031,22 +842,6 @@
         "TABLE_ID = \"gsod\"\n",
         "\n",
         "\n",
-        "def create_bigquery_dataset(dataset_id):\n",
-        "    dataset = bigquery.Dataset(\n",
-        "        bigquery.dataset.DatasetReference(PROJECT_ID, dataset_id)\n",
-        "    )\n",
-        "    dataset.location = \"us\"\n",
-        "\n",
-        "    try:\n",
-        "        dataset = bqclient.create_dataset(dataset)  # API request\n",
-        "        return True\n",
-        "    except Exception as err:\n",
-        "        print(err)\n",
-        "        if err.code != 409:  # http_client.CONFLICT\n",
-        "            raise\n",
-        "    return False\n",
-        "\n",
-        "\n",
         "def load_data_into_bigquery(url, dataset_id, table_id):\n",
         "    create_bigquery_dataset(dataset_id)\n",
         "    dataset = bqclient.dataset(dataset_id)\n",
@@ -1079,13 +874,11 @@
       "source": [
         "### Read BigQuery table into XGboost DMatrix\n",
         "\n",
-        "Currently, there is no direct data feeding connector between BigQuery and the open source XGBoost.\n",
+        "Currently, there is no direct data feeding connector between BigQuery and the open source XGBoost. The BigQuery ML service has a built-in XGBoost training module.\n",
         "\n",
-        "The BigQuery ML service has XGBoost training builtin.\n",
+        "Alernatively, you extract the data either as a pandas dataframe or as CSV files. The extracted data is then given as an input to a `DMatrix` object when training the model.\n",
         "\n",
-        "Alernatively, you extract the data either as a pandas dataframe or as CSV files. The extracted data is then inputted to a `DMatrix` object when training the model.\n",
-        "\n",
-        "Learn more about [Getting started with builtin XGBoost](https://cloud.google.com/ai-platform/training/docs/algorithms/xgboost-start)"
+        "Learn more about [Getting started with built-in XGBoost](https://cloud.google.com/ai-platform/training/docs/algorithms/xgboost-start)."
       ]
     },
     {
@@ -1096,7 +889,7 @@
       "source": [
         "### Read pandas table into XGboost DMatrix\n",
         "\n",
-        "Next, you load the pandas dataframe into a `DMatrix` object. XGBoost does not support non-numeric inputs. Any column that is categorical will need to be one-hot encoded prior to loading the dataframe."
+        "Next, you load the pandas dataframe into a `DMatrix` object. XGBoost does not support non-numeric inputs. Any column that is categorical need to be one-hot encoded prior to loading the dataframe."
       ]
     },
     {
@@ -1109,7 +902,7 @@
       "source": [
         "dataframe[\"station_number\"] = pd.to_numeric(dataframe[\"station_number\"])\n",
         "labels = dataframe[\"mean_temp\"]\n",
-        "data = dataframe.drop(4)\n",
+        "data = dataframe.drop([\"mean_temp\"], axis=1)\n",
         "\n",
         "dtrain = xgb.DMatrix(data, label=labels)"
       ]
@@ -1122,7 +915,7 @@
       "source": [
         "### Read CSV files into XGboost DMatrix\n",
         "\n",
-        "Currently, there is no Cloud Storage support in XGBoost. If you use CSV files for input, you will need to download them locally."
+        "Currently, there is no Cloud Storage support in XGBoost. If you use CSV files for input, you need to download them locally."
       ]
     },
     {
@@ -1144,87 +937,42 @@
         "id": "cleanup:mbsdk"
       },
       "source": [
-        "# Cleaning up\n",
+        "# Clean up\n",
         "\n",
         "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
         "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
         "\n",
         "Otherwise, you can delete the individual resources you created in this tutorial:\n",
         "\n",
-        "- Dataset\n",
-        "- Pipeline\n",
-        "- Model\n",
-        "- Endpoint\n",
-        "- AutoML Training Job\n",
-        "- Batch Job\n",
-        "- Custom Job\n",
-        "- Hyperparameter Tuning Job\n",
-        "- Cloud Storage Bucket"
+        "- Vertex AI Dataset resource\n",
+        "- Cloud Storage Bucket\n",
+        "- BigQuery Dataset\n",
+        "\n",
+        "Set `delete_storage` to _True_ to delete the storage resources used in this notebook."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "cleanup:mbsdk"
+        "id": "47ad926d84e8"
       },
       "outputs": [],
       "source": [
-        "delete_all = True\n",
+        "import os\n",
         "\n",
-        "if delete_all:\n",
-        "    # Delete the dataset using the Vertex dataset object\n",
-        "    try:\n",
-        "        if \"dataset\" in globals():\n",
-        "            dataset.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the dataset using the Vertex dataset object\n",
+        "dataset.delete()\n",
         "\n",
-        "    # Delete the model using the Vertex model object\n",
-        "    try:\n",
-        "        if \"model\" in globals():\n",
-        "            model.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the temporary BigQuery dataset\n",
+        "! bq rm -r -f $PROJECT_ID:$DATASET_ID\n",
         "\n",
-        "    # Delete the endpoint using the Vertex endpoint object\n",
-        "    try:\n",
-        "        if \"endpoint\" in globals():\n",
-        "            endpoint.undeploy_all()\n",
-        "            endpoint.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the AutoML or Pipeline training job\n",
-        "    try:\n",
-        "        if \"dag\" in globals():\n",
-        "            dag.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the custom training job\n",
-        "    try:\n",
-        "        if \"job\" in globals():\n",
-        "            job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the batch prediction job using the Vertex batch prediction object\n",
-        "    try:\n",
-        "        if \"batch_predict_job\" in globals():\n",
-        "            batch_predict_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the hyperparameter tuning job using the Vertex hyperparameter tuning object\n",
-        "    try:\n",
-        "        if \"hpt_job\" in globals():\n",
-        "            hpt_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    if \"BUCKET_NAME\" in globals():\n",
-        "        ! gsutil rm -r $BUCKET_NAME"
+        "delete_storage = False\n",
+        "if delete_storage or os.getenv(\"IS_TESTING\"):\n",
+        "    # Delete the created GCS bucket\n",
+        "    ! gsutil rm -r $BUCKET_URI\n",
+        "    # Delete the created BigQuery datasets\n",
+        "    ! bq rm -r -f $PROJECT_ID:$BQ_MY_DATASET"
       ]
     }
   ],

notebooks/community/ml_ops/stage1/get_started_dataflow.ipynb

@@ -39,8 +39,14 @@
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/ai/platform/notebooks/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_dataflow.ipynb\">\n",
-        "      Open in Google Cloud Notebooks\n",
+        "        <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_dataflow.ipynb\">\n",
+        "        <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\\\" alt=\"Colab logo\"> Run in Colab\n",
+        "        </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage1/get_started_dataflow.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>\n",
         "</table>\n",
@@ -59,17 +65,6 @@
         "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 1 : data management: get started with Dataflow."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dataset:gsod,lrg"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "The dataset used for this tutorial is the GSOD dataset from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). The version of the dataset you use only the fields year, month and day to predict the value of mean daily temperature (mean_temp)."
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -131,6 +126,34 @@
         "Alternately for AutoML tabular model training, you can reconfigure the otherwise default preprocessing."
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dataset:gsod,lrg"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "The dataset used for this tutorial is the GSOD dataset from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). The version of the dataset you use only the fields year, month and day to predict the value of mean daily temperature (mean_temp)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9e483012a752"
+      },
+      "source": [
+        "### Costs\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "- Vertex AI\n",
+        "- Cloud Storage\n",
+        "- BigQuery\n",
+        "- Dataflow\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [BigQuery pricing](https://cloud.google.com/bigquery/pricing), and [Dataflow pricing](https://cloud.google.com/dataflow/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -139,7 +162,7 @@
       "source": [
         "## Installations\n",
         "\n",
-        "Install *one time* the packages for executing the MLOps notebooks."
+        "Install the following packages to execute this notebook."
       ]
     },
     {
@@ -150,20 +173,22 @@
       },
       "outputs": [],
       "source": [
-        "ONCE_ONLY = False\n",
-        "if ONCE_ONLY:\n",
-        "    ! pip3 install -U tensorflow==2.5 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-data-validation==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-transform==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-io==0.18 $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-bigquery $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-logging $USER_FLAG\n",
-        "    ! pip3 install --upgrade apache-beam[gcp] $USER_FLAG\n",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
-        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG\n",
-        "    ! pip3 install --upgrade kfp $USER_FLAG"
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\") and not os.getenv(\"VIRTUAL_ENV\")\n",
+        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
+        "    \"/opt/deeplearning/metadata/env_version\"\n",
+        ")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
+        "extra_pkgs = \"tensorflow==2.5 tensorflow-data-validation==1.2 tensorflow-transform==1.2 \\\n",
+        "              tensorflow-io==0.18 pyarrow pandas apache-beam[gcp] google-cloud-bigquery\"\n",
+        "! pip3 install --upgrade --quiet {USER_FLAG} google-cloud-aiplatform $extra_pkgs"
       ]
     },
     {
@@ -185,9 +210,9 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "import sys\n",
         "\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
+        "if \"google.colab\" in sys.modules:\n",
         "    # Automatically restart kernel after installs\n",
         "    import IPython\n",
         "\n",
@@ -198,187 +223,42 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "project_id"
+        "id": "fc8fb52b5cca"
       },
       "source": [
-        "#### Set your project ID\n",
+        "### Common setup\n",
         "\n",
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+        "Now, execute the common setup for the notebook tutorials."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "set_project_id"
+        "id": "001a0fcd5d78"
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
+        "# Common code setup for notebook tutorials\n",
+        "\n",
+        "! wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/setup.py -O setup.py\n",
+        "\n",
+        "%run setup.py --bucket"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "autoset_project_id"
+        "id": "d809f07a8935"
       },
       "outputs": [],
       "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
-        "    # Get your GCP project id from gcloud\n",
-        "    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID:\", PROJECT_ID)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "set_gcloud_project_id"
-      },
-      "outputs": [],
-      "source": [
-        "! gcloud config set project $PROJECT_ID"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "region"
-      },
-      "source": [
-        "#### Region\n",
+        "# Other Common setup instructions for notebook tutorials\n",
         "\n",
-        "You can also change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
+        "! wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/setup.md -O setup.md\n",
         "\n",
-        "- Americas: `us-central1`\n",
-        "- Europe: `europe-west4`\n",
-        "- Asia Pacific: `asia-east1`\n",
-        "\n",
-        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
-        "\n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "region"
-      },
-      "outputs": [],
-      "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "timestamp"
-      },
-      "source": [
-        "#### Timestamp\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append the timestamp onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "timestamp"
-      },
-      "outputs": [],
-      "source": [
-        "from datetime import datetime\n",
-        "\n",
-        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bucket:custom"
-      },
-      "source": [
-        "### Create a Cloud Storage bucket\n",
-        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "When you submit a custom training job using the Vertex SDK, you upload a Python package\n",
-        "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
-        "the code from this package. In this tutorial, Vertex AI also saves the\n",
-        "trained model that results from your job in the same bucket. You can then\n",
-        "create an `Endpoint` resource based on this output in order to serve\n",
-        "online predictions.\n",
-        "\n",
-        "Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "bucket"
-      },
-      "outputs": [],
-      "source": [
-        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "autoset_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "create_bucket"
-      },
-      "source": [
-        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "create_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "validate_bucket"
-      },
-      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "validate_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al $BUCKET_NAME"
+        "%load setup.md "
       ]
     },
     {
@@ -401,7 +281,7 @@
       },
       "outputs": [],
       "source": [
-        "import google.cloud.aiplatform as aip"
+        "import google.cloud.aiplatform as aiplatform"
       ]
     },
     {
@@ -555,7 +435,7 @@
       },
       "outputs": [],
       "source": [
-        "aip.init(project=PROJECT_ID, location=REGION)"
+        "aiplatform.init(project=PROJECT_ID, location=REGION)"
       ]
     },
     {
@@ -676,6 +556,31 @@
         "dataframe[\"station_number\"] = pd.to_numeric(dataframe[\"station_number\"])"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bqml_create_dataset"
+      },
+      "source": [
+        "### Create BQ dataset resource\n",
+        "\n",
+        "First, you create an empty dataset resource in your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bqml_create_dataset"
+      },
+      "outputs": [],
+      "source": [
+        "BQ_MY_DATASET = 'samples'\n",
+        "BQ_MY_TABLE = 'gsod'\n",
+        "! bq --location=US mk -d \\\n",
+        "$PROJECT_ID:$BQ_MY_DATASET"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -979,7 +884,7 @@
       },
       "outputs": [],
       "source": [
-        "SCHEMA_LOCATION = BUCKET_NAME + \"/schema.txt\"\n",
+        "SCHEMA_LOCATION = BUCKET_URI + \"/schema.txt\"\n",
         "\n",
         "# When running Apache Beam directly (file is directly accessed)\n",
         "tfdv.write_schema_text(output_path=SCHEMA_LOCATION, schema=schema)\n",
@@ -1124,7 +1029,7 @@
         "            )\n",
         "\n",
         "\n",
-        "EXPORTED_DATA_PREFIX = os.path.join(BUCKET_NAME, \"exported_data\")\n",
+        "EXPORTED_DATA_PREFIX = os.path.join(BUCKET_URI, \"exported_data\")\n",
         "\n",
         "QUERY_STRING = \"SELECT {},{} FROM {} LIMIT 500\".format(\n",
         "    \"CAST(station_number as STRING) AS station_number,year,month,day\",\n",
@@ -1137,7 +1042,7 @@
         "    \"runner\": RUNNER,\n",
         "    \"raw_data_query\": QUERY_STRING,\n",
         "    \"exported_data_prefix\": EXPORTED_DATA_PREFIX,\n",
-        "    \"temp_location\": os.path.join(BUCKET_NAME, \"temp\"),\n",
+        "    \"temp_location\": os.path.join(BUCKET_URI, \"temp\"),\n",
         "    \"project\": PROJECT_ID,\n",
         "    \"region\": REGION,\n",
         "    \"setup_file\": \"./setup.py\",\n",
@@ -1162,17 +1067,7 @@
         "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
         "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
         "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
-        "\n",
-        "- Dataset\n",
-        "- Pipeline\n",
-        "- Model\n",
-        "- Endpoint\n",
-        "- AutoML Training Job\n",
-        "- Batch Job\n",
-        "- Custom Job\n",
-        "- Hyperparameter Tuning Job\n",
-        "- Cloud Storage Bucket"
+        "Otherwise, you can delete the individual resources you created in this tutorial."
       ]
     },
     {
@@ -1183,61 +1078,11 @@
       },
       "outputs": [],
       "source": [
-        "delete_all = True\n",
+        "delete_storage = False\n",
         "\n",
-        "if delete_all:\n",
-        "    # Delete the dataset using the Vertex dataset object\n",
-        "    try:\n",
-        "        if \"dataset\" in globals():\n",
-        "            dataset.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the model using the Vertex model object\n",
-        "    try:\n",
-        "        if \"model\" in globals():\n",
-        "            model.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the endpoint using the Vertex endpoint object\n",
-        "    try:\n",
-        "        if \"endpoint\" in globals():\n",
-        "            endpoint.undeploy_all()\n",
-        "            endpoint.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the AutoML or Pipeline training job\n",
-        "    try:\n",
-        "        if \"dag\" in globals():\n",
-        "            dag.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the custom training job\n",
-        "    try:\n",
-        "        if \"job\" in globals():\n",
-        "            job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the batch prediction job using the Vertex batch prediction object\n",
-        "    try:\n",
-        "        if \"batch_predict_job\" in globals():\n",
-        "            batch_predict_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the hyperparameter tuning job using the Vertex hyperparameter tuning object\n",
-        "    try:\n",
-        "        if \"hpt_job\" in globals():\n",
-        "            hpt_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    if \"BUCKET_NAME\" in globals():\n",
-        "        ! gsutil rm -r $BUCKET_NAME"
+        "if delete_storage or os.getenv(\"IS_TESTING\"):\n",
+        "    if \"BUCKET_URI\" in globals():\n",
+        "        ! gsutil rm -r $BUCKET_URI"
       ]
     }
   ],

notebooks/community/ml_ops/stage1/get_started_vertex_datasets.ipynb

@@ -8,7 +8,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2021 Google LLC\n",
+        "# Copyright 2022 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -29,20 +29,26 @@
         "id": "title:generic,gcp"
       },
       "source": [
-        "# E2E ML on GCP: MLOps stage 1 : data management: get started with Vertex datasets\n",
+        "# E2E ML on GCP: MLOps stage 1 : data management: get started with Vertex AI datasets\n",
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_vertex_datasets.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
-        "      View on GitHub\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_vertex_datasets.ipynb\">\n",
+        "     <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/ai/platform/notebooks/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_vertex_datasets.ipynb\">\n",
-        "      Open in Google Cloud Notebooks\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/get_started_vertex_datasets.ipynb\">\n",
+        "<img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
+        "   <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage1/get_started_vertex_datasets.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>   \n",
         "</table>\n",
         "<br/><br/><br/>"
       ]
@@ -133,6 +139,33 @@
         "            - Create a tf.data.Dataset from the TFRecords."
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "533dd6fe83c8"
+      },
+      "source": [
+        "### Datasets\n",
+        "\n",
+        "This tutorial uses a variety of public datasets to demonstrate using a `Vertex AI` managed dataset."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9e483012a752"
+      },
+      "source": [
+        "### Costs\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "- Vertex AI\n",
+        "- Cloud Storage\n",
+        "- BigQuery\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing) and [BigQuery pricing](https://cloud.google.com/bigquery/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -141,7 +174,7 @@
       "source": [
         "## Installations\n",
         "\n",
-        "Install *one time* the packages for executing the MLOps notebooks."
+        "Install the packages required for executing this notebook."
       ]
     },
     {
@@ -152,20 +185,27 @@
       },
       "outputs": [],
       "source": [
-        "ONCE_ONLY = False\n",
-        "if ONCE_ONLY:\n",
-        "    ! pip3 install -U tensorflow==2.5 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-data-validation==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-transform==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-io==0.18 $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-bigquery $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-logging $USER_FLAG\n",
-        "    ! pip3 install --upgrade apache-beam[gcp] $USER_FLAG\n",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
-        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG\n",
-        "    ! pip3 install --upgrade kfp $USER_FLAG"
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\") and not os.getenv(\"VIRTUAL_ENV\")\n",
+        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
+        "    \"/opt/deeplearning/metadata/env_version\"\n",
+        ")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
+        "! pip3 install -U tensorflow $USER_FLAG -q\n",
+        "! pip3 install -U tensorflow-data-validation $USER_FLAG -q\n",
+        "! pip3 install -U tensorflow-transform $USER_FLAG -q\n",
+        "! pip3 install -U tensorflow-io $USER_FLAG -q\n",
+        "! pip3 install --upgrade google-cloud-aiplatform $USER_FLAG -q\n",
+        "! pip3 install --upgrade google-cloud-bigquery $USER_FLAG -q\n",
+        "! pip3 install -U tensorflow-io==0.18 $USER_FLAG -q\n",
+        "! pip3 install --upgrade db-dtypes $USER_FLAG -q! pip3 install --upgrade future $USER_FLAG -q"
       ]
     },
     {
@@ -197,6 +237,32 @@
         "    app.kernel.do_shutdown(True)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cb082379ed5b"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). \n",
+        "\n",
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -223,7 +289,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "autoset_project_id"
+        "id": "nWlzLu5ELxWd"
       },
       "outputs": [],
       "source": [
@@ -238,7 +304,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "set_gcloud_project_id"
+        "id": "c021ca495967"
       },
       "outputs": [],
       "source": [
@@ -273,7 +339,10 @@
       },
       "outputs": [],
       "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if REGION == \"[your-region]\":\n",
+        "    REGION = \"us-central1\""
       ]
     },
     {
@@ -300,6 +369,66 @@
         "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "927085b84a07"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "**If you are using Vertex AI Workbench Notebooks**, your environment is already authenticated. Skip this step.\n",
+        "\n",
+        "**If you are using Colab**, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.\n",
+        "\n",
+        "**Otherwise**, follow these steps:\n",
+        "\n",
+        "In the Cloud Console, go to the [Create service account key](https://console.cloud.google.com/apis/credentials/serviceaccountkey) page.\n",
+        "\n",
+        "**Click Create service account**.\n",
+        "\n",
+        "In the **Service account name** field, enter a name, and click **Create**.\n",
+        "\n",
+        "In the **Grant this service account access to project** section, click the Role drop-down list. Type \"Vertex\" into the filter box, and select **Vertex Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
+        "\n",
+        "Click Create. A JSON file that contains your key downloads to your local environment.\n",
+        "\n",
+        "Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "89788a802687"
+      },
+      "outputs": [],
+      "source": [
+        "# If you are running this notebook in Colab, run this cell and follow the\n",
+        "# instructions to authenticate your GCP account. This provides access to your\n",
+        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
+        "# requests.\n",
+        "\n",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "# If on Vertex AI Workbench, then don't execute this code\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
+        "    \"DL_ANACONDA_HOME\"\n",
+        "):\n",
+        "    if \"google.colab\" in sys.modules:\n",
+        "        from google.colab import auth as google_auth\n",
+        "\n",
+        "        google_auth.authenticate_user()\n",
+        "\n",
+        "    # If you are running this notebook locally, replace the string below with the\n",
+        "    # path to your service account key and run this cell to authenticate your GCP\n",
+        "    # account.\n",
+        "    elif not os.getenv(\"IS_TESTING\"):\n",
+        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -328,7 +457,7 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
+        "BUCKET_URI = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -339,8 +468,8 @@
       },
       "outputs": [],
       "source": [
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
+        "if BUCKET_URI == \"\" or BUCKET_URI is None or BUCKET_URI == \"gs://[your-bucket-name]\":\n",
+        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
       ]
     },
     {
@@ -360,7 +489,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
+        "! gsutil mb -l $REGION $BUCKET_URI"
       ]
     },
     {
@@ -380,7 +509,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls -al $BUCKET_NAME"
+        "! gsutil ls -al $BUCKET_URI"
       ]
     },
     {
@@ -392,7 +521,11 @@
         "### Set up variables\n",
         "\n",
         "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
+        "### Import libraries and define constants\n",
+        "\n",
+        "Import the BigQuery package, TensorFlow Data Validation (TFDV) package and TensorFlow Data Validation package into your Python environment. \n",
+        "\n",
+        "Import TensorFlow Transform (TFT) package and pandas into your Python environment."
       ]
     },
     {
@@ -403,97 +536,13 @@
       },
       "outputs": [],
       "source": [
-        "import google.cloud.aiplatform as aip"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_bq"
-      },
-      "source": [
-        "#### Import BigQuery\n",
-        "\n",
-        "Import the BigQuery package into your Python environment."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_bq"
-      },
-      "outputs": [],
-      "source": [
+        "import google.cloud.aiplatform as aip\n",
+        "import pandas as pd\n",
+        "import tensorflow_data_validation as tfdv\n",
+        "import tensorflow_transform as tft\n",
         "from google.cloud import bigquery"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_tfdv"
-      },
-      "source": [
-        "#### Import TensorFlow Data Validation\n",
-        "\n",
-        "Import the TensorFlow Data Validation (TFDV) package into your Python environment."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_tfdv"
-      },
-      "outputs": [],
-      "source": [
-        "import tensorflow_data_validation as tfdv"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_tft"
-      },
-      "source": [
-        "#### Import TensorFlow Transform\n",
-        "\n",
-        "Import the TensorFlow Transform (TFT) package into your Python environment."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_tft"
-      },
-      "outputs": [],
-      "source": [
-        "import tensorflow_transform as tft"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_pandas"
-      },
-      "source": [
-        "#### Import pandas\n",
-        "\n",
-        "Import the pandas package into your Python environment."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_pandas"
-      },
-      "outputs": [],
-      "source": [
-        "import pandas as pd"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -513,7 +562,7 @@
       },
       "outputs": [],
       "source": [
-        "aip.init(project=PROJECT_ID, location=REGION)"
+        "aip.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
       ]
     },
     {
@@ -567,26 +616,13 @@
         "Learn more about [All dataset documentation](https://cloud.google.com/vertex-ai/docs/datasets/datasets)"
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_file:flowers,csv,icn"
-      },
-      "outputs": [],
-      "source": [
-        "IMPORT_FILE = (\n",
-        "    \"gs://cloud-samples-data/vision/automl_classification/flowers/all_data_v2.csv\"\n",
-        ")"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "create_dataset:image,icn"
       },
       "source": [
-        "### Create the Dataset\n",
+        "### Create an Image Dataset\n",
         "\n",
         "Next, create the `Dataset` resource using the `create` method for the `ImageDataset` class, which takes the following parameters:\n",
         "\n",
@@ -601,6 +637,19 @@
         "Learn more about [ImageDataset](https://cloud.google.com/vertex-ai/docs/datasets/prepare-image)."
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "import_file:flowers,csv,icn"
+      },
+      "outputs": [],
+      "source": [
+        "IMPORT_FILE = (\n",
+        "    \"gs://cloud-samples-data/vision/automl_classification/flowers/all_data_v2.csv\"\n",
+        ")"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -618,24 +667,13 @@
         "print(dataset.resource_name)"
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_file:hmdb,csv,vcn"
-      },
-      "outputs": [],
-      "source": [
-        "IMPORT_FILE = \"gs://automl-video-demo-data/hmdb_split1_5classes_train_inf.csv\""
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "create_dataset:video,vcn"
       },
       "source": [
-        "### Create the Dataset\n",
+        "### Create a Video Dataset\n",
         "\n",
         "Next, create the `Dataset` resource using the `create` method for the `VideoDataset` class, which takes the following parameters:\n",
         "\n",
@@ -649,6 +687,17 @@
         "Learn more about [VideoDataset](https://cloud.google.com/vertex-ai/docs/datasets/prepare-video)."
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "import_file:hmdb,csv,vcn"
+      },
+      "outputs": [],
+      "source": [
+        "IMPORT_FILE = \"gs://automl-video-demo-data/hmdb_split1_5classes_train_inf.csv\""
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -666,24 +715,13 @@
         "print(dataset.resource_name)"
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_file:happydb,csv,tcn"
-      },
-      "outputs": [],
-      "source": [
-        "IMPORT_FILE = \"gs://cloud-ml-data/NL-classification/happiness.csv\""
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "create_dataset:text,tcn"
       },
       "source": [
-        "### Create the Dataset\n",
+        "### Create a Text Dataset\n",
         "\n",
         "Next, create the `Dataset` resource using the `create` method for the `TextDataset` class, which takes the following parameters:\n",
         "\n",
@@ -698,6 +736,17 @@
         "Learn more about [TextDataset](https://cloud.google.com/vertex-ai/docs/datasets/prepare-text)."
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "import_file:happydb,csv,tcn"
+      },
+      "outputs": [],
+      "source": [
+        "IMPORT_FILE = \"gs://cloud-ml-data/NL-classification/happiness.csv\""
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -715,6 +764,24 @@
         "print(dataset.resource_name)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "create_dataset:tabular,bq,lrg,v2"
+      },
+      "source": [
+        "### Create a Tabular Dataset\n",
+        "\n",
+        "#### CSV input data\n",
+        "\n",
+        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class for CSV input data, which takes the following parameters:\n",
+        "\n",
+        "- `display_name`: The human readable name for the `Dataset` resource.\n",
+        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
+        "\n",
+        "Learn more about [TabularDataset from CSV files](https://cloud.google.com/vertex-ai/docs/datasets/create-dataset-api#aiplatform_create_dataset_tabular_gcs_sample-python)"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -723,27 +790,50 @@
       },
       "outputs": [],
       "source": [
-        "IMPORT_FILE = \"bq://bigquery-public-data.samples.gsod\"\n",
-        "BQ_TABLE = \"bigquery-public-data.samples.gsod\""
+        "IMPORT_FILE = \"gs://cloud-samples-data/tables/iris_1000.csv\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "create_dataset:tabular,bq,lrg,v2"
+      },
+      "outputs": [],
+      "source": [
+        "dataset = aip.TabularDataset.create(\n",
+        "    display_name=\"example\" + \"_\" + TIMESTAMP, gcs_source=[IMPORT_FILE]\n",
+        ")\n",
+        "\n",
+        "print(dataset.resource_name)"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "create_dataset:tabular,bq,lrg,v2"
+        "id": "854dd1e0195c"
       },
       "source": [
-        "### Create the Dataset\n",
+        "#### BigQuery input data\n",
         "\n",
-        "#### CSV input data\n",
-        "\n",
-        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
+        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class for BigQuery table input, which takes the following parameters:\n",
         "\n",
         "- `display_name`: The human readable name for the `Dataset` resource.\n",
-        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
-        "- `labels`: User defined metadata. In this example, you store the location of the Cloud Storage bucket containing the user defined data.\n",
+        "- `bq_source`: A list of one or more BigQuery tables to import the data items into the `Dataset` resource.\n",
         "\n",
-        "Learn more about [TabularDataset from CSV files](https://cloud.google.com/vertex-ai/docs/datasets/create-dataset-api#aiplatform_create_dataset_tabular_gcs_sample-python)"
+        "Learn more about [TabularDataset from BigQuery table](https://cloud.google.com/vertex-ai/docs/datasets/create-dataset-api#aiplatform_create_dataset_tabular_bigquery_sample-pythonn)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "86343c146300"
+      },
+      "outputs": [],
+      "source": [
+        "IMPORT_FILE = \"bq://bigquery-public-data.samples.gsod\"\n",
+        "BQ_TABLE = \"bigquery-public-data.samples.gsod\""
       ]
     },
     {
@@ -761,15 +851,63 @@
         "print(dataset.resource_name)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "82e9fe20ce71"
+      },
+      "source": [
+        "#### Dataframe input data\n",
+        "\n",
+        "Next, create the `Dataset` resource using the `create_from_dataframe` method for the `TabularDataset` class for pandas dataframe input, which takes the following parameters:\n",
+        "\n",
+        "- `display_name`: The human readable name for the `Dataset` resource.\n",
+        "- `df_source`: The pandas dataframe to import the data items into the `Dataset` resource.\n",
+        "- `staging_path`: The BigQuery table to store the imported data."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "import_file:covid,csv,forecast"
+        "id": "3805f945ffdd"
       },
       "outputs": [],
       "source": [
-        "IMPORT_FILE = \"gs://cloud-samples-data/ai-platform/covid/bigquery-public-covid-nyt-us-counties-train.csv\""
+        "# Download the table.\n",
+        "table = bigquery.TableReference.from_string(BQ_TABLE)\n",
+        "\n",
+        "rows = bqclient.list_rows(\n",
+        "    table,\n",
+        "    max_results=10000,\n",
+        "    selected_fields=[\n",
+        "        bigquery.SchemaField(\"station_number\", \"STRING\"),\n",
+        "        bigquery.SchemaField(\"year\", \"INTEGER\"),\n",
+        "        bigquery.SchemaField(\"month\", \"INTEGER\"),\n",
+        "        bigquery.SchemaField(\"day\", \"INTEGER\"),\n",
+        "        bigquery.SchemaField(\"mean_temp\", \"FLOAT\"),\n",
+        "    ],\n",
+        ")\n",
+        "\n",
+        "dataframe = rows.to_dataframe()\n",
+        "print(dataframe.head())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "create_dataset:tabular,bq,lrg,v2"
+      },
+      "outputs": [],
+      "source": [
+        "dataset = aip.TabularDataset.create_from_dataframe(\n",
+        "    display_name=\"example\" + \"_\" + TIMESTAMP,\n",
+        "    df_source=dataframe,\n",
+        "    staging_path=f\"bq://{PROJECT_ID}.samples.gsod\",\n",
+        ")\n",
+        "\n",
+        "print(dataset.resource_name)"
       ]
     },
     {
@@ -778,7 +916,7 @@
         "id": "create_dataset:tabular,forecast,v2"
       },
       "source": [
-        "### Create the Dataset\n",
+        "### Create a Time Series Dataset\n",
         "\n",
         "Next, create the `Dataset` resource using the `create` method for the `TimeSeriesDataset` class, which takes the following parameters:\n",
         "\n",
@@ -789,6 +927,17 @@
         "Learn more about [TimeSeriesDataset](https://cloud.google.com/vertex-ai/docs/datasets/prepare-tabular)."
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "import_file:covid,csv,forecast"
+      },
+      "outputs": [],
+      "source": [
+        "IMPORT_FILE = \"gs://cloud-samples-data/ai-platform/covid/bigquery-public-covid-nyt-us-counties-train.csv\""
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -1149,9 +1298,9 @@
         "comps = BQ_TABLE.split(\".\")\n",
         "BQ_PROJECT_DATASET_TABLE = comps[0] + \":\" + comps[1] + \".\" + comps[2]\n",
         "\n",
-        "! bq --location=us extract --destination_format CSV $BQ_PROJECT_DATASET_TABLE $BUCKET_NAME/mydata*.csv\n",
+        "! bq --location=us extract --destination_format CSV $BQ_PROJECT_DATASET_TABLE $BUCKET_URI/mydata*.csv\n",
         "\n",
-        "IMPORT_FILES = ! gsutil ls $BUCKET_NAME/mydata*.csv\n",
+        "IMPORT_FILES = ! gsutil ls $BUCKET_URI/mydata*.csv\n",
         "\n",
         "print(IMPORT_FILES)\n",
         "\n",
@@ -1209,6 +1358,38 @@
         "To create a dataframe from multiple CSV sources, you read each CSV file and concatenate the dataframes together."
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bcd2e4e0703b"
+      },
+      "source": [
+        "If you are running this notebook on Colab, run the following cell to install packages fsspec and gcsfs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "927bd3f92268"
+      },
+      "outputs": [],
+      "source": [
+        "# If you are running this notebook in Colab, run this cell and follow the\n",
+        "# instructions to authenticate your GCP account. This provides access to your\n",
+        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
+        "# requests.\n",
+        "\n",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "# If on Workbench AI Notebook, then don't execute this code\n",
+        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\"):\n",
+        "    if \"google.colab\" in sys.modules:\n",
+        "        ! pip3 install fsspec\n",
+        "        ! pip3 install gcsfs"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -1269,7 +1450,7 @@
       },
       "outputs": [],
       "source": [
-        "EXPORTED_DIR = f\"{BUCKET_NAME}/exported\"\n",
+        "EXPORTED_DIR = f\"{BUCKET_URI}/exported\"\n",
         "exported_files = dataset.export_data(output_dir=EXPORTED_DIR)\n",
         "\n",
         "! gsutil ls $EXPORTED_DIR"
@@ -1498,7 +1679,7 @@
         "    data = f.readlines()\n",
         "\n",
         "# The path to the TFRecord cached file.\n",
-        "GCS_TFRECORD_URI = BUCKET_NAME + \"/flowers.tfrecord\"\n",
+        "GCS_TFRECORD_URI = BUCKET_URI + \"/flowers.tfrecord\"\n",
         "\n",
         "# Create the TFRecord cached file\n",
         "with tf.io.TFRecordWriter(GCS_TFRECORD_URI) as writer:\n",
@@ -1532,14 +1713,7 @@
         "Otherwise, you can delete the individual resources you created in this tutorial:\n",
         "\n",
         "- Dataset\n",
-        "- Pipeline\n",
-        "- Model\n",
-        "- Endpoint\n",
-        "- AutoML Training Job\n",
-        "- Batch Job\n",
-        "- Custom Job\n",
-        "- Hyperparameter Tuning Job\n",
-        "- Cloud Storage Bucket"
+        "- Bucket"
       ]
     },
     {
@@ -1550,61 +1724,16 @@
       },
       "outputs": [],
       "source": [
-        "delete_all = True\n",
+        "delete_bucket = False\n",
         "\n",
-        "if delete_all:\n",
-        "    # Delete the dataset using the Vertex dataset object\n",
-        "    try:\n",
-        "        if \"dataset\" in globals():\n",
-        "            dataset.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the dataset using the Vertex dataset object\n",
+        "datasets = aip.TabularDataset.list(filter=f'display_name=\"example_{TIMESTAMP}\"')\n",
+        "for dataset in datasets:\n",
+        "    dataset.delete()\n",
         "\n",
-        "    # Delete the model using the Vertex model object\n",
-        "    try:\n",
-        "        if \"model\" in globals():\n",
-        "            model.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the endpoint using the Vertex endpoint object\n",
-        "    try:\n",
-        "        if \"endpoint\" in globals():\n",
-        "            endpoint.undeploy_all()\n",
-        "            endpoint.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the AutoML or Pipeline training job\n",
-        "    try:\n",
-        "        if \"dag\" in globals():\n",
-        "            dag.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the custom training job\n",
-        "    try:\n",
-        "        if \"job\" in globals():\n",
-        "            job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the batch prediction job using the Vertex batch prediction object\n",
-        "    try:\n",
-        "        if \"batch_predict_job\" in globals():\n",
-        "            batch_predict_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the hyperparameter tuning job using the Vertex hyperparameter tuning object\n",
-        "    try:\n",
-        "        if \"hpt_job\" in globals():\n",
-        "            hpt_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    if \"BUCKET_NAME\" in globals():\n",
-        "        ! gsutil rm -r $BUCKET_NAME"
+        "# Delete the bucket\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil rm -r $BUCKET_URI"
       ]
     }
   ],

notebooks/community/ml_ops/stage1/mlops_data_management.ipynb

@@ -39,12 +39,15 @@
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/ai/platform/notebooks/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage1/mlops_data_management.ipynb\">\n",
-        "      Open in Google Cloud Notebooks\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage1/mlops_data_management.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>\n",
         "</table>\n",
-        "<br/><br/><br/>"
+        "<br/><br/><br/>\n",
+        "\n",
+        "*Note: This notebook is not supported for execution in Colab*"
       ]
     },
     {
@@ -59,17 +62,6 @@
         "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 1 : data management."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dataset:bq,chicago,lbn"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "The dataset used for this tutorial is the [Chicago Taxi](https://www.kaggle.com/chicago/chicago-taxi-trips-bq). The version of the dataset you will use in this tutorial is stored in a public BigQuery table. The trained model predicts whether someone would leave a tip for a taxi fare."
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -114,6 +106,34 @@
         " - Preprocess the data with `Dataflow`"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dataset:bq,chicago,lbn"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "The dataset used for this tutorial is the [Chicago Taxi](https://www.kaggle.com/chicago/chicago-taxi-trips-bq). The version of the dataset used in this tutorial is stored in a public BigQuery table. The trained model predicts whether someone leaves a tip for a taxi fare."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9e483012a752"
+      },
+      "source": [
+        "### Costs\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "- Vertex AI\n",
+        "- Cloud Storage\n",
+        "- BigQuery\n",
+        "- Dataflow\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [BigQuery pricing](https://cloud.google.com/bigquery/pricing), and [Dataflow pricing](https://cloud.google.com/dataflow/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -133,20 +153,33 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\") and not os.getenv(\"VIRTUAL_ENV\")\n",
+        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
+        "    \"/opt/deeplearning/metadata/env_version\"\n",
+        ")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
         "ONCE_ONLY = False\n",
         "if ONCE_ONLY:\n",
-        "    ! pip3 install -U tensorflow==2.5 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-data-validation==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-transform==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-io==0.18 $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-bigquery $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-logging $USER_FLAG\n",
-        "    ! pip3 install --upgrade apache-beam[gcp] $USER_FLAG\n",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
-        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG\n",
-        "    ! pip3 install --upgrade kfp $USER_FLAG"
+        "    ! pip3 install -U {USER_FLAG} -q tensorflow==2.5 \\\n",
+        "                                     tensorflow-data-validation==1.2 \\\n",
+        "                                     tensorflow-transform==1.2 \\\n",
+        "                                     tensorflow-io==0.18 \n",
+        "    \n",
+        "    ! pip3 install --upgrade {USER_FLAG} -q google-cloud-aiplatform[tensorboard] \\\n",
+        "                                            google-cloud-pipeline-components \\\n",
+        "                                            google-cloud-bigquery \\\n",
+        "                                            google-cloud-logging \\\n",
+        "                                            apache-beam[gcp] \\\n",
+        "                                            pyarrow \\\n",
+        "                                            cloudml-hypertune\n"
       ]
     },
     {
@@ -178,6 +211,32 @@
         "    app.kernel.do_shutdown(True)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "84cd83853240"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI, BigQuery, Compute Engine and Cloud Storage APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,bigquery,compute_component,storage_component).\n",
+        "\n",
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -254,7 +313,10 @@
       },
       "outputs": [],
       "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if REGION == \"[your-region]\":\n",
+        "    REGION = \"us-central1\""
       ]
     },
     {
@@ -281,6 +343,66 @@
         "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "77c385f0db59"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "**If you are using Vertex AI Workbench Notebooks**, your environment is already authenticated. \n",
+        "\n",
+        "**If you are using Colab**, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.\n",
+        "\n",
+        "**Otherwise**, follow these steps:\n",
+        "\n",
+        "In the Cloud Console, go to the [Create service account key](https://console.cloud.google.com/apis/credentials/serviceaccountkey) page.\n",
+        "\n",
+        "1. **Click Create service account**.\n",
+        "\n",
+        "2. In the **Service account name** field, enter a name, and click **Create**.\n",
+        "\n",
+        "3. In the **Grant this service account access to project** section, click the Role drop-down list. Type \"Vertex AI\" into the filter box, and select **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
+        "\n",
+        "4. Click Create. A JSON file that contains your key downloads to your local environment.\n",
+        "\n",
+        "5. Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "535223fa4b84"
+      },
+      "outputs": [],
+      "source": [
+        "# If you are running this notebook in Colab, run this cell and follow the\n",
+        "# instructions to authenticate your GCP account. This provides access to your\n",
+        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
+        "# requests.\n",
+        "\n",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "# If on Vertex AI Workbench, then don't execute this code\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
+        "    \"DL_ANACONDA_HOME\"\n",
+        "):\n",
+        "    if \"google.colab\" in sys.modules:\n",
+        "        from google.colab import auth as google_auth\n",
+        "\n",
+        "        google_auth.authenticate_user()\n",
+        "\n",
+        "    # If you are running this notebook locally, replace the string below with the\n",
+        "    # path to your service account key and run this cell to authenticate your GCP\n",
+        "    # account.\n",
+        "    elif not os.getenv(\"IS_TESTING\"):\n",
+        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -291,7 +413,7 @@
         "\n",
         "**The following steps are required, regardless of your notebook environment.**\n",
         "\n",
-        "When you submit a custom training job using the Vertex SDK, you upload a Python package\n",
+        "When you submit a custom training job using the Vertex AI SDK, you upload a Python package\n",
         "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
         "the code from this package. In this tutorial, Vertex AI also saves the\n",
         "trained model that results from your job in the same bucket. You can then\n",
@@ -309,7 +431,8 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
+        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
+        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
       ]
     },
     {
@@ -320,8 +443,9 @@
       },
       "outputs": [],
       "source": [
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
+        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
+        "    BUCKET_NAME = PROJECT_ID + \"aip-\" + TIMESTAMP\n",
+        "    BUCKET_URI = \"gs://\" + BUCKET_NAME"
       ]
     },
     {
@@ -341,7 +465,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
+        "! gsutil mb -l $REGION $BUCKET_URI"
       ]
     },
     {
@@ -361,7 +485,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls -al $BUCKET_NAME"
+        "! gsutil ls -al $BUCKET_URI"
       ]
     },
     {
@@ -516,7 +640,7 @@
       },
       "outputs": [],
       "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
+        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
       ]
     },
     {
@@ -538,7 +662,7 @@
       },
       "outputs": [],
       "source": [
-        "bqclient = bigquery.Client()"
+        "bqclient = bigquery.Client(project=PROJECT_ID)"
       ]
     },
     {
@@ -661,6 +785,11 @@
         "LIMIT = 300000\n",
         "YEAR = 2020\n",
         "\n",
+        "# First, create the dataset entry\n",
+        "dataset = bigquery.Dataset(f\"{PROJECT_ID}.{BQ_DATASET}\")\n",
+        "dataset.location = \"US\"\n",
+        "dataset = bqclient.create_dataset(dataset, timeout=30)\n",
+        "\n",
         "query = f\"\"\"\n",
         "CREATE OR REPLACE TABLE `{BQ_TABLE_COPY}`\n",
         "AS (\n",
@@ -757,7 +886,7 @@
         "dataset = aip.TabularDataset.create(\n",
         "    display_name=\"Chicago Taxi\" + \"_\" + TIMESTAMP,\n",
         "    bq_source=[IMPORT_FILE],\n",
-        "    labels={\"user_metadata\": BUCKET_NAME[5:]},\n",
+        "    labels={\"user_metadata\": BUCKET_NAME},\n",
         ")\n",
         "\n",
         "label_column = \"tip_bin\"\n",
@@ -949,9 +1078,9 @@
       },
       "outputs": [],
       "source": [
-        "STATISTICS_SCHEMA = BUCKET_NAME + \"/statistics.jsonl\"\n",
+        "STATISTICS_SCHEMA = BUCKET_URI + \"/statistics.jsonl\"\n",
         "\n",
-        "tfdv.write_stats_text(stats, BUCKET_NAME + \"/statistics.jsonl\")\n",
+        "tfdv.write_stats_text(stats, BUCKET_URI + \"/statistics.jsonl\")\n",
         "\n",
         "with tf.io.gfile.GFile(\n",
         "    \"gs://\" + dataset.labels[\"user_metadata\"] + \"/metadata.jsonl\", \"r\"\n",
@@ -964,7 +1093,7 @@
         ") as f:\n",
         "    json.dump(metadata, f)\n",
         "\n",
-        "!gsutil cat $BUCKET_NAME/metadata.jsonl"
+        "! gsutil cat $BUCKET_URI/metadata.jsonl"
       ]
     },
     {
@@ -1011,7 +1140,7 @@
       },
       "outputs": [],
       "source": [
-        "SCHEMA_LOCATION = BUCKET_NAME + \"/schema.txt\"\n",
+        "SCHEMA_LOCATION = BUCKET_URI + \"/schema.txt\"\n",
         "\n",
         "# When running Apache Beam directly (file is directly accessed)\n",
         "tfdv.write_schema_text(output_path=SCHEMA_LOCATION, schema=schema)\n",
@@ -1049,7 +1178,7 @@
         ") as f:\n",
         "    json.dump(metadata, f)\n",
         "\n",
-        "!gsutil cat $BUCKET_NAME/metadata.jsonl"
+        "! gsutil cat $BUCKET_URI/metadata.jsonl"
       ]
     },
     {
@@ -1101,7 +1230,7 @@
         "import setuptools\n",
         "\n",
         "REQUIRED_PACKAGES = [\n",
-        "    \"google-cloud-aiplatform==1.4.2\",\n",
+        "    \"google-cloud-aiplatform\",\n",
         "    \"tensorflow-transform==1.2.0\",\n",
         "    \"tensorflow-data-validation==1.2.0\",\n",
         "]\n",
@@ -1372,10 +1501,10 @@
         "            )\n",
         "\n",
         "\n",
-        "EXPORTED_JSONL_PREFIX = os.path.join(BUCKET_NAME, \"exported_data/jsonl\")\n",
-        "EXPORTED_TFREC_PREFIX = os.path.join(BUCKET_NAME, \"exported_data/tfrec\")\n",
-        "TRANSFORMED_DATA_PREFIX = os.path.join(BUCKET_NAME, \"transformed_data\")\n",
-        "TRANSFORM_ARTIFACTS_DIR = os.path.join(BUCKET_NAME, \"transformed_artifacts\")\n",
+        "EXPORTED_JSONL_PREFIX = os.path.join(BUCKET_URI, \"exported_data/jsonl\")\n",
+        "EXPORTED_TFREC_PREFIX = os.path.join(BUCKET_URI, \"exported_data/tfrec\")\n",
+        "TRANSFORMED_DATA_PREFIX = os.path.join(BUCKET_URI, \"transformed_data\")\n",
+        "TRANSFORM_ARTIFACTS_DIR = os.path.join(BUCKET_URI, \"transformed_artifacts\")\n",
         "\n",
         "QUERY_STRING = \"SELECT * FROM {} LIMIT 300000\".format(BQ_TABLE)\n",
         "JOB_NAME = \"chicago\" + TIMESTAMP\n",
@@ -1388,7 +1517,7 @@
         "    \"transform_artifact_dir\": TRANSFORM_ARTIFACTS_DIR,\n",
         "    \"exported_jsonl_prefix\": EXPORTED_JSONL_PREFIX,\n",
         "    \"exported_tfrec_prefix\": EXPORTED_TFREC_PREFIX,\n",
-        "    \"temp_location\": os.path.join(BUCKET_NAME, \"temp\"),\n",
+        "    \"temp_location\": os.path.join(BUCKET_URI, \"temp\"),\n",
         "    \"project\": PROJECT_ID,\n",
         "    \"region\": REGION,\n",
         "    \"setup_file\": \"./setup.py\",\n",
@@ -1459,7 +1588,7 @@
         ") as f:\n",
         "    json.dump(metadata, f)\n",
         "\n",
-        "!gsutil cat $BUCKET_NAME/metadata.jsonl"
+        "! gsutil cat $BUCKET_URI/metadata.jsonl"
       ]
     },
     {
@@ -1473,17 +1602,9 @@
         "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
         "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
         "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial.\n",
         "\n",
-        "- Dataset\n",
-        "- Pipeline\n",
-        "- Model\n",
-        "- Endpoint\n",
-        "- AutoML Training Job\n",
-        "- Batch Job\n",
-        "- Custom Job\n",
-        "- Hyperparameter Tuning Job\n",
-        "- Cloud Storage Bucket"
+        "*Note:* stage2/mlops_experimentation is dependent on the resources created by this stage1 notebook."
       ]
     },
     {
@@ -1504,8 +1625,8 @@
         "    except Exception as e:\n",
         "        print(e)\n",
         "\n",
-        "    if \"BUCKET_NAME\" in globals():\n",
-        "        ! gsutil rm -r $BUCKET_NAME"
+        "    if \"BUCKET_URI\" in globals():\n",
+        "        ! gsutil rm -r $BUCKET_URI"
       ]
     }
   ],

notebooks/community/ml_ops/stage2/README.md

@@ -25,65 +25,313 @@ The second stage in MLOps is experimenting in developing one or more baseline mo
 - Use the What-if-Tool (WIT) to explore how the trained model would make predictions in different scenarios.
 
 
-<img src='stage2.png'>
+<img src='stage2v3.png'>
+<br/>
+<br/>
+<br/>
+<img src='stage2.2v1.png'>
 
 ## Notebooks
 
 ### Get Started
 
-[Get Started with Vertex Experiments and Vertex ML Metadata](get_started_vertex_experiments.ipynb)
+[Get started with Vertex AI Training for R](community/ml_ops/stage2/get_started_vertex_training_r.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Training` for training a R custom model.
+
+The steps performed include:
+
+- Locally train an R model in a notebook using %%R magic commands
+- Create a deployment image with trained R model and serving functions.
+- Test the deployment image locally.
+- Create a `Vertex AI Model` resource for the deployment image with embedded R model.
+- Deploy the deployment image with embedded R model to a `Vertex AI Endpoint` resource.
+- Test the deployment image with embedded R model.
+- Create a R-to-Python training package.
+- Create a training image for training the model.
+- Train a R model using `Vertex AI Trainingh` service with the R-to-Python training package.
+
+[Get started with Logging](community/ml_ops/stage2/get_started_with_logging.ipynb)
+
+In this tutorial, you learn how to use Python and Cloud logging awhen training with `Vertex AI`.
 
-```
 The steps performed include:
 
 - Use Python logging to log training configuration/results locally.
 - Use Google Cloud Logging to log training configuration/results in cloud storage.
-- Create a Vertex AI `Experiment` resource.
-- Instantiate an experiment run.
-- Log parameters for the run.
-- Log metrics for the run.
-- Display the logged experiment run.
-```
 
-[Get Started with Vertex TensorBoard](get_started_vertex_tensorboard.ipynb)
+[Get started with Vertex AI Hyperparameter Tuning for XGBoost] (community/ml_ops/stage2/get_started_vertex_hpt_xgboost.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Hyperparameter Tuning` for training a XGBoost custom model.
+
+The steps performed include:
+
+- Training using a Python package.
+- Report accuracy when hyperparameter tuning.
+- Save the model artifacts to Cloud Storage using GCSFuse.
+- Create a `Vertex AI Model` resource.
+
+[Get started with Vertex AI Training for XGBoost](community/ml_ops/stage2/get_started_vertex_training_xgboost.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Training` for training a XGBoost custom model.
+
+The steps performed include:
+
+- Training using a Python package.
+- Report accuracy when hyperparameter tuning.
+- Save the model artifacts to Cloud Storage using GCSFuse.
+- Create a `Vertex AI Model` resource.
+
+[Get started with TabNet builtin algorithm for training tabular models](community/ml_ops/stage2/get_started_with_tabnet.ipynb)
+
+In this notebook, you learn how to run `Vertex AI TabNet` built algorithm for training custom tabular models.
+
+The steps performed include:
+
+- Get the training data.
+- Configure training parameters for the `Vertex AI TabNet` container.
+- Train the model using `Vertex AI Training` using CSV data.
+- Upload the model as a `Vertex AI Model` resource.
+- Deploy the `Vertex AI Model` resource to a `Vertex AI Endpoint` resource.
+- Make a prediction with the deployed model.
+- Hyperparameter tuning the `Vertex AI TabNet` model.
+- Train the model using `Vertex AI Training` using BigQuery table.
+
+[Get started with prebuilt TFHub models](community/ml_ops/stage2/get_started_with_tfhub_models.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Training` with prebuilt models from TensorFlow Hub.
+
+The steps performed include:
+
+- Download a TensorFlow Hub prebuilt model.
+- Add the task component as a classifier for the CIFAR-10 dataset.
+- Fine tune locally the model with transfer learning training.
+- Construct a custom training script:
+    - Get training data from TensorFlow Datasets
+    - Get model architecture from TensorFlow Hub
+    - Train then model
+    - Save model artifacts and upload as Vertex AI Model resource.
+
+[Get started with BigQuery ML Training](community/ml_ops/stage2/get_started_bqml_training.ipynb)
+
+In this tutorial, you learn how to use `BigQueryML` (BQML) for training with `Vertex AI`.
+
+The steps performed include:
+
+- Create a local BigQuery table in your project
+- Train a BQML model
+- Evaluate the BQML model
+- Export the BQML model as a cloud model
+- Upload the exported model as a `Vertex AI Model` resource
+- Hyperparameter tune a BQML model with `Vertex AI Vizier`
+- Automatically register a BQML model to `Vertex AI Model Registry`
+
+[Get started with Vertex AI Vizier](community/ml_ops/stage2/get_started_vertex_vizier.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Vizier` for when training with `Vertex AI`.
+
+The steps performed include:
+
+- Hyperparameter tuning with Random algorithm.
+- Hyperparameter tuning with Vizier (Bayesian) algorithm.
+- Suggesting trials and updating results for Vizier study
+
+[Get started with distributed training using DASK](community/ml_ops/stage2/get_started_with_distributed_training_xgboost.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Training` for distributed training of XGBoost model using the OSS package DASK. Additionally, you learn to construct and deploy a custom serving container using a Flask web server.
+
+The steps performed include:
+
+- Construct an XGBoost training script using DASK for distributed training.
+- Construct a custom training container.
+- Configure a distributed custom training job.
+- Execute the custom training job.
+- Construct a custom serving container using Flask.
+- Upload the trained XGBoost model as a `Vertex AI Model` resource.
+- Create a `Vertex AI Endpoint` resource.
+- Deploy the `Vertex AI Model` resource to `Vertex AI Endpoint` resource.
+- Make a prediction.
+
+[Get started with Vertex AI TensorBoard](community/ml_ops/stage2/get_started_vertex_tensorboard.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI TensorBoard` when training with `Vertex AI`.
 
-```
 The steps performed include:
 
 - Create a TensorBoard callback when training a model.
-- Using Tensorboard with locally trained model.
+- Using TensorBoard with locally trained model.
 - Using Vertex AI TensorBoard with Vertex AI Training.
-```
 
-[Get Started with Custom Training Packages (Tensorflow)](get_started_vertex_training.ipynb)
+[Get started with Vertex AI Training for R using R Kernel](community/ml_ops/stage2/get_started_vertex_training_r_using_r_kernel.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI`, using an R kernel, for training and deploying an R custom model.
+
+The steps performed include:
+
+- Create a custom R training script
+- Create a custom R serving script
+- Create a custom R deployment (serving) container.
+- Train the model using `Vertex AI` custom training.
+- Create an `Endpoint` resouce.
+- Deploy the `Model` resource (trained R model) to the `Endpoint` resource.
+- Make an online prediction.
+
+
+[Get started Vision API test preprocessing and AutoML text model generation](community/ml_ops/stage2/get_started_with_visionapi_and_automl.ipynb)
+
+In this tutorial, you create an `AutoML` text entity extraction model pre-existing extracted data by generating a custom import file. You deploy this mode for online prediction from a Python script using the `BigQuery`, `Vision AI`, Cloud Storage and `Vertex AI SDK` for Python. 
+
+The steps performed include:
+
+- Preprocess training files using `Vision AI` APIs to extract the text from PDF files.
+- Create a custom import file that includes annotation data based on the sample `BigQuery` dataset.
+- Create a `Vertex AI Dataset` resource.
+- Train the model.
+- View the model evaluation.
+- Deploy the `Vertex AI Model` resource to a serving `Endpoint` resource.
+- Make a prediction.
+- Undeploy the `Model`.
+
+[Get started with Vertex AI Experiments](community/ml_ops/stage2/get_started_vertex_experiments.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Experiments` when training with `Vertex AI`.
+
+The steps performed include:
+
+- Local (notebook) Training
+    - Create an experiment
+    - Create a first run in the experiment
+    - Log parameters and metrics
+    - Create artifact lineage
+    - Visualize the experiment results
+    - Execute a second run
+    - Compare the two runs in the experiment
+- Cloud (`Vertex AI`) Training
+    - Within the training script:
+        - Create an experiment
+        - Log parameters and metrics
+        - Create artifact lineage
+    - Create a `Vertex AI Training` custom job
+    - Execute the custom job
+    - Visualize the experiment results
+
+[AutoML Image Classfication Training with Customer Managed Encryption Keys (CMEK)](community/ml_ops/stage2/get_started_with_cmek_training.ipynb)
+
+In this tutorial, you learn how to use a customer managed encryption key (CMEK) for `Vertex AI AutoML` training.
+
+The steps performed include:
+
+- Creating a customer managed encryption key.
+- Creating an image dataset with CMEK encryption.
+- Train an AutoML model with CMEK encryption.
+
+[Get started with Vertex AI Feature Store](community/ml_ops/stage2/get_started_vertex_feature_store.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Feature Store` when training and predicting with `Vertex AI`.
+
+The steps performed include:
+
+- Creating a Vertex AI `Featurestore` resource.
+    - Creating `EntityType` resources for the `Featurestore` resource.
+    - Creating `Feature` resources for each `EntityType` resource.
+- Import feature values (entity data items) into `Featurestore` resource.
+    - From a Cloud Storage location.
+    - From a pandas DataFrame.
+- Perform online serving from a `Featurestore` resource.
+- Perform batch serving from a `Featurestore` resource.
+
+[Get started with AutoML Training](community/ml_ops/stage2/get_started_automl_training.ipynb)
+
+In this tutorial, you learn how to use `AutoML` for training with `Vertex AI`.
+
+The steps performed include:
+
+- Train an image model
+- Export the image model as an edge model
+- Train a tabular model
+- Export the tabular model as a cloud model
+- Train a text model
+- Train a video model
+
+[Get started with Vertex AI Training for LightGBM](community/ml_ops/stage2/get_started_vertex_training_lightgbm.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Training` for training a LightGBM custom model.
+
+The steps performed include:
+
+- Training using a Python package.
+- Save the model artifacts to Cloud Storage using GCSFuse.
+- Construct a FastAPI prediction server.
+- Construct a Dockerfile deployment image.
+- Test the deployment image locally.
+- Create a `Vertex AI Model` resource.
+
+[Get started with Vertex AI Training for Scikit-Learn](community/ml_ops/stage2/get_started_vertex_training_sklearn.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Training` for training a Scikit-Learn custom model.
+
+The steps performed include:
+
+- Training using a Python package.
+- Report accuracy when hyperparameter tuning.
+- Save the model artifacts to Cloud Storage using GCSFuse.
+- Create a `Vertex AI Model` resource.
+
+[Get started with Vertex AI Training](community/ml_ops/stage2/get_started_vertex_training.ipynb)
+
+In this tutorial, you learn how to use `Vertex AI Training` for custom models when training with `Vertex AI`.
 
-```
 The steps performed include:
 
 - Training using a single Python script.
 - Training using a Python package.
 - Training using a custom training image.
 - Laying out a training package.
-```
 
-[Get Started with Custom Training Packages (Scikit-Learn)](get_started_vertex_training_sklearn.ipynb)
 
-[Get Started with Custom Training Packages (XGBoost)](get_started_vertex_training_xgboost.ipynb)
+[Get started with Vertex AI Training for Pytorch](community/ml_ops/stage2/get_started_vertex_training_pytorch.ipynb)
 
-[Get Started with Custom Training Packages (Pytorch)](get_started_vertex_training_pytorch.ipynb)
+In this tutorial, you learn how to use `Vertex AI Training` for training a Pytorch custom model.
 
-[Get Started with Custom Training Packages (R)](get_started_vertex_training_r.ipynb)
+The steps performed include:
 
-[Get Started with Distributed Training](get_started_vertex_distributed_training.ipynb)
+- Single node training using a Python package.
+- Report accuracy when hyperparameter tuning.
+- Save the model artifacts to Cloud Storage using GCSFuse.
+- Create a `Vertex AI Model` resource.
 
-[Get Started with Vizier Hyperparameter Tuning](get_started_vertex_vizier.ipynb)
+[Get started with Vertex AI Distributed Training](community/ml_ops/stage2/get_started_vertex_distributed_training.ipynb)
 
-[Get Started with AutoML Training](get_started_automl_training.ipynb)
+In this tutorial, you learn how to use `Vertex AI Distributed Training` for when training with `Vertex AI`.
 
-[Get Started with BQML Training](get_started_bqml_training.ipynb)
+The steps performed include:
 
-[Get Started with Vertex Feature Store](get_started_vertex_feature_store.ipynb)
+- `MirroredStrategy`: Train on a single VM with multiple GPUs.
+- `MultiWorkerMirroredStrategy`: Train on multiple VMs with automatic setup of replicas.
+- `MultiWorkerMirroredStrategy`: Train on multiple VMs with fine grain control of replicas.
+- `ReductionServer`: Train on multiple VMS and sync updates across VMS with `Vertex AI Reduction Server`.
+- `TPUTraining`: Train with multiple Cloud TPUs.
 
 ### E2E Stage Example
 
 [Stage 2: Experimentation](mlops_experimentation.ipynb)
+
+```
+The steps performed include:
+- Review the `Dataset` resource created during stage 1.
+- Train an AutoML tabular binary classifier model in the background.
+- Build the experimental model architecture.
+- Construct a custom training package for the `Dataset` resource.
+- Test the custom training package locally.
+- Test the custom training package in the cloud with Vertex AI Training.
+- Hyperparameter tune the model training with Vertex AI Vizier.
+- Train the custom model with Vertex AI Training.
+- Add a serving function for online/batch prediction to the custom model.
+- Test the custom model with the serving function.
+- Evaluate the custom model using Vertex AI Batch Prediction
+- Wait for the AutoML training job to complete.
+- Evaluate the AutoML model using Vertex AI Batch Prediction with the same evaluation slices as the custom model.
+- Set the evaluation results of the AutoML model as the baseline.
+- If the evaluation of the custom model is below baseline, continue to experiment with the custom model.
+- If the evaluation of the custom model is above baseline, save the model as the first best model.
+```

notebooks/community/ml_ops/stage2/get_started_bqml_training.ipynb

@@ -8,7 +8,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2021 Google LLC\n",
+        "# Copyright 2022 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -38,9 +38,15 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
+        "    <td>\n",
+        "        <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage2/get_started_bqml_training.ipynb\">\n",
+        "        <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "        </a>\n",
+        "  </td>\n",
         "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/ai/platform/notebooks/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage2/get_started_bqml_training.ipynb\">\n",
-        "      Open in Google Cloud Notebooks\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage2/get_started_bqml_training.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>\n",
         "</table>\n",
@@ -59,17 +65,6 @@
         "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 2 : experimentation: get started with BigQuery ML Training."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dataset:penguins,lcn,bq"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "The dataset used for this tutorial is the Penguins dataset from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). The version of the dataset predicts the species."
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -78,22 +73,50 @@
       "source": [
         "### Objective\n",
         "\n",
-        "In this tutorial, you learn how to use `BigQueryML` (BQML) for training with `Vertex AI`.\n",
+        "In this tutorial, you learn how to use `BigQueryML` for training with `Vertex AI`.\n",
         "\n",
         "This tutorial uses the following Google Cloud ML services:\n",
         "\n",
         "- `BigQueryML Training`\n",
         "- `Vertex AI Model resource`\n",
-        "- `Vertex AI Vizier.\n",
+        "- `Vertex AI Vizier`\n",
         "\n",
         "The steps performed include:\n",
         "\n",
-        "- Create a local BQ table in your project.\n",
-        "- Train a BQML model.\n",
-        "- Evaluate the BQML model.\n",
-        "- Export the BQML model as a cloud model.\n",
-        "- Upload the exported model as a Vertex AI Model resource.\n",
-        "- Hyperparameter tune a BQML model with Vertex AI Vizier."
+        "- Create a local BigQuery table in your project\n",
+        "- Train a BigQuery ML model\n",
+        "- Evaluate the BigQuery ML model\n",
+        "- Export the BigQuery ML model as a cloud model\n",
+        "- Upload the exported model as a `Vertex AI Model` resource\n",
+        "- Hyperparameter tune a BigQuery ML model with `Vertex AI Vizier`\n",
+        "- Automatically register a BigQuery ML model to `Vertex AI Model Registry`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dataset:penguins,lcn,bq"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "The dataset used for this tutorial is the Penguins dataset from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). This version of the dataset is used to predict the species of penguins from the available features like culmen-length, flipper-depth etc."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "81c777b8ad32"
+      },
+      "source": [
+        "### Costs\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "- Vertex AI\n",
+        "- Cloud Storage\n",
+        "- BigQuery\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing) and [BigQuery pricing](https://cloud.google.com/bigquery/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
       ]
     },
     {
@@ -104,7 +127,7 @@
       "source": [
         "## Installations\n",
         "\n",
-        "Install *one time* the packages for executing the MLOps notebooks."
+        "Install the following packages for executing this notebook."
       ]
     },
     {
@@ -115,20 +138,22 @@
       },
       "outputs": [],
       "source": [
-        "ONCE_ONLY = False\n",
-        "if ONCE_ONLY:\n",
-        "    ! pip3 install -U tensorflow==2.5 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-data-validation==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-transform==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-io==0.18 $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-bigquery $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-logging $USER_FLAG\n",
-        "    ! pip3 install --upgrade apache-beam[gcp] $USER_FLAG\n",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
-        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG\n",
-        "    ! pip3 install --upgrade kfp $USER_FLAG"
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\") and not os.getenv(\"VIRTUAL_ENV\")\n",
+        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
+        "# Install the packages\n",
+        "! pip3 install --upgrade pyarrow \\\n",
+        "                        google-cloud-aiplatform \\\n",
+        "                        google-cloud-bigquery \\\n",
+        "                        google-cloud-bigquery-storage $USER_FLAG -q"
       ]
     },
     {
@@ -165,6 +190,32 @@
       "metadata": {
         "id": "project_id"
       },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI, BigQuery, Compute Engine and Cloud Storage APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,bigquery,compute_component,storage_component).\n",
+        "\n",
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "56d591439df1"
+      },
       "source": [
         "#### Set your project ID\n",
         "\n",
@@ -236,7 +287,10 @@
       },
       "outputs": [],
       "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if REGION == \"[your-region]\":\n",
+        "    REGION = \"us-central1\""
       ]
     },
     {
@@ -263,6 +317,67 @@
         "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f3bd8c0d0469"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "**If you are using Vertex AI Workbench Notebooks**, your environment is already authenticated. Skip this step.\n",
+        "\n",
+        "**If you are using Colab**, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.\n",
+        "\n",
+        "**Otherwise**, follow these steps:\n",
+        "\n",
+        "In the Cloud Console, go to the [Create service account key](https://console.cloud.google.com/apis/credentials/serviceaccountkey) page.\n",
+        "\n",
+        "1. **Click Create service account**.\n",
+        "\n",
+        "2. In the **Service account name** field, enter a name, and click **Create**.\n",
+        "\n",
+        "3. In the **Grant this service account access to project** section, click the Role drop-down list. Type \"Vertex AI\" into the filter box, and select **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
+        "\n",
+        "4. Click Create. A JSON file that contains your key downloads to your local environment.\n",
+        "\n",
+        "5. Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e0953a00668e"
+      },
+      "outputs": [],
+      "source": [
+        "# If you are running this notebook in Colab, run this cell and follow the\n",
+        "# instructions to authenticate your GCP account. This provides access to your\n",
+        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
+        "# requests.\n",
+        "\n",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "# If on Vertex AI Workbench, then don't execute this code\n",
+        "IS_COLAB = False\n",
+        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
+        "    \"DL_ANACONDA_HOME\"\n",
+        "):\n",
+        "    if \"google.colab\" in sys.modules:\n",
+        "        IS_COLAB = True\n",
+        "        from google.colab import auth as google_auth\n",
+        "\n",
+        "        google_auth.authenticate_user()\n",
+        "\n",
+        "    # If you are running this notebook locally, replace the string below with the\n",
+        "    # path to your service account key and run this cell to authenticate your GCP\n",
+        "    # account.\n",
+        "    elif not os.getenv(\"IS_TESTING\"):\n",
+        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -286,7 +401,8 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
+        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
+        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
       ]
     },
     {
@@ -297,8 +413,9 @@
       },
       "outputs": [],
       "source": [
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
+        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
+        "    BUCKET_NAME = PROJECT_ID + \"aip-\" + TIMESTAMP\n",
+        "    BUCKET_URI = f\"gs://{BUCKET_NAME}\""
       ]
     },
     {
@@ -318,7 +435,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
+        "! gsutil mb -l $REGION $BUCKET_URI"
       ]
     },
     {
@@ -338,7 +455,55 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls -al $BUCKET_NAME"
+        "! gsutil ls -al $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "set_service_account"
+      },
+      "source": [
+        "#### Service Account\n",
+        "\n",
+        "You use a service account to create Vertex AI Pipeline jobs. If you do not want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "set_service_account"
+      },
+      "outputs": [],
+      "source": [
+        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "autoset_service_account"
+      },
+      "outputs": [],
+      "source": [
+        "if (\n",
+        "    SERVICE_ACCOUNT == \"\"\n",
+        "    or SERVICE_ACCOUNT is None\n",
+        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
+        "):\n",
+        "    # Get your service account from gcloud\n",
+        "    if not IS_COLAB:\n",
+        "        shell_output = !gcloud auth list 2>/dev/null\n",
+        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
+        "\n",
+        "    if IS_COLAB:\n",
+        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
+        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
       ]
     },
     {
@@ -347,9 +512,6 @@
         "id": "setup_vars"
       },
       "source": [
-        "### Set up variables\n",
-        "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
         "### Import libraries and define constants"
       ]
     },
@@ -361,28 +523,7 @@
       },
       "outputs": [],
       "source": [
-        "import google.cloud.aiplatform as aip"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "import_bq"
-      },
-      "source": [
-        "#### Import BigQuery\n",
-        "\n",
-        "Import the BigQuery package into your Python environment."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_bq"
-      },
-      "outputs": [],
-      "source": [
+        "import google.cloud.aiplatform as aiplatform\n",
         "from google.cloud import bigquery"
       ]
     },
@@ -392,7 +533,7 @@
         "id": "init_aip:mbsdk"
       },
       "source": [
-        "### Initialize Vertex AI SDK for Python\n",
+        "### Initialize Vertex AI and BigQuery SDKs for Python\n",
         "\n",
         "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
       ]
@@ -405,7 +546,7 @@
       },
       "outputs": [],
       "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
+        "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
       ]
     },
     {
@@ -414,8 +555,6 @@
         "id": "init_bq"
       },
       "source": [
-        "### Create BigQuery client\n",
-        "\n",
         "Create the BigQuery client."
       ]
     },
@@ -427,7 +566,7 @@
       },
       "outputs": [],
       "source": [
-        "bqclient = bigquery.Client()"
+        "bqclient = bigquery.Client(project=PROJECT_ID)"
       ]
     },
     {
@@ -436,17 +575,17 @@
         "id": "accelerators:prediction,mbsdk"
       },
       "source": [
-        "#### Set hardware accelerators\n",
+        "### Set hardware accelerators\n",
         "\n",
         "You can set hardware accelerators for prediction.\n",
         "\n",
         "Set the variable `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
         "\n",
-        "    (aip.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
+        "    (aiplatform.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
         "\n",
         "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
         "\n",
-        "Learn more [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your region"
+        "Learn more [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your region."
       ]
     },
     {
@@ -457,13 +596,15 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
         "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
-        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
+        "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
         "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
         "    )\n",
         "else:\n",
-        "    DEPLOY_GPU, DEPLOY_NGPU = (aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 1)"
+        "    DEPLOY_GPU, DEPLOY_NGPU = (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80, 1)"
       ]
     },
     {
@@ -472,7 +613,7 @@
         "id": "container:prediction"
       },
       "source": [
-        "#### Set pre-built containers\n",
+        "### Set pre-built containers\n",
         "\n",
         "Set the pre-built Docker container image for prediction.\n",
         "\n",
@@ -519,11 +660,11 @@
         "id": "machine:prediction"
       },
       "source": [
-        "#### Set machine type\n",
+        "### Set machine type\n",
         "\n",
         "Next, set the machine type to use for prediction.\n",
         "\n",
-        "- Set the variable `DEPLOY_COMPUTE` to configure the compute resources for the VM you will use for prediction.\n",
+        "- Set the variable `DEPLOY_COMPUTE` to configure the compute resources for the VM which is used for prediction.\n",
         " - `machine type`\n",
         "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
         "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
@@ -557,7 +698,7 @@
         "id": "bqml_intro"
       },
       "source": [
-        "## Bigquery ML introduction\n",
+        "## BigQuery ML introduction\n",
         "\n",
         "BigQuery ML (BQML) provides the capability to train ML tabular models, such as classification and regression, in BigQuery using SQL syntax.\n",
         "\n",
@@ -582,9 +723,9 @@
         "id": "bqml_create_dataset"
       },
       "source": [
-        "### Create BQ dataset/model resource\n",
+        "### Create BQ dataset resource\n",
         "\n",
-        "First, you create a empty dataset/model resource in your project."
+        "First, you create an empty dataset resource in your project."
       ]
     },
     {
@@ -608,9 +749,9 @@
         "id": "bqml_create_model"
       },
       "source": [
-        "### Train BQML model\n",
+        "### Train BigQuery ML model\n",
         "\n",
-        "Next, you create and train a BQML tabular classification model from the public dataset penguins and store the model in your project using the `CREATE MODEL` statement. The model configuration is specified in the `OPTIONS` statement as follows:\n",
+        "Next, you create and train a BigQuery ML tabular classification model from the public dataset penguins and store the model in your project using the `CREATE MODEL` statement. The model configuration is specified in the `OPTIONS` statement as follows:\n",
         "\n",
         "- `model_type`: The type and archictecture of tabular model to train, e.g., DNN classification.\n",
         "- `labels`: The column which are the labels.\n",
@@ -659,9 +800,9 @@
         "id": "bqml_eval_model"
       },
       "source": [
-        "### Evaluate the BQML trained model\n",
+        "### Evaluate the trained BigQuery ML model\n",
         "\n",
-        "Next, retrieve the model evaluation for the trained BQML model.\n",
+        "Next, retrieve the model evaluation for the trained BigQuery ML model.\n",
         "\n",
         "Learn more about [The ML.EVALUATE function](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-evaluate)."
       ]
@@ -692,9 +833,9 @@
         "id": "bqml_export_model"
       },
       "source": [
-        "### Export the model from BQML\n",
+        "### Export the model from BigQuery ML\n",
         "\n",
-        "The model you trained in BQML is a TensorFlow model. Next, you will export the TensorFlow model artifacts in TF.SavedModel format."
+        "The model you trained in BigQuery ML is a TensorFlow model. Next, you export the TensorFlow model artifacts in TF.SavedModel format."
       ]
     },
     {
@@ -705,10 +846,10 @@
       },
       "outputs": [],
       "source": [
-        "param = f\"{PROJECT_ID}:{BQ_DATASET_NAME}.{MODEL_NAME} {BUCKET_NAME}/{MODEL_NAME}\"\n",
+        "param = f\"{PROJECT_ID}:{BQ_DATASET_NAME}.{MODEL_NAME} {BUCKET_URI}/{MODEL_NAME}\"\n",
         "! bq extract -m $param\n",
         "\n",
-        "MODEL_DIR = f\"{BUCKET_NAME}/{BQ_DATASET_NAME}\"\n",
+        "MODEL_DIR = f\"{BUCKET_URI}/{BQ_DATASET_NAME}\"\n",
         "! gsutil ls $MODEL_DIR"
       ]
     },
@@ -718,9 +859,25 @@
         "id": "upload_bqml_model"
       },
       "source": [
-        "## Upload the BQML model to a Model resource\n",
+        "## Upload the BigQuery ML model to a Vertex AI Model resource\n",
         "\n",
-        "Finally, now that you have the BQML model exported as a TF.SavedModel format, you upload the model artifacts to Vertex AI Model resource, in the same way as if you were uploading a custom trained model."
+        "Finally, now that you have the BigQuery ML model exported, you upload the model artifacts to Vertex AI Model resource, in the same way as if you were uploading a custom trained model.\n",
+        "\n",
+        "Below is a partial list of mapping BigQuery ML model types to their corresponding exported model format:\n",
+        "\n",
+        "'LINEAR_REG'<br/>\n",
+        "'LOGISTIC_REG' --> TensorFlow SavedFormat\n",
+        "\n",
+        "'AUTOML_CLASSIFIER'<br/>\n",
+        "'AUTOML_REGRESSOR' --> TensorFlow SavedFormat\n",
+        "\n",
+        "'BOOSTED_TREE_CLASSIFIER'<br/>\n",
+        "'BOOSTED_TREE_REGRESSOR' --> XGBoost format\n",
+        "\n",
+        "'DNN_CLASSIFIER'<br/>\n",
+        "'DNN_REGRESSOR'<br/>\n",
+        "'DNN_LINEAR_COMBINED_CLASSIFIER'<br/>\n",
+        "'DNN_LINEAR_COMBINED_REGRESSOR' --> TensorFlow Estimator"
       ]
     },
     {
@@ -731,7 +888,7 @@
       },
       "outputs": [],
       "source": [
-        "model = aip.Model.upload(\n",
+        "model = aiplatform.Model.upload(\n",
         "    display_name=\"penguins_\" + TIMESTAMP,\n",
         "    artifact_uri=MODEL_DIR,\n",
         "    serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -752,7 +909,7 @@
         "- `deployed_model_display_name`: A human readable name for the deployed model.\n",
         "- `traffic_split`: Percent of traffic at the endpoint that goes to this model, which is specified as a dictionary of one or more key/value pairs.\n",
         "If only one model, then specify as { \"0\": 100 }, where \"0\" refers to this model being uploaded and 100 means 100% of the traffic.\n",
-        "If there are existing models on the endpoint, for which the traffic will be split, then use model_id to specify as { \"0\": percent, model_id: percent, ... }, where model_id is the model id of an existing model to the deployed endpoint. The percents must add up to 100.\n",
+        "If there are existing models on the endpoint, for which the traffic needs to be split, then use model_id to specify as { \"0\": percent, model_id: percent, ... }, where model_id is the model id of an existing model to the deployed endpoint. The percents must add up to 100.\n",
         "- `machine_type`: The type of machine to use for training.\n",
         "- `accelerator_type`: The hardware accelerator type.\n",
         "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
@@ -801,7 +958,7 @@
         "id": "undeploy_model:mbsdk"
       },
       "source": [
-        "## Undeploy the model\n",
+        "#### Undeploy the model\n",
         "\n",
         "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
       ]
@@ -823,9 +980,9 @@
         "id": "model_delete:mbsdk"
       },
       "source": [
-        "#### Delete the model\n",
+        "#### Delete the `Vertex AI Model` resource\n",
         "\n",
-        "The method 'delete()' will delete the model."
+        "The method 'delete()' deletes the model."
       ]
     },
     {
@@ -839,15 +996,41 @@
         "model.delete()"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7890ae6f6410"
+      },
+      "source": [
+        "### Delete the `BigQuery ML` model\n",
+        "\n",
+        "Next, delete the `BigQuery ML` instance of the model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f0b6163e70c0"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_QUERY = f\"\"\"\n",
+        "DROP MODEL `{BQ_DATASET_NAME}.{MODEL_NAME}`\n",
+        "\"\"\"\n",
+        "\n",
+        "job = bqclient.query(MODEL_QUERY)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "bqml_create_model:vizier"
       },
       "source": [
-        "### Hyperparameter Tune and train a BQML model\n",
+        "### Hyperparameter Tune and train a BigQuery ML model\n",
         "\n",
-        "Next, you train a BQML tabular classification model with hyperparameter tuning using the Vertex AI Vizier service. The hyperparameter settings are specified in the `OPTIONS` statement as follows:\n",
+        "Next, you train a BigQuery ML tabular classification model with hyperparameter tuning using the `Vertex AI Vizier` service. The hyperparameter settings are specified in the `OPTIONS` statement as follows:\n",
         "\n",
         "- `HPARAM_TUNING_ALGORITHM`: The algorithm for selecting the next trial parameters.\n",
         "- `num_trials`: The number of trials.\n",
@@ -900,9 +1083,9 @@
         "id": "bqml_eval_model"
       },
       "source": [
-        "### Evaluate the BQML trained model\n",
+        "### Evaluate the BigQuery ML trained model\n",
         "\n",
-        "Next, retrieve the model evaluation for the trained BQML model.\n",
+        "Next, retrieve the model evaluation results for the trained BigQuery ML model.\n",
         "\n",
         "Learn more about [The ML.EVALUATE function](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-evaluate)."
       ]
@@ -927,15 +1110,41 @@
         "print(results)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3f3cee1236b1"
+      },
+      "source": [
+        "### Delete the `BigQuery ML` model\n",
+        "\n",
+        "Next, delete the `BigQuery ML` instance of the model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "957b7d841502"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_QUERY = f\"\"\"\n",
+        "DROP MODEL `{BQ_DATASET_NAME}.{MODEL_NAME}`\n",
+        "\"\"\"\n",
+        "\n",
+        "job = bqclient.query(MODEL_QUERY)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "bqml_create_model:xai"
       },
       "source": [
-        "### Train a BQML model with Explainability\n",
+        "### Train a BigQuery ML model with Explainability\n",
         "\n",
-        "Next, you train the same BQML model, but this time you enable Vertex AI Explainability on the model predictions by adding the option:\n",
+        "Next, you train the same BigQuery ML model, but this time you enable Vertex AI Explainability on the model predictions by adding the option:\n",
         "\n",
         "- `ENABLE_GLOBAL_EXPLAIN`"
       ]
@@ -976,6 +1185,179 @@
         "print(\"{} created in {}\".format(tblname, job.ended - job.started))"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4def8aaf3398"
+      },
+      "source": [
+        "### Delete the `BigQuery ML` model\n",
+        "\n",
+        "Next, delete the `BigQuery ML` instance of the model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ff5b32618018"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_QUERY = f\"\"\"\n",
+        "DROP MODEL `{BQ_DATASET_NAME}.{MODEL_NAME}`\n",
+        "\"\"\"\n",
+        "\n",
+        "job = bqclient.query(MODEL_QUERY)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2b4498ca6fea"
+      },
+      "source": [
+        "## Model Registry\n",
+        "\n",
+        "Alternatively, you can implicitly upload your BigQuery ML model as a `Vertex AI Model` resource with exporting and importing the model artifacts. In this method, you add additional options when training the model that tells BigQuery ML to automatically upload and register the trained model as a `Model` resource.\n",
+        "\n",
+        "### Setting permissions to automatically register the model\n",
+        "\n",
+        "You need to set some additional IAM permissions for BigQuery ML to automatically upload and register the model after training. Depending on your service account, the setting of the permissions below may fail. In this case, we recommend executing the permissions in a Cloud Shell.\n",
+        "\n",
+        "Learn more about [Setting permissions for Model Registry](https://cloud.google.com/bigquery-ml/docs/managing-models-vertex)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "29229f72d13d"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
+        "    --member=serviceAccount:$SERVICE_ACCOUNT --role=roles/aiplatform.admin --condition=None"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6c390ee7c11a"
+      },
+      "source": [
+        "### Training and registering the model\n",
+        "\n",
+        "Next, you train the model and automatically register the model to the `Vertex AI Model Registry`, by adding the following parameters as options:\n",
+        "\n",
+        "- `model_registry`: Set to \"vertex_ai\" to indicate automatic registation to `Vertex AI Model Registry`.\n",
+        "- `vertex_ai_model_id`: The human readable display name for the registered model.\n",
+        "- `vertex_ai_model_version_aliases`: Alternate names for the model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "57db464f4c42"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_NAME = \"penguins\"\n",
+        "MODEL_QUERY = f\"\"\"\n",
+        "CREATE OR REPLACE MODEL `{BQ_DATASET_NAME}.{MODEL_NAME}`\n",
+        "OPTIONS(\n",
+        "    model_type='DNN_CLASSIFIER',\n",
+        "    labels = ['species'],\n",
+        "    model_registry=\"vertex_ai\",\n",
+        "    vertex_ai_model_id=\"bqml_model_{TIMESTAMP}\", \n",
+        "    vertex_ai_model_version_aliases=[\"1\"]\n",
+        "    )\n",
+        "AS\n",
+        "SELECT *\n",
+        "FROM `{BQ_TABLE}`\n",
+        "\"\"\"\n",
+        "\n",
+        "job = bqclient.query(MODEL_QUERY)\n",
+        "print(job.errors, job.state)\n",
+        "\n",
+        "while job.running():\n",
+        "    from time import sleep\n",
+        "\n",
+        "    sleep(30)\n",
+        "    print(\"Running ...\")\n",
+        "print(job.errors, job.state)\n",
+        "\n",
+        "tblname = job.ddl_target_table\n",
+        "tblname = \"{}.{}\".format(tblname.dataset_id, tblname.table_id)\n",
+        "print(\"{} created in {}\".format(tblname, job.ended - job.started))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5b4970272040"
+      },
+      "source": [
+        "### Find the model in the `Vertex Model Registry`\n",
+        "\n",
+        "Finally, you can use the `Vertex AI Model` list() method with a filter query to find the automatically registered model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "76c22674ba99"
+      },
+      "outputs": [],
+      "source": [
+        "models = aiplatform.Model.list(filter=\"display_name=bqml_model_\" + TIMESTAMP)\n",
+        "model = models[0]\n",
+        "\n",
+        "print(model.gca_resource)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "48e6ef5d5ffa"
+      },
+      "outputs": [],
+      "source": [
+        "models = aiplatform.Model.list()\n",
+        "for model in models:\n",
+        "    if model.gca_resource.display_name.startswith(\"bqml\"):\n",
+        "        print(model.gca_resource.display_name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ef61354b1a5f"
+      },
+      "source": [
+        "### Delete the `BigQuery ML` model\n",
+        "\n",
+        "Next, delete the `BigQuery ML` instance of the model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f6004d1ce59d"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_QUERY = f\"\"\"\n",
+        "DROP MODEL `{BQ_DATASET_NAME}.{MODEL_NAME}`\n",
+        "\"\"\"\n",
+        "\n",
+        "job = bqclient.query(MODEL_QUERY)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -987,17 +1369,9 @@
         "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
         "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
         "\n",
-        "Otherwise, you can delete the individual resources you created in this tutorial:\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial.\n",
         "\n",
-        "- Dataset\n",
-        "- Pipeline\n",
-        "- Model\n",
-        "- Endpoint\n",
-        "- AutoML Training Job\n",
-        "- Batch Job\n",
-        "- Custom Job\n",
-        "- Hyperparameter Tuning Job\n",
-        "- Cloud Storage Bucket"
+        "Set `delete_storage` to `True` to delete the Cloud Storage bucket used in this notebook."
       ]
     },
     {
@@ -1008,61 +1382,23 @@
       },
       "outputs": [],
       "source": [
-        "delete_all = True\n",
+        "# Delete the endpoint using the Vertex endpoint object\n",
+        "endpoint.undeploy_all()\n",
+        "endpoint.delete()\n",
         "\n",
-        "if delete_all:\n",
-        "    # Delete the dataset using the Vertex dataset object\n",
-        "    try:\n",
-        "        if \"dataset\" in globals():\n",
-        "            dataset.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the model using the Vertex model object\n",
+        "try:\n",
+        "    model.delete()\n",
+        "except Exception as e:\n",
+        "    print(e)\n",
         "\n",
-        "    # Delete the model using the Vertex model object\n",
-        "    try:\n",
-        "        if \"model\" in globals():\n",
-        "            model.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
+        "# Delete the created BigQuery dataset\n",
+        "! bq rm -r -f $PROJECT_ID:$BQ_DATASET_NAME\n",
         "\n",
-        "    # Delete the endpoint using the Vertex endpoint object\n",
-        "    try:\n",
-        "        if \"endpoint\" in globals():\n",
-        "            endpoint.undeploy_all()\n",
-        "            endpoint.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the AutoML or Pipeline training job\n",
-        "    try:\n",
-        "        if \"dag\" in globals():\n",
-        "            dag.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the custom training job\n",
-        "    try:\n",
-        "        if \"job\" in globals():\n",
-        "            job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the batch prediction job using the Vertex batch prediction object\n",
-        "    try:\n",
-        "        if \"batch_predict_job\" in globals():\n",
-        "            batch_predict_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    # Delete the hyperparameter tuning job using the Vertex hyperparameter tuning object\n",
-        "    try:\n",
-        "        if \"hpt_job\" in globals():\n",
-        "            hpt_job.delete()\n",
-        "    except Exception as e:\n",
-        "        print(e)\n",
-        "\n",
-        "    if \"BUCKET_NAME\" in globals():\n",
-        "        ! gsutil rm -r $BUCKET_NAME"
+        "delete_storage = False\n",
+        "if delete_storage or os.getenv(\"IS_TESTING\"):\n",
+        "    # Delete the created GCS bucket\n",
+        "    ! gsutil rm -r $BUCKET_URI"
       ]
     }
   ],

notebooks/community/ml_ops/stage2/get_started_vertex_distributed_training.ipynb

@@ -8,7 +8,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2021 Google LLC\n",
+        "# Copyright 2022 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -29,9 +29,14 @@
         "id": "title:generic,gcp"
       },
       "source": [
-        "# E2E ML on GCP: MLOps stage 2 : experimentation: get started with Vertex Distributed Training\n",
+        "# E2E ML on GCP: MLOps stage 2 : experimentation: get started with Vertex AI Distributed Training\n",
         "\n",
         "<table align=\"left\">\n",
+        "      <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage2/get_started_vertex_distributed_training.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage2/get_started_vertex_distributed_training.ipynb\">\n",
         "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
@@ -39,8 +44,9 @@
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/ai/platform/notebooks/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage2/get_started_vertex_distributed_training.ipynb\">\n",
-        "      Open in Google Cloud Notebooks\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage2/get_started_vertex_distributed_training.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>\n",
         "</table>\n",
@@ -56,18 +62,7 @@
         "## Overview\n",
         "\n",
         "\n",
-        "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 2 : experimentation: get started with Vertex Distributed Training."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dataset:custom,boston,lrg"
-      },
-      "source": [
-        "### Dataset\n",
-        "\n",
-        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you will use in this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
+        "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 2 : experimentation: get started with Vertex Distributed Training. Please note: There are incompatibilities between Colab and Docker and the Docker section may not work until resolved by the platform."
       ]
     },
     {
@@ -126,59 +121,86 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "install_mlops"
+        "id": "dataset:custom,boston,lrg"
       },
       "source": [
-        "## Installations\n",
+        "### Dataset\n",
         "\n",
-        "Install *one time* the packages for executing the MLOps notebooks."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "install_mlops"
-      },
-      "outputs": [],
-      "source": [
-        "ONCE_ONLY = False\n",
-        "if ONCE_ONLY:\n",
-        "    ! pip3 install -U tensorflow==2.5 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-data-validation==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-transform==1.2 $USER_FLAG\n",
-        "    ! pip3 install -U tensorflow-io==0.18 $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-bigquery $USER_FLAG\n",
-        "    ! pip3 install --upgrade google-cloud-logging $USER_FLAG\n",
-        "    ! pip3 install --upgrade apache-beam[gcp] $USER_FLAG\n",
-        "    ! pip3 install --upgrade pyarrow $USER_FLAG\n",
-        "    ! pip3 install --upgrade cloudml-hypertune $USER_FLAG\n",
-        "    ! pip3 install --upgrade kfp $USER_FLAG\n",
-        "    ! pip3 install --upgrade torchvision $USER_FLAG\n",
-        "    ! pip3 install --upgrade rpy2 $USER_FLAG"
+        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you use in this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "restart"
+        "id": "d10166df7141"
       },
       "source": [
-        "### Restart the kernel\n",
+        "### Costs\n",
+        " \n",
+        "This tutorial uses billable components of Google Cloud:\n",
         "\n",
-        "Once you've installed the additional packages, you need to restart the notebook kernel so it can find the packages."
+        "Vertex AI\n",
+        "Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/),\n",
+        "        to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XkYpRvOQyVYb"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the packages required for executing this notebook."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "restart"
+        "id": "xs_Kt8RcyXTC"
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\") and not os.getenv(\"VIRTUAL_ENV\")\n",
+        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
+        "    \"/opt/deeplearning/metadata/env_version\"\n",
+        ")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
+        "! pip3 install {USER_FLAG} --upgrade google-cloud-aiplatform -q"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "oQhwq1iozAxh"
+      },
+      "source": [
+        "### Restart the kernel\n",
+        "\n",
+        "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "zo3YFZXLzCRJ"
+      },
+      "outputs": [],
+      "source": [
+        "# Automatically restart kernel after installs\n",
         "import os\n",
         "\n",
         "if not os.getenv(\"IS_TESTING\"):\n",
@@ -189,6 +211,32 @@
         "    app.kernel.do_shutdown(True)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "84cd83853240"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI, BigQuery, Compute Engine and Cloud Storage APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,bigquery,compute_component,storage_component).\n",
+        "\n",
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -261,11 +309,14 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "region"
+        "id": "qohAA9fJulvP"
       },
       "outputs": [],
       "source": [
-        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if REGION == \"[your-region]\":\n",
+        "    REGION = \"us-central1\""
       ]
     },
     {
@@ -283,7 +334,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "timestamp"
+        "id": "8NKwwe7aulvQ"
       },
       "outputs": [],
       "source": [
@@ -292,6 +343,82 @@
         "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "poKeKYG8ulvQ"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "**If you are using Vertex AI Workbench Notebooks**, your environment is already\n",
+        "authenticated. Skip this step."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "MIpJGzF9ulvQ"
+      },
+      "source": [
+        "**If you are using Colab**, run the cell below and follow the instructions\n",
+        "when prompted to authenticate your account via oAuth.\n",
+        "\n",
+        "**Otherwise**, follow these steps:\n",
+        "\n",
+        "1. In the Cloud Console, go to the [**Create service account key**\n",
+        "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
+        "\n",
+        "2. Click **Create service account**.\n",
+        "\n",
+        "3. In the **Service account name** field, enter a name, and\n",
+        "   click **Create**.\n",
+        "\n",
+        "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"Vertex AI\"\n",
+        "into the filter box, and select\n",
+        "   **Vertex AI Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
+        "\n",
+        "5. Click *Create*. A JSON file that contains your key downloads to your\n",
+        "local environment.\n",
+        "\n",
+        "6. Enter the path to your service account key as the\n",
+        "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Vh6KDXB5ulvQ"
+      },
+      "outputs": [],
+      "source": [
+        "# If you are running this notebook in Colab, run this cell and follow the\n",
+        "# instructions to authenticate your GCP account. This provides access to your\n",
+        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
+        "# requests.\n",
+        "\n",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "# If on Vertex AI Workbench, then don't execute this code\n",
+        "IS_COLAB = False\n",
+        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
+        "    \"DL_ANACONDA_HOME\"\n",
+        "):\n",
+        "    if \"google.colab\" in sys.modules:\n",
+        "        IS_COLAB = True\n",
+        "        from google.colab import auth as google_auth\n",
+        "\n",
+        "        google_auth.authenticate_user()\n",
+        "\n",
+        "    # If you are running this notebook locally, replace the string below with the\n",
+        "    # path to your service account key and run this cell to authenticate your GCP\n",
+        "    # account.\n",
+        "    elif not os.getenv(\"IS_TESTING\"):\n",
+        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -315,7 +442,8 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
+        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
+        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
       ]
     },
     {
@@ -326,8 +454,9 @@
       },
       "outputs": [],
       "source": [
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
+        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
+        "    BUCKET_NAME = PROJECT_ID + \"aip-\" + TIMESTAMP\n",
+        "    BUCKET_URI = f\"gs://{BUCKET_NAME}\""
       ]
     },
     {
@@ -343,11 +472,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "create_bucket"
+        "id": "Moosy2rOulvR"
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
+        "! gsutil mb -l $REGION $BUCKET_URI"
       ]
     },
     {
@@ -363,11 +492,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "validate_bucket"
+        "id": "56irx2CvulvS"
       },
       "outputs": [],
       "source": [
-        "! gsutil ls -al $BUCKET_NAME"
+        "! gsutil ls -al $BUCKET_URI"
       ]
     },
     {
@@ -408,11 +537,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "init_aip:mbsdk"
+        "id": "wbvYPSTDulvS"
       },
       "outputs": [],
       "source": [
-        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
+        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
       ]
     },
     {
@@ -441,7 +570,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "accelerators:training,prediction,ngpu,mbsdk"
+        "id": "PryARdnoulvT"
       },
       "outputs": [],
       "source": [
@@ -483,14 +612,14 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "container:training,prediction"
+        "id": "LhhUFw2nulvT"
       },
       "outputs": [],
       "source": [
         "if os.getenv(\"IS_TESTING_TF\"):\n",
         "    TF = os.getenv(\"IS_TESTING_TF\")\n",
         "else:\n",
-        "    TF = \"2.1\".replace(\".\", \"-\")\n",
+        "    TF = \"2.5\".replace(\".\", \"-\")\n",
         "\n",
         "if TF[0] == \"2\":\n",
         "    if TRAIN_GPU:\n",
@@ -551,7 +680,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "machine:training"
+        "id": "vytMaukeulvT"
       },
       "outputs": [],
       "source": [
@@ -613,7 +742,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "create_custom_pp_training_job:mbsdk"
+        "id": "mhw34XoOulvU"
       },
       "outputs": [],
       "source": [
@@ -621,7 +750,7 @@
         "\n",
         "job = aip.CustomPythonPackageTrainingJob(\n",
         "    display_name=DISPLAY_NAME,\n",
-        "    python_package_gcs_uri=f\"{BUCKET_NAME}/trainer_boston.tar.gz\",\n",
+        "    python_package_gcs_uri=f\"{BUCKET_URI}/trainer_boston.tar.gz\",\n",
         "    python_module_name=\"trainer.task\",\n",
         "    container_uri=TRAIN_IMAGE,\n",
         "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -662,7 +791,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "examine_training_package"
+        "id": "IAaZpZyyulvU"
       },
       "outputs": [],
       "source": [
@@ -711,7 +840,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "taskpy_contents:mirrored,boston"
+        "id": "zKzddzl6ulvV"
       },
       "outputs": [],
       "source": [
@@ -767,6 +896,13 @@
         "    strategy = tf.distribute.MultiWorkerMirroredStrategy()\n",
         "    logging.info(\"Multi-worker Strategy distributed training\")\n",
         "    logging.info('TF_CONFIG = {}'.format(os.environ.get('TF_CONFIG', 'Not found')))\n",
+        "    # Single Machine, multiple TPU devices\n",
+        "elif args.distribute == 'tpu':\n",
+        "    cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=\"local\")\n",
+        "    tf.config.experimental_connect_to_cluster(cluster_resolver)\n",
+        "    tf.tpu.experimental.initialize_tpu_system(cluster_resolver)\n",
+        "    strategy = tf.distribute.TPUStrategy(cluster_resolver)\n",
+        "    print(\"All devices: \", tf.config.list_logical_devices('TPU'))\n",
         "\n",
         "logging.info('num_replicas_in_sync = {}'.format(strategy.num_replicas_in_sync))\n",
         "\n",
@@ -825,8 +961,11 @@
         "    else:\n",
         "        task_type, task_id = None, None\n",
         "\n",
+        "    if args.distribute==\"tpu\":\n",
+        "        save_locally = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')\n",
+        "        model.save(args.model_dir, options=save_locally)\n",
         "    # single, mirrored or primary for multiworker\n",
-        "    if _is_chief(task_type, task_id):\n",
+        "    elif _is_chief(task_type, task_id):\n",
         "        model.save(args.model_dir)\n",
         "    # non-primary workers for multi-workers\n",
         "    else:\n",
@@ -860,14 +999,14 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "tarball_training_script"
+        "id": "LFUHioqTulvV"
       },
       "outputs": [],
       "source": [
         "! rm -f custom.tar custom.tar.gz\n",
         "! tar cvf custom.tar custom\n",
         "! gzip custom.tar\n",
-        "! gsutil cp custom.tar.gz $BUCKET_NAME/trainer_boston.tar.gz"
+        "! gsutil cp custom.tar.gz $BUCKET_URI/trainer_boston.tar.gz"
       ]
     },
     {
@@ -885,11 +1024,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "run_custom_pp_training_job:mirrored"
+        "id": "LnUX0UkvulvV"
       },
       "outputs": [],
       "source": [
-        "MODEL_DIR = BUCKET_NAME\n",
+        "MODEL_DIR = BUCKET_URI\n",
         "\n",
         "CMDARGS = [\"--epochs=5\", \"--batch_size=16\", \"--distribute=mirrored\"]\n",
         "\n",
@@ -920,7 +1059,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "delete_job"
+        "id": "iUWHFpPoulvW"
       },
       "outputs": [],
       "source": [
@@ -942,7 +1081,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "model_delete:mbsdk"
+        "id": "-0gqCUTEulvW"
       },
       "outputs": [],
       "source": [
@@ -1027,7 +1166,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "create_custom_pp_training_job:mbsdk"
+        "id": "aXvPN8P6ulvX"
       },
       "source": [
         "### Create and run custom training job\n",
@@ -1053,7 +1192,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "create_custom_pp_training_job:mbsdk"
+        "id": "kYcFsVSEulvX"
       },
       "outputs": [],
       "source": [
@@ -1061,7 +1200,7 @@
         "\n",
         "job = aip.CustomPythonPackageTrainingJob(\n",
         "    display_name=DISPLAY_NAME,\n",
-        "    python_package_gcs_uri=f\"{BUCKET_NAME}/trainer_boston.tar.gz\",\n",
+        "    python_package_gcs_uri=f\"{BUCKET_URI}/trainer_boston.tar.gz\",\n",
         "    python_module_name=\"trainer.task\",\n",
         "    container_uri=TRAIN_IMAGE,\n",
         "    model_serving_container_image_uri=DEPLOY_IMAGE,\n",
@@ -1084,11 +1223,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "run_custom_pp_training_job:multiworker"
+        "id": "GHRxPU32ulvX"
       },
       "outputs": [],
       "source": [
-        "MODEL_DIR = BUCKET_NAME\n",
+        "MODEL_DIR = BUCKET_URI\n",
         "\n",
         "CMDARGS = [\"--epochs=5\", \"--batch_size=16\", \"--distribute=multiworker\"]\n",
         "\n",
@@ -1111,7 +1250,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "delete_job"
+        "id": "92D_hbuVulvX"
       },
       "source": [
         "### Delete a custom training job\n",
@@ -1123,7 +1262,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "delete_job"
+        "id": "CqrfWkB3ulvX"
       },
       "outputs": [],
       "source": [
@@ -1175,14 +1314,13 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "write_docker_file:training,multiworker"
+        "id": "pGI2viDAulvY"
       },
       "outputs": [],
       "source": [
         "%%writefile custom/Dockerfile\n",
         "\n",
         "FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-5\n",
-        "WORKDIR /root\n",
         "\n",
         "WORKDIR /\n",
         "\n",
@@ -1208,7 +1346,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "name_container:training"
+        "id": "7P8cdlFtulvY"
       },
       "outputs": [],
       "source": [
@@ -1228,11 +1366,15 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "build_container:training"
+        "id": "jmw5cakNulvY"
       },
       "outputs": [],
       "source": [
-        "! docker build custom -t $TRAIN_IMAGE"
+        "if not IS_COLAB:\n",
+        "    ! docker build custom -t $TRAIN_IMAGE\n",
+        "else:\n",
+        "    # install docker daemon\n",
+        "    ! apt-get -qq install docker.io"
       ]
     },
     {
@@ -1250,11 +1392,12 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "test_container:training"
+        "id": "jJGLjU-TulvZ"
       },
       "outputs": [],
       "source": [
-        "! docker run $TRAIN_IMAGE --epochs=5 --model-dir=./"
+        "if not IS_COLAB:\n",
+        "    ! docker run $TRAIN_IMAGE --epochs=5 --model-dir=./"
       ]
     },
     {
@@ -1272,11 +1415,42 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "register_container:training"
+        "id": "GAXGjae7ulvZ"
       },
       "outputs": [],
       "source": [
-        "! docker push $TRAIN_IMAGE"
+        "if not IS_COLAB:\n",
+        "    ! docker push $TRAIN_IMAGE"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f50e9c553fb7"
+      },
+      "source": [
+        "*Executes in Colab*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a7e8c98f1e56"
+      },
+      "outputs": [],
+      "source": [
+        "%%bash -s $IS_COLAB $TRAIN_IMAGE\n",
+        "if [ $1 == \"False\" ]; then\n",
+        "  exit 0\n",
+        "fi\n",
+        "set -x\n",
+        "dockerd -b none --iptables=0 -l warn &\n",
+        "for i in $(seq 5); do [ ! -S \"/var/run/docker.sock\" ] && sleep 2 || break; done\n",
+        "docker build custom -t $2\n",
+        "docker run $2 --epochs=5 --model-dir=./\n",
+        "docker push $2\n",
+        "kill $(jobs -p)"
       ]
     },
     {
@@ -1296,13 +1470,13 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "worker_pool_primary"
+        "id": "CEAnXBzCulvZ"
       },
       "outputs": [],
       "source": [
         "PRIMARY_COMPUTE = \"n2-highcpu-64\"\n",
         "\n",
-        "MODEL_DIR = BUCKET_NAME\n",
+        "MODEL_DIR = BUCKET_URI\n",
         "\n",
         "CMDARGS = [\n",
         "    \"--model-dir=\" + MODEL_DIR,\n",
@@ -1339,7 +1513,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "worker_pool_training"
+        "id": "6dchPSfNulvZ"
       },
       "outputs": [],
       "source": [
@@ -1375,7 +1549,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "custom_job:worker_pool"
+        "id": "m2VgmqEOulva"
       },
       "outputs": [],
       "source": [
@@ -1399,7 +1573,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "run_custom_job:multiworker"
+        "id": "hg8vnI_Wulva"
       },
       "outputs": [],
       "source": [
@@ -1413,7 +1587,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "delete_job"
+        "id": "WT76Sc-culva"
       },
       "source": [
         "### Delete a custom training job\n",
@@ -1425,7 +1599,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "delete_job"
+        "id": "I_IxVfuDulva"
       },
       "outputs": [],
       "source": [
@@ -1474,7 +1648,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "custom_job:worker_pool"
+        "id": "L8Av8ATVulvb"
       },
       "source": [
         "### Create CustomJob with worker pool specifications\n",
@@ -1490,7 +1664,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "custom_job:worker_pool"
+        "id": "TUWEP1Lmulvb"
       },
       "outputs": [],
       "source": [
@@ -1502,7 +1676,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "run_custom_job:multiworker"
+        "id": "_95FH8jeulvb"
       },
       "source": [
         "### Run the CustomJob\n",
@@ -1514,7 +1688,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "run_custom_job:multiworker"
+        "id": "IEbrY05Gulvb"
       },
       "outputs": [],
       "source": [
@@ -1528,7 +1702,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "delete_job"
+        "id": "8R2Bnmwmulvb"
       },
       "source": [
         "### Delete a custom training job\n",
@@ -1540,7 +1714,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "delete_job"
+        "id": "s1geVE3Lulvb"
       },
       "outputs": [],
       "source": [
@@ -1583,14 +1757,14 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "docker_write:tpu"
+        "id": "nQVPtknpulvb"
       },
       "outputs": [],
       "source": [
         "%%writefile custom/Dockerfile\n",
         "FROM python:3.8\n",
         "\n",
-        "WORKDIR /root\n",
+        "WORKDIR /\n",
         "\n",
         "# Copies the trainer code to the docker image.\n",
         "COPY trainer /trainer\n",
@@ -1622,11 +1796,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "docker_push:tpu"
+        "id": "J_d_zEXUulvc"
       },
       "outputs": [],
       "source": [
-        "TRAIN_IMAGE = f\"gcr.io/\" + PROJECT_ID + \"/tpu-train:latest\"\n",
+        "TRAIN_IMAGE = \"gcr.io/\" + PROJECT_ID + \"/tpu-train:latest\"\n",
         "\n",
         "os.chdir(\"custom\")\n",
         "! docker build --quiet --tag={TRAIN_IMAGE} .\n",
@@ -1653,7 +1827,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "worker_pool_tpu"
+        "id": "d514eU7lulvc"
       },
       "outputs": [],
       "source": [
@@ -1701,7 +1875,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "custom_job:worker_pool"
+        "id": "RruSqNfrulvc"
       },
       "source": [
         "### Create CustomJob with worker pool specifications\n",
@@ -1717,7 +1891,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "custom_job:worker_pool"
+        "id": "2QvSqbbHulvc"
       },
       "outputs": [],
       "source": [
@@ -1729,7 +1903,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "run_custom_job:multiworker"
+        "id": "Iw4L3UIfulvd"
       },
       "source": [
         "### Run the CustomJob\n",
@@ -1741,7 +1915,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "run_custom_job:multiworker"
+        "id": "zmqCNS78ulvd"
       },
       "outputs": [],
       "source": [
@@ -1755,7 +1929,7 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "delete_job"
+        "id": "gWZoH9QKulvd"
       },
       "source": [
         "### Delete a custom training job\n",
@@ -1767,7 +1941,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "delete_job"
+        "id": "Lt8BJ4iBulvd"
       },
       "outputs": [],
       "source": [
@@ -1787,13 +1961,7 @@
         "\n",
         "Otherwise, you can delete the individual resources you created in this tutorial:\n",
         "\n",
-        "- Dataset\n",
-        "- Pipeline\n",
-        "- Model\n",
-        "- Endpoint\n",
-        "- Batch Job\n",
-        "- Custom Job\n",
-        "- Hyperparameter Tuning Job\n",
+        "\n",
         "- Cloud Storage Bucket"
       ]
     },
@@ -1801,70 +1969,15 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "cleanup"
+        "id": "U98Wzc01ulvd"
       },
       "outputs": [],
       "source": [
-        "delete_dataset = True\n",
-        "delete_pipeline = True\n",
-        "delete_model = True\n",
-        "delete_endpoint = True\n",
-        "delete_batchjob = True\n",
-        "delete_customjob = True\n",
-        "delete_hptjob = True\n",
-        "delete_bucket = True\n",
+        "# Set this to true only if you'd like to delete your bucket\n",
+        "delete_bucket = False\n",
         "\n",
-        "# Delete the dataset using the Vertex fully qualified identifier for the dataset\n",
-        "try:\n",
-        "    if delete_dataset and \"dataset_id\" in globals():\n",
-        "        clients[\"dataset\"].delete_dataset(name=dataset_id)\n",
-        "except Exception as e:\n",
-        "    print(e)\n",
-        "\n",
-        "# Delete the training pipeline using the Vertex fully qualified identifier for the pipeline\n",
-        "try:\n",
-        "    if delete_pipeline and \"pipeline_id\" in globals():\n",
-        "        clients[\"pipeline\"].delete_training_pipeline(name=pipeline_id)\n",
-        "except Exception as e:\n",
-        "    print(e)\n",
-        "\n",
-        "# Delete the model using the Vertex fully qualified identifier for the model\n",
-        "try:\n",
-        "    if delete_model and \"model_to_deploy_id\" in globals():\n",
-        "        clients[\"model\"].delete_model(name=model_to_deploy_id)\n",
-        "except Exception as e:\n",
-        "    print(e)\n",
-        "\n",
-        "# Delete the endpoint using the Vertex fully qualified identifier for the endpoint\n",
-        "try:\n",
-        "    if delete_endpoint and \"endpoint_id\" in globals():\n",
-        "        clients[\"endpoint\"].delete_endpoint(name=endpoint_id)\n",
-        "except Exception as e:\n",
-        "    print(e)\n",
-        "\n",
-        "# Delete the batch job using the Vertex fully qualified identifier for the batch job\n",
-        "try:\n",
-        "    if delete_batchjob and \"batch_job_id\" in globals():\n",
-        "        clients[\"job\"].delete_batch_prediction_job(name=batch_job_id)\n",
-        "except Exception as e:\n",
-        "    print(e)\n",
-        "\n",
-        "# Delete the custom job using the Vertex fully qualified identifier for the custom job\n",
-        "try:\n",
-        "    if delete_customjob and \"job_id\" in globals():\n",
-        "        clients[\"job\"].delete_custom_job(name=job_id)\n",
-        "except Exception as e:\n",
-        "    print(e)\n",
-        "\n",
-        "# Delete the hyperparameter tuning job using the Vertex fully qualified identifier for the hyperparameter tuning job\n",
-        "try:\n",
-        "    if delete_hptjob and \"hpt_job_id\" in globals():\n",
-        "        clients[\"job\"].delete_hyperparameter_tuning_job(name=hpt_job_id)\n",
-        "except Exception as e:\n",
-        "    print(e)\n",
-        "\n",
-        "if delete_bucket and \"BUCKET_NAME\" in globals():\n",
-        "    ! gsutil rm -r $BUCKET_NAME"
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil rm -r $BUCKET_URI"
       ]
     }
   ],