Fixed missing return

.cloud-build/execute_changed_notebooks_cli.py

@@ -0,0 +1,100 @@
+#!/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.
+
+"""A CLI to process changed notebooks and execute them on Google Cloud Build"""
+
+import argparse
+import pathlib
+import execute_changed_notebooks_helper
+
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+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(
+    "--container_uri",
+    type=str,
+    help="The container uri to run each notebook in.",
+    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,
+)
+parser.add_argument(
+    "--staging_bucket",
+    type=str,
+    help="The GCP directory for staging temporary files.",
+    required=True,
+)
+parser.add_argument(
+    "--artifacts_bucket",
+    type=str,
+    help="The GCP directory for storing executed notebooks.",
+    required=True,
+)
+parser.add_argument(
+    "--should_parallelize",
+    type=str2bool,
+    nargs="?",
+    const=True,
+    default=True,
+    help="Should run notebooks in parallel.",
+)
+
+args = parser.parse_args()
+
+notebooks = execute_changed_notebooks_helper.get_changed_notebooks(
+    test_paths_file=args.test_paths_file,
+    base_branch=args.base_branch,
+)
+
+execute_changed_notebooks_helper.process_and_execute_notebooks(
+    notebooks=notebooks,
+    container_uri=args.container_uri,
+    staging_bucket=args.staging_bucket,
+    artifacts_bucket=args.artifacts_bucket,
+    variable_project_id=args.variable_project_id,
+    variable_region=args.variable_region,
+    should_parallelize=args.should_parallelize,
+)

.cloud-build/execute_changed_notebooks_helper.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import argparse
 import concurrent
 import dataclasses
 import datetime
@@ -32,17 +31,6 @@ from utils import util, NotebookProcessors
 from google.cloud.devtools.cloudbuild_v1.types import BuildOperationMetadata
 
 
-def str2bool(v):
-    if isinstance(v, bool):
-        return v
-    if v.lower() in ("yes", "true", "t", "y", "1"):
-        return True
-    elif v.lower() in ("no", "false", "f", "n", "0"):
-        return False
-    else:
-        raise argparse.ArgumentTypeError("Boolean value expected.")
-
-
 def format_timedelta(delta: datetime.timedelta) -> str:
     """Formats a timedelta duration to [N days] %H:%M:%S format"""
     seconds = int(delta.total_seconds())
@@ -115,7 +103,7 @@ def _create_tag(filepath: str) -> str:
     return tag
 
 
-def execute_notebook(
+def process_and_execute_notebook(
     container_uri: str,
     staging_bucket: str,
     artifacts_bucket: str,
@@ -202,41 +190,13 @@ def execute_notebook(
     return result
 
 
-def run_changed_notebooks(
+def get_changed_notebooks(
     test_paths_file: str,
-    container_uri: str,
-    staging_bucket: str,
-    artifacts_bucket: str,
-    variable_project_id: str,
-    variable_region: str,
-    should_parallelize: bool,
     base_branch: Optional[str] = None,
-):
+) -> List[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 artifacts_bucket.
-
-    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.
-        staging_bucket (str):
-            Required. The GCS staging bucket to write source code to.
-        artifacts_bucket (str):
-            Required. The GCS staging bucket 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.
-        should_parallelize (bool):
-            Required. Should run notebooks in parallel using a thread pool as opposed to in sequence.
+    Get the notebooks that exist under the folders defined in the test_paths_file.
+    It only returns notebooks that have differences from the Git base_branch.
     """
 
     test_paths = []
@@ -266,6 +226,44 @@ def run_changed_notebooks(
     notebooks = [notebook for notebook in notebooks if len(notebook) > 0]
     notebooks = [notebook for notebook in notebooks if pathlib.Path(notebook).exists()]
 
+    return notebooks
+
+
+def process_and_execute_notebooks(
+    notebooks: List[str],
+    container_uri: str,
+    staging_bucket: str,
+    artifacts_bucket: str,
+    variable_project_id: str,
+    variable_region: str,
+    should_parallelize: bool,
+):
+    """
+    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 artifacts_bucket.
+
+    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.
+        staging_bucket (str):
+            Required. The GCS staging bucket to write source code to.
+        artifacts_bucket (str):
+            Required. The GCS staging bucket 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.
+        should_parallelize (bool):
+            Required. Should run notebooks in parallel using a thread pool as opposed to in sequence.
+    """
     notebook_execution_results: List[NotebookExecutionResult] = []
 
     if len(notebooks) > 0:
@@ -279,7 +277,7 @@ def run_changed_notebooks(
                 notebook_execution_results = list(
                     executor.map(
                         functools.partial(
-                            execute_notebook,
+                            process_and_execute_notebook,
                             container_uri,
                             staging_bucket,
                             artifacts_bucket,
@@ -291,7 +289,7 @@ def run_changed_notebooks(
                 )
         else:
             notebook_execution_results = [
-                execute_notebook(
+                process_and_execute_notebook(
                     container_uri=container_uri,
                     staging_bucket=staging_bucket,
                     artifacts_bucket=artifacts_bucket,
@@ -340,68 +338,4 @@ def run_changed_notebooks(
 
     # 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")
-
-
-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(
-    "--container_uri",
-    type=str,
-    help="The container uri to run each notebook in.",
-    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,
-)
-parser.add_argument(
-    "--staging_bucket",
-    type=str,
-    help="The GCP directory for staging temporary files.",
-    required=True,
-)
-parser.add_argument(
-    "--artifacts_bucket",
-    type=str,
-    help="The GCP directory for storing executed notebooks.",
-    required=True,
-)
-parser.add_argument(
-    "--should_parallelize",
-    type=str2bool,
-    nargs="?",
-    const=True,
-    default=True,
-    help="Should run notebooks in parallel.",
-)
-
-args = parser.parse_args()
-run_changed_notebooks(
-    test_paths_file=args.test_paths_file,
-    container_uri=args.container_uri,
-    staging_bucket=args.staging_bucket,
-    artifacts_bucket=args.artifacts_bucket,
-    variable_project_id=args.variable_project_id,
-    variable_region=args.variable_region,
-    should_parallelize=args.should_parallelize,
-    base_branch=args.base_branch,
-)
+        raise RuntimeError("Notebook failures detected. See logs for details")

.cloud-build/execute_notebook_cli.py

@@ -13,10 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import argparse
-import ExecuteNotebook
+"""A CLI to download (optional) and run a single notebook locally"""
 
-parser = argparse.ArgumentParser(description="Run changed notebooks.")
+import argparse
+import execute_notebook_helper
+
+parser = argparse.ArgumentParser(description="Run a single notebook locally.")
 parser.add_argument(
     "--notebook_source",
     type=str,
@@ -31,7 +33,7 @@ parser.add_argument(
 )
 
 args = parser.parse_args()
-ExecuteNotebook.execute_notebook(
+execute_notebook_helper.execute_notebook(
     notebook_source=args.notebook_source,
     output_file_or_uri=args.output_file_or_uri,
     should_log_output=True,

.cloud-build/execute_notebook_helper.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Methods to run a notebook locally"""
+
 import sys
 import os
 import errno
@@ -30,6 +32,7 @@ def execute_notebook(
     output_file_or_uri: str,
     should_log_output: bool,
 ):
+    """Execute a single notebook using Papermill"""
     file_name = os.path.basename(os.path.normpath(notebook_source))
 
     # Download notebook if it's a GCS URI

.cloud-build/execute_notebook_remote.py

@@ -1,3 +1,20 @@
+#!/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.
+
+"""Methods to run a notebook on Google Cloud Build"""
+
 from google.protobuf import duration_pb2
 from yaml.loader import FullLoader
 
@@ -22,8 +39,7 @@ def execute_notebook_remote(
     container_uri: str,
     tag: Optional[str],
 ) -> operation.Operation:
-    """Create and execute a simple Google Cloud Build configuration,
-    print the in-progress status and print the completed status."""
+    """Create and execute a single notebook on Google Cloud Build"""
 
     # Authorize the client with Google defaults
     credentials, project_id = google.auth.default()

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

@@ -32,7 +32,7 @@ steps:
     entrypoint: /bin/sh
     args:
     - -c
-    - 'python3 -m pip install -U pip && python3 -m pip freeze && python3 .cloud-build/ExecuteChangedNotebooks.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}'
+    - '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}'
     env:
     - 'IS_TESTING=1'
 timeout: 86400s

.cloud-build/utils/util.py

@@ -11,7 +11,7 @@ import uuid
 
 
 def download_file(bucket_name: str, blob_name: str, destination_file: str) -> str:
-    """Copies a remote GCS file to a local path."""
+    """Copies a remote GCS file to a local path"""
     remote_file_path = "".join(["gs://", "/".join([bucket_name, blob_name])])
 
     subprocess.check_output(
@@ -25,7 +25,7 @@ def upload_file(
     local_file_path: str,
     remote_file_path: str,
 ) -> str:
-    """Copies a local file to a GCS path."""
+    """Copies a local file to a GCS path"""
     subprocess.check_output(
         ["gsutil", "cp", local_file_path, remote_file_path], encoding="UTF-8"
     )