From 4fca7d98b2a239988a832ac674c2de2911d09f4a Mon Sep 17 00:00:00 2001 From: Ivan Cheung Date: Wed, 18 May 2022 18:58:54 -0400 Subject: [PATCH] Increases notebook concurrency and linted CI files (#512) * Fixed private pool issues * Ran linter * Added worker timeouts * Tweaked timeout * Removed gcloud requirement * Removed unneeded file --- .cloud-build/cleanup/cleanup.py | 53 +++++++++---------- .../cleanup/resource_cleanup_manager.py | 5 +- .cloud-build/execute_changed_notebooks_cli.py | 11 +++- .../execute_changed_notebooks_helper.py | 40 +++++++++++--- .cloud-build/execute_notebook_cli.py | 1 + .cloud-build/execute_notebook_helper.py | 8 +-- .cloud-build/execute_notebook_remote.py | 37 +++++++------ ...book-execution-test-cloudbuild-single.yaml | 5 +- .cloud-build/requirements.txt | 2 +- .cloud-build/utils/NotebookProcessors.py | 6 ++- .cloud-build/utils/UpdateNotebookVariables.py | 2 +- .cloud-build/utils/util.py | 14 ++--- 12 files changed, 110 insertions(+), 74 deletions(-) diff --git a/.cloud-build/cleanup/cleanup.py b/.cloud-build/cleanup/cleanup.py index b31acf607..1a88c747a 100644 --- a/.cloud-build/cleanup/cleanup.py +++ b/.cloud-build/cleanup/cleanup.py @@ -1,45 +1,44 @@ from typing import List -from resource_cleanup_manager import ( - ResourceCleanupManager, - DatasetResourceCleanupManager, - EndpointResourceCleanupManager, - ModelResourceCleanupManager, -) + +from resource_cleanup_manager import (DatasetResourceCleanupManager, + EndpointResourceCleanupManager, + ModelResourceCleanupManager, + ResourceCleanupManager) 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: + 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: + try: + 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(), ] run_cleanup_managers(managers=managers, is_dry_run=is_dry_run) diff --git a/.cloud-build/cleanup/resource_cleanup_manager.py b/.cloud-build/cleanup/resource_cleanup_manager.py index ae6eb804a..84c61a32c 100644 --- a/.cloud-build/cleanup/resource_cleanup_manager.py +++ b/.cloud-build/cleanup/resource_cleanup_manager.py @@ -1,8 +1,9 @@ import abc -from google.cloud import aiplatform from typing import Any -from proto.datetime_helpers import DatetimeWithNanoseconds + +from google.cloud import aiplatform 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 diff --git a/.cloud-build/execute_changed_notebooks_cli.py b/.cloud-build/execute_changed_notebooks_cli.py index 28c9ed772..c4fac41f7 100755 --- a/.cloud-build/execute_changed_notebooks_cli.py +++ b/.cloud-build/execute_changed_notebooks_cli.py @@ -17,6 +17,7 @@ import argparse import pathlib + import execute_changed_notebooks_helper @@ -73,6 +74,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, @@ -102,6 +110,7 @@ execute_changed_notebooks_helper.process_and_execute_notebooks( artifacts_bucket=args.artifacts_bucket, variable_project_id=args.variable_project_id, variable_region=args.variable_region, - private_pool_id=args.private_pool_id if not "default" else None, + private_pool_id=args.private_pool_id, should_parallelize=args.should_parallelize, + timeout=args.timeout, ) diff --git a/.cloud-build/execute_changed_notebooks_helper.py b/.cloud-build/execute_changed_notebooks_helper.py index 90604f892..fc95461ba 100755 --- a/.cloud-build/execute_changed_notebooks_helper.py +++ b/.cloud-build/execute_changed_notebooks_helper.py @@ -17,18 +17,22 @@ import concurrent import dataclasses import datetime import functools +import operator import os import pathlib -import nbformat import re import subprocess from typing import List, Optional -from tabulate import tabulate -import operator 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 def format_timedelta(delta: datetime.timedelta) -> str: @@ -103,6 +107,9 @@ 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, @@ -110,9 +117,12 @@ def process_and_execute_notebook( variable_project_id: str, variable_region: str, private_pool_id: Optional[str], + deadline: datetime, notebook: str, should_get_tail_logs: bool = False, ) -> NotebookExecutionResult: + rate_limit.wait() # wait before creating the task + print(f"Running notebook: {notebook}") # Create paths @@ -145,14 +155,20 @@ def process_and_execute_notebook( # 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, ) operation_metadata = BuildOperationMetadata(mapping=operation.metadata) @@ -241,6 +257,7 @@ def process_and_execute_notebooks( variable_region: str, private_pool_id: Optional[str], should_parallelize: bool, + timeout: int, ): """ Run the notebooks that exist under the folders defined in the test_paths_file. @@ -267,9 +284,16 @@ 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] = [] + # Calculate deadline + deadline = datetime.datetime.now() + datetime.timedelta( + seconds=max(timeout - WORKER_TIMEOUT_BUFFER_IN_SECONDS, 0) + ) + if len(notebooks) > 0: print(f"Found {len(notebooks)} modified notebooks: {notebooks}") @@ -277,7 +301,9 @@ def process_and_execute_notebooks( 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( @@ -288,6 +314,7 @@ def process_and_execute_notebooks( variable_project_id, variable_region, private_pool_id, + deadline, ), notebooks, ) @@ -301,6 +328,7 @@ def process_and_execute_notebooks( variable_project_id=variable_project_id, variable_region=variable_region, private_pool_id=private_pool_id, + deadline=deadline, notebook=notebook, ) for notebook in notebooks diff --git a/.cloud-build/execute_notebook_cli.py b/.cloud-build/execute_notebook_cli.py index 9545f9c4c..6c9aa55bc 100644 --- a/.cloud-build/execute_notebook_cli.py +++ b/.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.") diff --git a/.cloud-build/execute_notebook_helper.py b/.cloud-build/execute_notebook_helper.py index d59b7b616..3665bbbd0 100644 --- a/.cloud-build/execute_notebook_helper.py +++ b/.cloud-build/execute_notebook_helper.py @@ -15,14 +15,14 @@ """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. diff --git a/.cloud-build/execute_notebook_remote.py b/.cloud-build/execute_notebook_remote.py index 3ca1b3069..f15b01605 100644 --- a/.cloud-build/execute_notebook_remote.py +++ b/.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,12 +36,14 @@ 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, ) -> 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 = { @@ -57,13 +55,14 @@ def execute_notebook_remote( 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 +84,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] diff --git a/.cloud-build/notebook-execution-test-cloudbuild-single.yaml b/.cloud-build/notebook-execution-test-cloudbuild-single.yaml index 5e12755b1..8c4b5c2cc 100644 --- a/.cloud-build/notebook-execution-test-cloudbuild-single.yaml +++ b/.cloud-build/notebook-execution-test-cloudbuild-single.yaml @@ -25,7 +25,4 @@ steps: - '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}"' env: - 'IS_TESTING=1' -timeout: 86400s -options: - pool: - name: ${_PRIVATE_POOL_NAME} \ No newline at end of file +timeout: 86400s \ No newline at end of file diff --git a/.cloud-build/requirements.txt b/.cloud-build/requirements.txt index 45da344a4..9346077c2 100644 --- a/.cloud-build/requirements.txt +++ b/.cloud-build/requirements.txt @@ -9,4 +9,4 @@ tabulate google-cloud-aiplatform google-cloud-storage google-cloud-build -gcloud +ratemate diff --git a/.cloud-build/utils/NotebookProcessors.py b/.cloud-build/utils/NotebookProcessors.py index 04f10906c..2ca802507 100644 --- a/.cloud-build/utils/NotebookProcessors.py +++ b/.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 \ No newline at end of file + return notebook, resources diff --git a/.cloud-build/utils/UpdateNotebookVariables.py b/.cloud-build/utils/UpdateNotebookVariables.py index d4f6d3ae7..d36fbf825 100644 --- a/.cloud-build/utils/UpdateNotebookVariables.py +++ b/.cloud-build/utils/UpdateNotebookVariables.py @@ -78,4 +78,4 @@ def test_region(): variable_name="REGION", variable_value="us-central1", ) - assert new_content == 'REGION = "us-central1" # @param {type:"string"}' \ No newline at end of file + assert new_content == 'REGION = "us-central1" # @param {type:"string"}' diff --git a/.cloud-build/utils/util.py b/.cloud-build/utils/util.py index f9cb6f05e..86eadece6 100644 --- a/.cloud-build/utils/util.py +++ b/.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 + +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,4 @@ def archive_code_and_upload(staging_bucket: str): print(f"Uploaded source code archive to {source_archived_file_gcs}") - return source_archived_file_gcs \ No newline at end of file + return source_archived_file_gcs