Files
model_garden/.cloud-build/execute_notebook_remote.py
f184411666 Create a function to get notebook python version for execution (#828)
* Create a function to get notebook python version for execution

* Inject python version to yaml file

* Fix python version references

* Add python string to python version variable

* add python version extraction script (untested)

* Create a function to get notebook python version for execution

* Inject python version to yaml file

* Fix python version references

* Add python string to python version variable

* Remove one notebook condition

* Remove extra check and use python 3 as default version

* Use python3.9 as default version

* Update python version notebook parser

* Add python version to the notebook template

* Fix bug

* Update python version parser function

* Add python version to a notebook for testing

* Run linter

* Use regex in python version parser function

* Add new notebook for testing

* Use better variable name

* fix typo

* Use f string instead +

* Use python from env instead of using _PYTHON_VERSION

* Use simpler regex

* Add python version test notebook

* Fixed a mistake

* Updated notebook template with python version

* Fixed python version format

* Print log contents to stdout

* Remove failing notebook

* Run linter

* Remove extra steps in the printed log

* Edit comments

* Run linter

* Run linter

* Revert test changes

Co-authored-by: AG Sol <aarongabriel@google.com>
2022-09-26 14:32:53 -07:00

106 lines
3.5 KiB
Python

#!/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 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
CLOUD_BUILD_FILEPATH = ".cloud-build/notebook-execution-test-cloudbuild-single.yaml"
SERVICE_BASE_PATH = "cloudbuild.googleapis.com"
def execute_notebook_remote(
code_archive_uri: str,
notebook_uri: str,
notebook_output_uri: str,
container_uri: 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 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"{private_pool_region}-{SERVICE_BASE_PATH}"
)
# Authorize the client with Google defaults
credentials, project_id = google.auth.default()
client = cloudbuild_v1.services.cloud_build.CloudBuildClient(client_options=options)
(
source_archived_file_gcs_bucket,
source_archived_file_gcs_object,
) = utils.extract_bucket_and_prefix_from_gcs_path(code_archive_uri)
build.source = Source(
storage_source=StorageSource(
bucket=source_archived_file_gcs_bucket,
object_=source_archived_file_gcs_object,
)
)
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)
if tag:
build.tags = [tag]
operation = client.create_build(project_id=project_id, build=build)
# Print the in-progress operation
# print("IN PROGRESS:")
# print(operation.metadata)
# Print the completed status
# print("RESULT:", result.status)
return operation