model_garden/.cloud-build/execute_changed_notebooks_helper.py

#!/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 concurrent
import dataclasses
import datetime
import functools
import git
import operator
import os
import pathlib
import re
import subprocess
from typing import List, Optional

import execute_notebook_helper
import execute_notebook_remote
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:
    """Formats a timedelta duration to [N days] %H:%M:%S format"""
    seconds = int(delta.total_seconds())

    secs_in_a_day = 86400
    secs_in_a_hour = 3600
    secs_in_a_min = 60

    days, seconds = divmod(seconds, secs_in_a_day)
    hours, seconds = divmod(seconds, secs_in_a_hour)
    minutes, seconds = divmod(seconds, secs_in_a_min)

    time_fmt = f"{hours:02d}:{minutes:02d}:{seconds:02d}"

    if days > 0:
        suffix = "s" if days > 1 else ""
        return f"{days} day{suffix} {time_fmt}"

    return time_fmt


@dataclasses.dataclass
class NotebookExecutionResult:
    name: str
    duration: datetime.timedelta
    is_pass: bool
    log_url: str
    output_uri: str
    build_id: 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,
):
    # Read notebook
    with open(notebook_path) as f:
        nb = nbformat.read(f, as_version=4)

    # Create preprocessors
    remove_no_execute_cells_preprocessor = NotebookProcessors.RemoveNoExecuteCells()
    update_variables_preprocessor = NotebookProcessors.UpdateVariablesPreprocessor(
        replacement_map={
            "PROJECT_ID": variable_project_id,
            "REGION": variable_region,
            "SERVICE_ACCOUNT": variable_service_account,
        },
    )

    # Use no-execute preprocessor
    (
        nb,
        resources,
    ) = remove_no_execute_cells_preprocessor.preprocess(nb)

    (nb, resources) = update_variables_preprocessor.preprocess(nb, resources)

    with open(notebook_path, mode="w", encoding="utf-8") as new_file:
        nbformat.write(nb, new_file)


def _create_tag(filepath: str) -> str:
    tag = os.path.basename(os.path.normpath(filepath))
    tag = re.sub("[^0-9a-zA-Z_.-]+", "-", tag)

    if tag.startswith(".") or tag.startswith("-"):
        tag = tag[1:]

    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,
    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
    notebook_output_uri = "/".join([artifacts_bucket, pathlib.Path(notebook).name])

    # Create tag from notebook
    tag = _create_tag(filepath=notebook)

    result = NotebookExecutionResult(
        name=tag,
        duration=datetime.timedelta(seconds=0),
        is_pass=False,
        output_uri=notebook_output_uri,
        log_url="",
        build_id="",
        error_message=None,
    )

    # TODO: Handle cases where multiple notebooks have the same name
    time_start = datetime.datetime.now()
    operation = None
    try:
        # 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,
        )

        # 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,
            private_pool_id=private_pool_id,
            private_pool_region=variable_region,
            timeout_in_seconds=timeout_in_seconds,
        )

        operation_metadata = BuildOperationMetadata(mapping=operation.metadata)
        result.build_id = operation_metadata.build.id
        result.log_url = operation_metadata.build.log_url

        # Block and wait for the result
        operation_result = operation.result()

        result.duration = datetime.datetime.now() - time_start
        result.is_pass = True
        print(f"{notebook} PASSED in {format_timedelta(result.duration)}.")
    except Exception as error:
        result.error_message = str(error)

        if operation and should_get_tail_logs:
            # Extract the logs
            logs_bucket = operation_metadata.build.logs_bucket

            # Download tail end of logs file
            log_file_uri = f"{logs_bucket}/log-{result.build_id}.txt"

            # Use gcloud to get tail
            try:
                result.error_message = subprocess.check_output(
                    ["gsutil", "cat", "-r", "-1000", log_file_uri], encoding="UTF-8"
                )
            except Exception as error:
                result.error_message = str(error)

        result.duration = datetime.datetime.now() - time_start
        result.is_pass = False

        print(
            f"{notebook} FAILED in {format_timedelta(result.duration)}: {result.error_message}"
        )

    return result


def get_changed_notebooks(
    test_paths_file: str,
    base_branch: Optional[str] = None,
) -> List[str]:
    """
    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 = []
    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 = []

    # Instantiate GitPython objects
    repo = git.Repo(os.getcwd())
    index = repo.index

    if base_branch:
        # 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 = 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


def process_and_execute_notebooks(
    notebooks: List[str],
    container_uri: str,
    staging_bucket: str,
    artifacts_bucket: str,
    variable_project_id: str,
    variable_region: str,
    variable_service_account: 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.
    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.
        timeout (str):
            Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
    """

    # 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=100) as executor:
                print(f"Max workers: {executor._max_workers}")

                notebook_execution_results = list(
                    executor.map(
                        functools.partial(
                            process_and_execute_notebook,
                            container_uri,
                            staging_bucket,
                            artifacts_bucket,
                            variable_project_id,
                            variable_region,
                            variable_service_account,
                            private_pool_id,
                            deadline,
                        ),
                        notebooks,
                    )
                )
        else:
            notebook_execution_results = [
                process_and_execute_notebook(
                    container_uri=container_uri,
                    staging_bucket=staging_bucket,
                    artifacts_bucket=artifacts_bucket,
                    variable_project_id=variable_project_id,
                    variable_region=variable_region,
                    variable_service_account=variable_service_account,
                    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
                    ]
                    for result in results_sorted
                ],
                headers=["build_tag", "status", "duration", "log_url", "output_uri", "output_uri_web"],
            )
        )

        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")

    elif len(notebooks) == 1:
        notebook = notebooks[0]

        # 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,
        )

        execute_notebook_helper.execute_notebook(
            notebook_source=notebook,
            output_file_or_uri="/".join(
                [artifacts_bucket, pathlib.Path(notebook).name]
            ),
            should_log_output=True,
        )
    else:
        print("No notebooks modified in this pull request.")