feat: python notebooks template repo

synthtool/gcp/common.py

@@ -205,6 +205,12 @@ def py_samples_override(
         overridden_samples_kwargs["subdir"] = override_path
         return self._generic_library("python_samples", **overridden_samples_kwargs)
 
+    def python_notebooks(self, **kwargs) -> Path:
+        # kwargs["metadata"] is required to load values from .repo-metadata.json
+        if "metadata" not in kwargs:
+            kwargs["metadata"] = {}
+        return self._generic_library("python_notebooks", **kwargs)
+
     def py_library(self, **kwargs) -> Path:
         # kwargs["metadata"] is required to load values from .repo-metadata.json
         if "metadata" not in kwargs:

synthtool/gcp/templates/python_notebooks/.cloud-build/check_python_version.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+# # Copyright 2021 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
+#
+#      http://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 sys
+
+MINIMUM_MAJOR_VERSION = 3
+MINIMUM_MINOR_VERSION = 3
+
+if (
+    sys.version_info.major < MINIMUM_MAJOR_VERSION
+    and sys.version_info.minor < MINIMUM_MINOR_VERSION
+):
+    print("Error: Python version less than 3.5")
+    exit(1)
+else:
+    print(f"Python version acceptable: {sys.version}")
+    exit(0)

synthtool/gcp/templates/python_notebooks/.cloud-build/cleanup/cleanup-cloudbuild.yaml

@@ -0,0 +1,8 @@
+steps:  
+  # Install Python dependencies and run cleanup script
+  - name: ${_PYTHON_IMAGE}
+    entrypoint: /bin/sh
+    args:
+    - -c
+    - 'python3 -m pip install -U -r .cloud-build/cleanup/cleanup-requirements.txt && python3 .cloud-build/cleanup/cleanup.py'
+timeout: 86400s

synthtool/gcp/templates/python_notebooks/.cloud-build/cleanup/cleanup-requirements.txt

@@ -0,0 +1 @@
+google-cloud-aiplatform

synthtool/gcp/templates/python_notebooks/.cloud-build/cleanup/cleanup.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+# # Copyright 2021 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
+#
+#      http://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.
+from typing import List
+from resource_cleanup_manager import (
+    ResourceCleanupManager,
+    DatasetResourceCleanupManager,
+    EndpointResourceCleanupManager,
+    ModelResourceCleanupManager,
+)
+
+
+def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: bool):
+    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
+
+            if is_dry_run:
+                resource_name = manager.resource_name(resource)
+                print(f"Will delete '{type_name}': {resource_name}")
+            else:
+                manager.delete(resource)
+
+        print("")
+
+
+is_dry_run = False
+
+if is_dry_run:
+    print("Starting cleanup in dry run mode...")
+
+# List of all cleanup managers
+managers = [
+    DatasetResourceCleanupManager(),
+    EndpointResourceCleanupManager(),
+    ModelResourceCleanupManager(),
+]
+
+run_cleanup_managers(managers=managers, is_dry_run=is_dry_run)

synthtool/gcp/templates/python_notebooks/.cloud-build/cleanup/resource_cleanup_manager.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python
+# # Copyright 2021 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
+#
+#      http://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 abc
+from google.cloud import aiplatform
+from typing import Any
+from proto.datetime_helpers import DatetimeWithNanoseconds
+from google.cloud.aiplatform import base
+
+# If a resource was updated within this number of seconds, do not delete.
+RESOURCE_UPDATE_BUFFER_IN_SECONDS = 60 * 60 * 8
+
+
+class ResourceCleanupManager(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def type_name(str) -> str:
+        pass
+
+    @abc.abstractmethod
+    def list(self) -> Any:
+        pass
+
+    @abc.abstractmethod
+    def resource_name(self, resource: Any) -> str:
+        pass
+
+    @abc.abstractmethod
+    def delete(self, resource: Any):
+        pass
+
+    @abc.abstractmethod
+    def get_seconds_since_modification(self, resource: Any) -> float:
+        pass
+
+    def is_deletable(self, resource: Any) -> bool:
+        time_difference = self.get_seconds_since_modification(resource)
+
+        if self.resource_name(resource).startswith("perm"):
+            print(f"Skipping '{resource}' due to name starting with 'perm'.")
+            return False
+
+        # 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}'."
+            )
+            return False
+
+        return True
+
+
+class VertexAIResourceCleanupManager(ResourceCleanupManager):
+    @property
+    @abc.abstractmethod
+    def vertex_ai_resource(self) -> base.VertexAiResourceNounWithFutureManager:
+        pass
+
+    @property
+    def type_name(self) -> str:
+        return self.vertex_ai_resource._resource_noun
+
+    def list(self) -> Any:
+        return self.vertex_ai_resource.list()
+
+    def resource_name(self, resource: Any) -> str:
+        return resource.display_name
+
+    def delete(self, resource):
+        resource.delete()
+
+    def get_seconds_since_modification(self, resource: Any) -> bool:
+        update_time = resource.update_time
+        current_time = DatetimeWithNanoseconds.now(tz=update_time.tzinfo)
+        return (current_time - update_time).total_seconds()
+
+
+class DatasetResourceCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.datasets._Dataset
+
+
+class EndpointResourceCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.Endpoint
+
+    def delete(self, resource):
+        resource.delete(force=True)
+
+
+class ModelResourceCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.Model

synthtool/gcp/templates/python_notebooks/.cloud-build/execute_changed_notebooks.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python
+# # Copyright 2021 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
+#
+#      http://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 argparse
+import pathlib
+import subprocess
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import execute_notebook
+
+
+def run_changed_notebooks(
+    test_paths_file: str,
+    output_folder: str,
+    variable_project_id: str,
+    variable_region: str,
+    base_branch: Optional[str],
+):
+    """
+    Run the notebooks that exist under the folders defined in the test_paths_file.
+    It only runs notebooks that have differences from the Git base_branch.
+    The executed notebooks are saved in the output_folder.
+    Variables are also injected into the notebooks such as the variable_project_id and variable_region.
+    Args:
+        test_paths_file (str):
+            Required. The new-line delimited file to folders and files that need checking.
+            Folders are checked recursively.
+        base_branch (str):
+            Optional. If provided, only the files that have changed from the base_branch will be checked.
+            If not provided, all files will be checked.
+        output_folder (str):
+            Required. The folder to write executed notebooks to.
+        variable_project_id (str):
+            Required. The value for PROJECT_ID to inject into notebooks.
+        variable_region (str):
+            Required. The value for REGION to inject into notebooks.
+    """
+
+    test_paths = []
+    with open(test_paths_file) as file:
+        lines = [line.strip() for line in file.readlines()]
+        lines = [line for line in lines if len(line) > 0]
+        test_paths = [line for line in lines]
+
+    if len(test_paths) == 0:
+        raise RuntimeError("No test folders found.")
+
+    print(f"Checking folders: {test_paths}")
+
+    # Find notebooks
+    notebooks = []
+    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
+        )
+    else:
+        print(f"Looking for all notebooks.")
+        notebooks = subprocess.check_output(["git", "ls-files"] + test_paths)
+
+    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 Path(notebook).exists()]
+
+    # Create paths
+    artifacts_path = Path(output_folder)
+    artifacts_path.mkdir(parents=True, exist_ok=True)
+    artifacts_path.joinpath("success").mkdir(parents=True, exist_ok=True)
+    artifacts_path.joinpath("failure").mkdir(parents=True, exist_ok=True)
+
+    passed_notebooks: List[str] = []
+    failed_notebooks: List[str] = []
+
+    if len(notebooks) > 0:
+        print(f"Found {len(notebooks)} modified notebooks: {notebooks}")
+
+        for notebook in notebooks:
+            print(f"Running notebook: {notebook}")
+
+            # TODO: Handle cases where multiple notebooks have the same name
+            try:
+                execute_notebook.execute_notebook(
+                    notebook_file_path=notebook,
+                    output_file_folder=artifacts_path,
+                    replacement_map={
+                        "PROJECT_ID": variable_project_id,
+                        "REGION": variable_region,
+                    },
+                )
+                print(f"Notebook finished successfully.")
+                passed_notebooks.append(notebook)
+            except Exception as error:
+                print(f"Notebook finished with failure: {error}")
+                failed_notebooks.append(notebook)
+    else:
+        print("No notebooks modified in this pull request.")
+
+    if len(failed_notebooks) > 0:
+        print(f"{len(failed_notebooks)} notebooks failed:")
+        print(failed_notebooks)
+        print(f"{len(passed_notebooks)} notebooks passed:")
+        print(passed_notebooks)
+    elif len(passed_notebooks) > 0:
+        print("All notebooks executed successfully:")
+        print(passed_notebooks)
+
+
+parser = argparse.ArgumentParser(description="Run changed notebooks.")
+parser.add_argument(
+    "--test_paths_file",
+    type=pathlib.Path,
+    help="The path to the file that has newline-limited folders of notebooks that should be tested.",
+    required=True,
+)
+parser.add_argument(
+    "--base_branch",
+    help="The base git branch to diff against to find changed files.",
+    required=False,
+)
+parser.add_argument(
+    "--output_folder",
+    type=pathlib.Path,
+    help="The path to the folder to store executed notebooks.",
+    required=True,
+)
+parser.add_argument(
+    "--variable_project_id",
+    type=str,
+    help="The GCP project id. This is used to inject a variable value into the notebook before running.",
+    required=True,
+)
+parser.add_argument(
+    "--variable_region",
+    type=str,
+    help="The GCP region. This is used to inject a variable value into the notebook before running.",
+    required=True,
+)
+
+args = parser.parse_args()
+run_changed_notebooks(
+    test_paths_file=args.test_paths_file,
+    base_branch=args.base_branch,
+    output_folder=args.output_folder,
+    variable_project_id=args.variable_project_id,
+    variable_region=args.variable_region,
+)