diff options
Diffstat (limited to 'crosperf/download_images.py')
| -rw-r--r-- | crosperf/download_images.py | 286 |
1 files changed, 286 insertions, 0 deletions
diff --git a/crosperf/download_images.py b/crosperf/download_images.py new file mode 100644 index 00000000..8ceaa874 --- /dev/null +++ b/crosperf/download_images.py @@ -0,0 +1,286 @@ +# Copyright (c) 2014-2015 The Chromium OS Authors. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. +"""Download images from Cloud Storage.""" + +from __future__ import print_function + +import ast +import os + +import test_flag + +from cros_utils import command_executer + +GS_UTIL = 'chromium/tools/depot_tools/gsutil.py' + + +class MissingImage(Exception): + """Raised when the requested image does not exist in gs://""" + + +class MissingFile(Exception): + """Raised when the requested file does not exist in gs://""" + + +class RunCommandExceptionHandler(object): + """Handle Exceptions from calls to RunCommand""" + + def __init__(self, logger_to_use, log_level, cmd_exec, command): + self.logger = logger_to_use + self.log_level = log_level + self.ce = cmd_exec + self.cleanup_command = command + + def HandleException(self, _, e): + # Exception handler, Run specified command + if self.log_level != 'verbose' and self.cleanup_command is not None: + self.logger.LogOutput('CMD: %s' % self.cleanup_command) + if self.cleanup_command is not None: + _ = self.ce.RunCommand(self.cleanup_command) + # Raise exception again + raise e + + +class ImageDownloader(object): + """Download images from Cloud Storage.""" + + def __init__(self, logger_to_use=None, log_level='verbose', cmd_exec=None): + self._logger = logger_to_use + self.log_level = log_level + self._ce = cmd_exec or command_executer.GetCommandExecuter( + self._logger, log_level=self.log_level) + + def GetBuildID(self, chromeos_root, xbuddy_label): + # Get the translation of the xbuddy_label into the real Google Storage + # image name. + command = ('cd ~/trunk/src/third_party/toolchain-utils/crosperf; ' + "python translate_xbuddy.py '%s'" % xbuddy_label) + _, build_id_tuple_str, _ = self._ce.ChrootRunCommandWOutput(chromeos_root, + command) + if not build_id_tuple_str: + raise MissingImage("Unable to find image for '%s'" % xbuddy_label) + + build_id_tuple = ast.literal_eval(build_id_tuple_str) + build_id = build_id_tuple[0] + + return build_id + + def DownloadImage(self, chromeos_root, build_id, image_name): + if self.log_level == 'average': + self._logger.LogOutput('Preparing to download %s image to local ' + 'directory.' % build_id) + + # Make sure the directory for downloading the image exists. + download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id) + image_path = os.path.join(download_path, 'chromiumos_test_image.bin') + if not os.path.exists(download_path): + os.makedirs(download_path) + + # Check to see if the image has already been downloaded. If not, + # download the image. + if not os.path.exists(image_path): + gsutil_cmd = os.path.join(chromeos_root, GS_UTIL) + command = '%s cp %s %s' % (gsutil_cmd, image_name, download_path) + + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % command) + status = self._ce.RunCommand(command) + downloaded_image_name = os.path.join(download_path, + 'chromiumos_test_image.tar.xz') + if status != 0 or not os.path.exists(downloaded_image_name): + raise MissingImage('Cannot download image: %s.' % downloaded_image_name) + + return image_path + + def UncompressImage(self, chromeos_root, build_id): + # Check to see if the file has already been uncompresssed, etc. + if os.path.exists( + os.path.join(chromeos_root, 'chroot/tmp', build_id, + 'chromiumos_test_image.bin')): + return + + # Uncompress and untar the downloaded image. + download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id) + command = ('cd %s ; tar -Jxf chromiumos_test_image.tar.xz ' % download_path) + # Cleanup command for exception handler + clean_cmd = ('cd %s ; rm -f chromiumos_test_image.bin ' % download_path) + exception_handler = RunCommandExceptionHandler(self._logger, self.log_level, + self._ce, clean_cmd) + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % command) + print('(Uncompressing and un-tarring may take a couple of minutes...' + 'please be patient.)') + retval = self._ce.RunCommand( + command, except_handler=exception_handler.HandleException) + if retval != 0: + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % clean_cmd) + print('(Removing file chromiumos_test_image.bin.)') + # Remove partially uncompressed file + _ = self._ce.RunCommand(clean_cmd) + # Raise exception for failure to uncompress + raise MissingImage('Cannot uncompress image: %s.' % build_id) + + # Remove compressed image + command = ('cd %s ; rm -f chromiumos_test_image.tar.xz; ' % download_path) + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % command) + print('(Removing file chromiumos_test_image.tar.xz.)') + # try removing file, its ok to have an error, print if encountered + retval = self._ce.RunCommand(command) + if retval != 0: + print('(Warning: Could not remove file chromiumos_test_image.tar.xz .)') + + def DownloadSingleAutotestFile(self, chromeos_root, build_id, + package_file_name): + # Verify if package files exist + status = 0 + gs_package_name = ('gs://chromeos-image-archive/%s/%s' % + (build_id, package_file_name)) + gsutil_cmd = os.path.join(chromeos_root, GS_UTIL) + if not test_flag.GetTestMode(): + cmd = '%s ls %s' % (gsutil_cmd, gs_package_name) + status = self._ce.RunCommand(cmd) + if status != 0: + raise MissingFile('Cannot find autotest package file: %s.' % + package_file_name) + + if self.log_level == 'average': + self._logger.LogOutput('Preparing to download %s package to local ' + 'directory.' % package_file_name) + + # Make sure the directory for downloading the package exists. + download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id) + package_path = os.path.join(download_path, package_file_name) + if not os.path.exists(download_path): + os.makedirs(download_path) + + # Check to see if the package file has already been downloaded. If not, + # download it. + if not os.path.exists(package_path): + command = '%s cp %s %s' % (gsutil_cmd, gs_package_name, download_path) + + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % command) + status = self._ce.RunCommand(command) + if status != 0 or not os.path.exists(package_path): + raise MissingFile('Cannot download package: %s .' % package_path) + + def UncompressSingleAutotestFile(self, chromeos_root, build_id, + package_file_name, uncompress_cmd): + # Uncompress file + download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id) + command = ('cd %s ; %s %s' % + (download_path, uncompress_cmd, package_file_name)) + + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % command) + print('(Uncompressing autotest file %s .)' % package_file_name) + retval = self._ce.RunCommand(command) + if retval != 0: + raise MissingFile('Cannot uncompress file: %s.' % package_file_name) + # Remove uncompressed downloaded file + command = ('cd %s ; rm -f %s' % (download_path, package_file_name)) + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % command) + print('(Removing processed autotest file %s .)' % package_file_name) + # try removing file, its ok to have an error, print if encountered + retval = self._ce.RunCommand(command) + if retval != 0: + print('(Warning: Could not remove file %s .)' % package_file_name) + + def VerifyAutotestFilesExist(self, chromeos_root, build_id, package_file): + # Quickly verify if the files are there + status = 0 + gs_package_name = ('gs://chromeos-image-archive/%s/%s' % + (build_id, package_file)) + gsutil_cmd = os.path.join(chromeos_root, GS_UTIL) + if not test_flag.GetTestMode(): + cmd = '%s ls %s' % (gsutil_cmd, gs_package_name) + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % cmd) + status = self._ce.RunCommand(cmd) + if status != 0: + print('(Warning: Could not find file %s )' % gs_package_name) + return 1 + # Package exists on server + return 0 + + def DownloadAutotestFiles(self, chromeos_root, build_id): + # Download autest package files (3 files) + autotest_packages_name = ('autotest_packages.tar') + autotest_server_package_name = ('autotest_server_package.tar.bz2') + autotest_control_files_name = ('control_files.tar') + + download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id) + # Autotest directory relative path wrt chroot + autotest_rel_path = os.path.join('/tmp', build_id, 'autotest_files') + # Absolute Path to download files + autotest_path = os.path.join(chromeos_root, 'chroot/tmp', build_id, + 'autotest_files') + + if not os.path.exists(autotest_path): + # Quickly verify if the files are present on server + # If not, just exit with warning + status = self.VerifyAutotestFilesExist(chromeos_root, build_id, + autotest_packages_name) + if status != 0: + default_autotest_dir = '~/trunk/src/third_party/autotest/files' + print('(Warning: Could not find autotest packages .)\n' + '(Warning: Defaulting autotest path to %s .' % + default_autotest_dir) + return default_autotest_dir + + # Files exist on server, download and uncompress them + self.DownloadSingleAutotestFile(chromeos_root, build_id, + autotest_packages_name) + self.DownloadSingleAutotestFile(chromeos_root, build_id, + autotest_server_package_name) + self.DownloadSingleAutotestFile(chromeos_root, build_id, + autotest_control_files_name) + + self.UncompressSingleAutotestFile(chromeos_root, build_id, + autotest_packages_name, 'tar -xvf ') + self.UncompressSingleAutotestFile(chromeos_root, build_id, + autotest_server_package_name, + 'tar -jxvf ') + self.UncompressSingleAutotestFile(chromeos_root, build_id, + autotest_control_files_name, + 'tar -xvf ') + # Rename created autotest directory to autotest_files + command = ('cd %s ; mv autotest autotest_files' % download_path) + if self.log_level != 'verbose': + self._logger.LogOutput('CMD: %s' % command) + print('(Moving downloaded autotest files to autotest_files)') + retval = self._ce.RunCommand(command) + if retval != 0: + raise MissingFile('Could not create directory autotest_files') + + return autotest_rel_path + + def Run(self, chromeos_root, xbuddy_label, autotest_path): + build_id = self.GetBuildID(chromeos_root, xbuddy_label) + image_name = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz' + % build_id) + + # Verify that image exists for build_id, before attempting to + # download it. + status = 0 + if not test_flag.GetTestMode(): + gsutil_cmd = os.path.join(chromeos_root, GS_UTIL) + cmd = '%s ls %s' % (gsutil_cmd, image_name) + status = self._ce.RunCommand(cmd) + if status != 0: + raise MissingImage('Cannot find official image: %s.' % image_name) + + image_path = self.DownloadImage(chromeos_root, build_id, image_name) + self.UncompressImage(chromeos_root, build_id) + + if self.log_level != 'quiet': + self._logger.LogOutput('Using image from %s.' % image_path) + + if autotest_path == '': + autotest_path = self.DownloadAutotestFiles(chromeos_root, build_id) + + return image_path, autotest_path |