| #!/usr/bin/env python |
| # |
| # Copyright 2016 - The Android Open Source Project |
| # |
| # 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. |
| |
| """Public Device Driver APIs. |
| |
| This module provides public device driver APIs that can be called |
| as a Python library. |
| |
| TODO(fdeng): The following APIs have not been implemented |
| - RebootAVD(ip): |
| - RegisterSshPubKey(username, key): |
| - UnregisterSshPubKey(username, key): |
| - CleanupStaleImages(): |
| - CleanupStaleDevices(): |
| """ |
| |
| import datetime |
| import logging |
| import os |
| |
| import dateutil.parser |
| import dateutil.tz |
| |
| from acloud.public import avd |
| from acloud.public import errors |
| from acloud.public import report |
| from acloud.internal import constants |
| from acloud.internal.lib import auth |
| from acloud.internal.lib import android_build_client |
| from acloud.internal.lib import android_compute_client |
| from acloud.internal.lib import gstorage_client |
| from acloud.internal.lib import utils |
| |
| logger = logging.getLogger(__name__) |
| |
| ALL_SCOPES = " ".join([android_build_client.AndroidBuildClient.SCOPE, |
| gstorage_client.StorageClient.SCOPE, |
| android_compute_client.AndroidComputeClient.SCOPE]) |
| |
| MAX_BATCH_CLEANUP_COUNT = 100 |
| |
| |
| class AndroidVirtualDevicePool(object): |
| """A class that manages a pool of devices.""" |
| |
| def __init__(self, cfg, devices=None): |
| self._devices = devices or [] |
| self._cfg = cfg |
| credentials = auth.CreateCredentials(cfg, ALL_SCOPES) |
| self._build_client = android_build_client.AndroidBuildClient( |
| credentials) |
| self._storage_client = gstorage_client.StorageClient(credentials) |
| self._compute_client = android_compute_client.AndroidComputeClient( |
| cfg, credentials) |
| |
| def _CreateGceImageWithBuildInfo(self, build_target, build_id): |
| """Creates a Gce image using build from Launch Control. |
| |
| Clone avd-system.tar.gz of a build to a cache storage bucket |
| using launch control api. And then create a Gce image. |
| |
| Args: |
| build_target: Target name, e.g. "gce_x86-userdebug" |
| build_id: Build id, a string, e.g. "2263051", "P2804227" |
| |
| Returns: |
| String, name of the Gce image that has been created. |
| """ |
| logger.info("Creating a new gce image using build: build_id %s, " |
| "build_target %s", build_id, build_target) |
| disk_image_id = utils.GenerateUniqueName( |
| suffix=self._cfg.disk_image_name) |
| self._build_client.CopyTo( |
| build_target, |
| build_id, |
| artifact_name=self._cfg.disk_image_name, |
| destination_bucket=self._cfg.storage_bucket_name, |
| destination_path=disk_image_id) |
| disk_image_url = self._storage_client.GetUrl( |
| self._cfg.storage_bucket_name, disk_image_id) |
| try: |
| image_name = self._compute_client.GenerateImageName(build_target, |
| build_id) |
| self._compute_client.CreateImage(image_name=image_name, |
| source_uri=disk_image_url) |
| finally: |
| self._storage_client.Delete(self._cfg.storage_bucket_name, |
| disk_image_id) |
| return image_name |
| |
| def _CreateGceImageWithLocalFile(self, local_disk_image): |
| """Create a Gce image with a local image file. |
| |
| The local disk image can be either a tar.gz file or a |
| raw vmlinux image. |
| e.g. /tmp/avd-system.tar.gz or /tmp/android_system_disk_syslinux.img |
| If a raw vmlinux image is provided, it will be archived into a tar.gz file. |
| |
| The final tar.gz file will be uploaded to a cache bucket in storage. |
| |
| Args: |
| local_disk_image: string, path to a local disk image, |
| |
| Returns: |
| String, name of the Gce image that has been created. |
| |
| Raises: |
| DriverError: if a file with an unexpected extension is given. |
| """ |
| logger.info("Creating a new gce image from a local file %s", |
| local_disk_image) |
| with utils.TempDir() as tempdir: |
| if local_disk_image.endswith(self._cfg.disk_raw_image_extension): |
| dest_tar_file = os.path.join(tempdir, |
| self._cfg.disk_image_name) |
| utils.MakeTarFile( |
| src_dict={local_disk_image: self._cfg.disk_raw_image_name}, |
| dest=dest_tar_file) |
| local_disk_image = dest_tar_file |
| elif not local_disk_image.endswith(self._cfg.disk_image_extension): |
| raise errors.DriverError( |
| "Wrong local_disk_image type, must be a *%s file or *%s file" |
| % (self._cfg.disk_raw_image_extension, |
| self._cfg.disk_image_extension)) |
| |
| disk_image_id = utils.GenerateUniqueName( |
| suffix=self._cfg.disk_image_name) |
| self._storage_client.Upload( |
| local_src=local_disk_image, |
| bucket_name=self._cfg.storage_bucket_name, |
| object_name=disk_image_id, |
| mime_type=self._cfg.disk_image_mime_type) |
| disk_image_url = self._storage_client.GetUrl( |
| self._cfg.storage_bucket_name, disk_image_id) |
| try: |
| image_name = self._compute_client.GenerateImageName() |
| self._compute_client.CreateImage(image_name=image_name, |
| source_uri=disk_image_url) |
| finally: |
| self._storage_client.Delete(self._cfg.storage_bucket_name, |
| disk_image_id) |
| return image_name |
| |
| def CreateDevices(self, |
| num, |
| build_target=None, |
| build_id=None, |
| gce_image=None, |
| local_disk_image=None, |
| cleanup=True, |
| extra_data_disk_size_gb=None, |
| precreated_data_image=None): |
| """Creates |num| devices for given build_target and build_id. |
| |
| - If gce_image is provided, will use it to create an instance. |
| - If local_disk_image is provided, will upload it to a temporary |
| caching storage bucket which is defined by user as |storage_bucket_name| |
| And then create an gce image with it; and then create an instance. |
| - If build_target and build_id are provided, will clone the disk image |
| via launch control to the temporary caching storage bucket. |
| And then create an gce image with it; and then create an instance. |
| |
| Args: |
| num: Number of devices to create. |
| build_target: Target name, e.g. "gce_x86-userdebug" |
| build_id: Build id, a string, e.g. "2263051", "P2804227" |
| gce_image: string, if given, will use this image |
| instead of creating a new one. |
| implies cleanup=False. |
| local_disk_image: string, path to a local disk image, e.g. |
| /tmp/avd-system.tar.gz |
| cleanup: boolean, if True clean up compute engine image after creating |
| the instance. |
| extra_data_disk_size_gb: Integer, size of extra disk, or None. |
| precreated_data_image: A string, the image to use for the extra disk. |
| |
| Raises: |
| errors.DriverError: If no source is specified for image creation. |
| """ |
| if gce_image: |
| # GCE image is provided, we can directly move to instance creation. |
| logger.info("Using existing gce image %s", gce_image) |
| image_name = gce_image |
| cleanup = False |
| elif local_disk_image: |
| image_name = self._CreateGceImageWithLocalFile(local_disk_image) |
| elif build_target and build_id: |
| image_name = self._CreateGceImageWithBuildInfo(build_target, |
| build_id) |
| else: |
| raise errors.DriverError( |
| "Invalid image source, must specify one of the following: gce_image, " |
| "local_disk_image, or build_target and build id.") |
| |
| # Create GCE instances. |
| try: |
| for _ in range(num): |
| instance = self._compute_client.GenerateInstanceName( |
| build_target, build_id) |
| extra_disk_name = None |
| if extra_data_disk_size_gb > 0: |
| extra_disk_name = self._compute_client.GetDataDiskName( |
| instance) |
| self._compute_client.CreateDisk(extra_disk_name, |
| precreated_data_image, |
| extra_data_disk_size_gb) |
| self._compute_client.CreateInstance(instance, image_name, |
| extra_disk_name) |
| ip = self._compute_client.GetInstanceIP(instance) |
| self.devices.append(avd.AndroidVirtualDevice( |
| ip=ip, instance_name=instance)) |
| finally: |
| if cleanup: |
| self._compute_client.DeleteImage(image_name) |
| |
| def DeleteDevices(self): |
| """Deletes devices. |
| |
| Returns: |
| A tuple, (deleted, failed, error_msgs) |
| deleted: A list of names of instances that have been deleted. |
| faild: A list of names of instances that we fail to delete. |
| error_msgs: A list of failure messages. |
| """ |
| instance_names = [device.instance_name for device in self._devices] |
| return self._compute_client.DeleteInstances(instance_names, |
| self._cfg.zone) |
| |
| def WaitForBoot(self): |
| """Waits for all devices to boot up. |
| |
| Returns: |
| A dictionary that contains all the failures. |
| The key is the name of the instance that fails to boot, |
| the value is an errors.DeviceBootTimeoutError object. |
| """ |
| failures = {} |
| for device in self._devices: |
| try: |
| self._compute_client.WaitForBoot(device.instance_name) |
| except errors.DeviceBootTimeoutError as e: |
| failures[device.instance_name] = e |
| return failures |
| |
| @property |
| def devices(self): |
| """Returns a list of devices in the pool. |
| |
| Returns: |
| A list of devices in the pool. |
| """ |
| return self._devices |
| |
| |
| def _AddDeletionResultToReport(report_obj, deleted, failed, error_msgs, |
| resource_name): |
| """Adds deletion result to a Report object. |
| |
| This function will add the following to report.data. |
| "deleted": [ |
| {"name": "resource_name", "type": "resource_name"}, |
| ], |
| "failed": [ |
| {"name": "resource_name", "type": "resource_name"}, |
| ], |
| This function will append error_msgs to report.errors. |
| |
| Args: |
| report_obj: A Report object. |
| deleted: A list of names of the resources that have been deleted. |
| failed: A list of names of the resources that we fail to delete. |
| error_msgs: A list of error message strings to be added to the report. |
| resource_name: A string, representing the name of the resource. |
| """ |
| for name in deleted: |
| report_obj.AddData(key="deleted", |
| value={"name": name, |
| "type": resource_name}) |
| for name in failed: |
| report_obj.AddData(key="failed", |
| value={"name": name, |
| "type": resource_name}) |
| report_obj.AddErrors(error_msgs) |
| if failed or error_msgs: |
| report_obj.SetStatus(report.Status.FAIL) |
| |
| |
| def _FetchSerialLogsFromDevices(compute_client, instance_names, output_file, |
| port): |
| """Fetch serial logs from a port for a list of devices to a local file. |
| |
| Args: |
| compute_client: An object of android_compute_client.AndroidComputeClient |
| instance_names: A list of instance names. |
| output_file: A path to a file ending with "tar.gz" |
| port: The number of serial port to read from, 0 for serial output, 1 for |
| logcat. |
| """ |
| with utils.TempDir() as tempdir: |
| src_dict = {} |
| for instance_name in instance_names: |
| serial_log = compute_client.GetSerialPortOutput( |
| instance=instance_name, port=port) |
| file_name = "%s.log" % instance_name |
| file_path = os.path.join(tempdir, file_name) |
| src_dict[file_path] = file_name |
| with open(file_path, "w") as f: |
| f.write(serial_log.encode("utf-8")) |
| utils.MakeTarFile(src_dict, output_file) |
| |
| |
| def _CreateSshKeyPairIfNecessary(cfg): |
| """Create ssh key pair if necessary. |
| |
| Args: |
| cfg: An Acloudconfig instance. |
| |
| Raises: |
| error.DriverError: If it falls into an unexpected condition. |
| """ |
| if not cfg.ssh_public_key_path: |
| logger.warning("ssh_public_key_path is not specified in acloud config. " |
| "Project-wide public key will " |
| "be used when creating AVD instances. " |
| "Please ensure you have the correct private half of " |
| "a project-wide public key if you want to ssh into the " |
| "instances after creation.") |
| elif cfg.ssh_public_key_path and not cfg.ssh_private_key_path: |
| logger.warning("Only ssh_public_key_path is specified in acloud config," |
| " but ssh_private_key_path is missing. " |
| "Please ensure you have the correct private half " |
| "if you want to ssh into the instances after creation.") |
| elif cfg.ssh_public_key_path and cfg.ssh_private_key_path: |
| utils.CreateSshKeyPairIfNotExist( |
| cfg.ssh_private_key_path, cfg.ssh_public_key_path) |
| else: |
| # Should never reach here. |
| raise errors.DriverError( |
| "Unexpected error in _CreateSshKeyPairIfNecessary") |
| |
| |
| def CreateAndroidVirtualDevices(cfg, |
| build_target=None, |
| build_id=None, |
| num=1, |
| gce_image=None, |
| local_disk_image=None, |
| cleanup=True, |
| serial_log_file=None, |
| logcat_file=None): |
| """Creates one or multiple android devices. |
| |
| Args: |
| cfg: An AcloudConfig instance. |
| build_target: Target name, e.g. "gce_x86-userdebug" |
| build_id: Build id, a string, e.g. "2263051", "P2804227" |
| num: Number of devices to create. |
| gce_image: string, if given, will use this gce image |
| instead of creating a new one. |
| implies cleanup=False. |
| local_disk_image: string, path to a local disk image, e.g. |
| /tmp/avd-system.tar.gz |
| cleanup: boolean, if True clean up compute engine image and |
| disk image in storage after creating the instance. |
| serial_log_file: A path to a file where serial output should |
| be saved to. |
| logcat_file: A path to a file where logcat logs should be saved. |
| |
| Returns: |
| A Report instance. |
| """ |
| r = report.Report(command="create") |
| credentials = auth.CreateCredentials(cfg, ALL_SCOPES) |
| compute_client = android_compute_client.AndroidComputeClient(cfg, |
| credentials) |
| try: |
| _CreateSshKeyPairIfNecessary(cfg) |
| device_pool = AndroidVirtualDevicePool(cfg) |
| device_pool.CreateDevices( |
| num, |
| build_target, |
| build_id, |
| gce_image, |
| local_disk_image, |
| cleanup, |
| extra_data_disk_size_gb=cfg.extra_data_disk_size_gb, |
| precreated_data_image=cfg.precreated_data_image_map.get( |
| cfg.extra_data_disk_size_gb)) |
| failures = device_pool.WaitForBoot() |
| # Write result to report. |
| for device in device_pool.devices: |
| device_dict = {"ip": device.ip, |
| "instance_name": device.instance_name} |
| if device.instance_name in failures: |
| r.AddData(key="devices_failing_boot", value=device_dict) |
| r.AddError(str(failures[device.instance_name])) |
| else: |
| r.AddData(key="devices", value=device_dict) |
| if failures: |
| r.SetStatus(report.Status.BOOT_FAIL) |
| else: |
| r.SetStatus(report.Status.SUCCESS) |
| |
| # Dump serial and logcat logs. |
| if serial_log_file: |
| _FetchSerialLogsFromDevices( |
| compute_client, |
| instance_names=[d.instance_name for d in device_pool.devices], |
| port=constants.DEFAULT_SERIAL_PORT, |
| output_file=serial_log_file) |
| if logcat_file: |
| _FetchSerialLogsFromDevices( |
| compute_client, |
| instance_names=[d.instance_name for d in device_pool.devices], |
| port=constants.LOGCAT_SERIAL_PORT, |
| output_file=logcat_file) |
| except errors.DriverError as e: |
| r.AddError(str(e)) |
| r.SetStatus(report.Status.FAIL) |
| return r |
| |
| |
| def DeleteAndroidVirtualDevices(cfg, instance_names): |
| """Deletes android devices. |
| |
| Args: |
| cfg: An AcloudConfig instance. |
| instance_names: A list of names of the instances to delete. |
| |
| Returns: |
| A Report instance. |
| """ |
| r = report.Report(command="delete") |
| credentials = auth.CreateCredentials(cfg, ALL_SCOPES) |
| compute_client = android_compute_client.AndroidComputeClient(cfg, |
| credentials) |
| try: |
| deleted, failed, error_msgs = compute_client.DeleteInstances( |
| instance_names, cfg.zone) |
| _AddDeletionResultToReport( |
| r, deleted, |
| failed, error_msgs, |
| resource_name="instance") |
| if r.status == report.Status.UNKNOWN: |
| r.SetStatus(report.Status.SUCCESS) |
| except errors.DriverError as e: |
| r.AddError(str(e)) |
| r.SetStatus(report.Status.FAIL) |
| return r |
| |
| |
| def _FindOldItems(items, cut_time, time_key): |
| """Finds items from |items| whose timestamp is earlier than |cut_time|. |
| |
| Args: |
| items: A list of items. Each item is a dictionary represent |
| the properties of the item. It should has a key as noted |
| by time_key. |
| cut_time: A datetime.datatime object. |
| time_key: String, key for the timestamp. |
| |
| Returns: |
| A list of those from |items| whose timestamp is earlier than cut_time. |
| """ |
| cleanup_list = [] |
| for item in items: |
| t = dateutil.parser.parse(item[time_key]) |
| if t < cut_time: |
| cleanup_list.append(item) |
| return cleanup_list |
| |
| |
| def Cleanup(cfg, expiration_mins): |
| """Cleans up stale gce images, gce instances, and disk images in storage. |
| |
| Args: |
| cfg: An AcloudConfig instance. |
| expiration_mins: Integer, resources older than |expiration_mins| will |
| be cleaned up. |
| |
| Returns: |
| A Report instance. |
| """ |
| r = report.Report(command="cleanup") |
| try: |
| cut_time = (datetime.datetime.now(dateutil.tz.tzlocal()) - |
| datetime.timedelta(minutes=expiration_mins)) |
| logger.info( |
| "Cleaning up any gce images/instances and cached build artifacts." |
| "in google storage that are older than %s", cut_time) |
| credentials = auth.CreateCredentials(cfg, ALL_SCOPES) |
| compute_client = android_compute_client.AndroidComputeClient( |
| cfg, credentials) |
| storage_client = gstorage_client.StorageClient(credentials) |
| |
| # Cleanup expired instances |
| items = compute_client.ListInstances(zone=cfg.zone) |
| cleanup_list = [ |
| item["name"] |
| for item in _FindOldItems(items, cut_time, "creationTimestamp") |
| ] |
| logger.info("Found expired instances: %s", cleanup_list) |
| for i in range(0, len(cleanup_list), MAX_BATCH_CLEANUP_COUNT): |
| result = compute_client.DeleteInstances( |
| instances=cleanup_list[i:i + MAX_BATCH_CLEANUP_COUNT], |
| zone=cfg.zone) |
| _AddDeletionResultToReport(r, *result, resource_name="instance") |
| |
| # Cleanup expired images |
| items = compute_client.ListImages() |
| skip_list = cfg.precreated_data_image_map.viewvalues() |
| cleanup_list = [ |
| item["name"] |
| for item in _FindOldItems(items, cut_time, "creationTimestamp") |
| if item["name"] not in skip_list |
| ] |
| logger.info("Found expired images: %s", cleanup_list) |
| for i in range(0, len(cleanup_list), MAX_BATCH_CLEANUP_COUNT): |
| result = compute_client.DeleteImages( |
| image_names=cleanup_list[i:i + MAX_BATCH_CLEANUP_COUNT]) |
| _AddDeletionResultToReport(r, *result, resource_name="image") |
| |
| # Cleanup expired disks |
| # Disks should have been attached to instances with autoDelete=True. |
| # However, sometimes disks may not be auto deleted successfully. |
| items = compute_client.ListDisks(zone=cfg.zone) |
| cleanup_list = [ |
| item["name"] |
| for item in _FindOldItems(items, cut_time, "creationTimestamp") |
| if not item.get("users") |
| ] |
| logger.info("Found expired disks: %s", cleanup_list) |
| for i in range(0, len(cleanup_list), MAX_BATCH_CLEANUP_COUNT): |
| result = compute_client.DeleteDisks( |
| disk_names=cleanup_list[i:i + MAX_BATCH_CLEANUP_COUNT], |
| zone=cfg.zone) |
| _AddDeletionResultToReport(r, *result, resource_name="disk") |
| |
| # Cleanup expired google storage |
| items = storage_client.List(bucket_name=cfg.storage_bucket_name) |
| cleanup_list = [ |
| item["name"] |
| for item in _FindOldItems(items, cut_time, "timeCreated") |
| ] |
| logger.info("Found expired cached artifacts: %s", cleanup_list) |
| for i in range(0, len(cleanup_list), MAX_BATCH_CLEANUP_COUNT): |
| result = storage_client.DeleteFiles( |
| bucket_name=cfg.storage_bucket_name, |
| object_names=cleanup_list[i:i + MAX_BATCH_CLEANUP_COUNT]) |
| _AddDeletionResultToReport( |
| r, *result, resource_name="cached_build_artifact") |
| |
| # Everything succeeded, write status to report. |
| if r.status == report.Status.UNKNOWN: |
| r.SetStatus(report.Status.SUCCESS) |
| except errors.DriverError as e: |
| r.AddError(str(e)) |
| r.SetStatus(report.Status.FAIL) |
| return r |
| |
| |
| def AddSshRsa(cfg, user, ssh_rsa_path): |
| """Add public ssh rsa key to the project. |
| |
| Args: |
| cfg: An AcloudConfig instance. |
| user: the name of the user which the key belongs to. |
| ssh_rsa_path: The absolute path to public rsa key. |
| |
| Returns: |
| A Report instance. |
| """ |
| r = report.Report(command="sshkey") |
| try: |
| credentials = auth.CreateCredentials(cfg, ALL_SCOPES) |
| compute_client = android_compute_client.AndroidComputeClient( |
| cfg, credentials) |
| compute_client.AddSshRsa(user, ssh_rsa_path) |
| r.SetStatus(report.Status.SUCCESS) |
| except errors.DriverError as e: |
| r.AddError(str(e)) |
| r.SetStatus(report.Status.FAIL) |
| return r |
| |
| |
| def CheckAccess(cfg): |
| """Check if user has access. |
| |
| Args: |
| cfg: An AcloudConfig instance. |
| """ |
| credentials = auth.CreateCredentials(cfg, ALL_SCOPES) |
| compute_client = android_compute_client.AndroidComputeClient( |
| cfg, credentials) |
| logger.info("Checking if user has access to project %s", cfg.project) |
| if not compute_client.CheckAccess(): |
| logger.error("User does not have access to project %s", cfg.project) |
| # Print here so that command line user can see it. |
| print "Looks like you do not have access to %s. " % cfg.project |
| if cfg.project in cfg.no_project_access_msg_map: |
| print cfg.no_project_access_msg_map[cfg.project] |