Home | History | Annotate | Download | only in crosperf
      1 # Copyright (c) 2014-2015 The Chromium OS Authors. All rights reserved.
      2 # Use of this source code is governed by a BSD-style license that can be
      3 # found in the LICENSE file.
      4 """Download images from Cloud Storage."""
      5 
      6 from __future__ import print_function
      7 
      8 import ast
      9 import os
     10 
     11 import test_flag
     12 
     13 from cros_utils import command_executer
     14 
     15 GS_UTIL = 'chromium/tools/depot_tools/gsutil.py'
     16 
     17 
     18 class MissingImage(Exception):
     19   """Raised when the requested image does not exist in gs://"""
     20 
     21 
     22 class MissingFile(Exception):
     23   """Raised when the requested file does not exist in gs://"""
     24 
     25 
     26 class RunCommandExceptionHandler(object):
     27   """Handle Exceptions from calls to RunCommand"""
     28 
     29   def __init__(self, logger_to_use, log_level, cmd_exec, command):
     30     self.logger = logger_to_use
     31     self.log_level = log_level
     32     self.ce = cmd_exec
     33     self.cleanup_command = command
     34 
     35   def HandleException(self, _, e):
     36     # Exception handler, Run specified command
     37     if self.log_level != 'verbose' and self.cleanup_command is not None:
     38       self.logger.LogOutput('CMD: %s' % self.cleanup_command)
     39     if self.cleanup_command is not None:
     40       _ = self.ce.RunCommand(self.cleanup_command)
     41     # Raise exception again
     42     raise e
     43 
     44 
     45 class ImageDownloader(object):
     46   """Download images from Cloud Storage."""
     47 
     48   def __init__(self, logger_to_use=None, log_level='verbose', cmd_exec=None):
     49     self._logger = logger_to_use
     50     self.log_level = log_level
     51     self._ce = cmd_exec or command_executer.GetCommandExecuter(
     52         self._logger, log_level=self.log_level)
     53 
     54   def GetBuildID(self, chromeos_root, xbuddy_label):
     55     # Get the translation of the xbuddy_label into the real Google Storage
     56     # image name.
     57     command = ('cd ~/trunk/src/third_party/toolchain-utils/crosperf; '
     58                "python translate_xbuddy.py '%s'" % xbuddy_label)
     59     _, build_id_tuple_str, _ = self._ce.ChrootRunCommandWOutput(chromeos_root,
     60                                                                 command)
     61     if not build_id_tuple_str:
     62       raise MissingImage("Unable to find image for '%s'" % xbuddy_label)
     63 
     64     build_id_tuple = ast.literal_eval(build_id_tuple_str)
     65     build_id = build_id_tuple[0]
     66 
     67     return build_id
     68 
     69   def DownloadImage(self, chromeos_root, build_id, image_name):
     70     if self.log_level == 'average':
     71       self._logger.LogOutput('Preparing to download %s image to local '
     72                              'directory.' % build_id)
     73 
     74     # Make sure the directory for downloading the image exists.
     75     download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
     76     image_path = os.path.join(download_path, 'chromiumos_test_image.bin')
     77     if not os.path.exists(download_path):
     78       os.makedirs(download_path)
     79 
     80     # Check to see if the image has already been downloaded.  If not,
     81     # download the image.
     82     if not os.path.exists(image_path):
     83       gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
     84       command = '%s cp %s %s' % (gsutil_cmd, image_name, download_path)
     85 
     86       if self.log_level != 'verbose':
     87         self._logger.LogOutput('CMD: %s' % command)
     88       status = self._ce.RunCommand(command)
     89       downloaded_image_name = os.path.join(download_path,
     90                                            'chromiumos_test_image.tar.xz')
     91       if status != 0 or not os.path.exists(downloaded_image_name):
     92         raise MissingImage('Cannot download image: %s.' % downloaded_image_name)
     93 
     94     return image_path
     95 
     96   def UncompressImage(self, chromeos_root, build_id):
     97     # Check to see if the file has already been uncompresssed, etc.
     98     if os.path.exists(
     99         os.path.join(chromeos_root, 'chroot/tmp', build_id,
    100                      'chromiumos_test_image.bin')):
    101       return
    102 
    103     # Uncompress and untar the downloaded image.
    104     download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
    105     command = ('cd %s ; tar -Jxf chromiumos_test_image.tar.xz ' % download_path)
    106     # Cleanup command for exception handler
    107     clean_cmd = ('cd %s ; rm -f chromiumos_test_image.bin ' % download_path)
    108     exception_handler = RunCommandExceptionHandler(self._logger, self.log_level,
    109                                                    self._ce, clean_cmd)
    110     if self.log_level != 'verbose':
    111       self._logger.LogOutput('CMD: %s' % command)
    112       print('(Uncompressing and un-tarring may take a couple of minutes...'
    113             'please be patient.)')
    114     retval = self._ce.RunCommand(
    115         command, except_handler=exception_handler.HandleException)
    116     if retval != 0:
    117       if self.log_level != 'verbose':
    118         self._logger.LogOutput('CMD: %s' % clean_cmd)
    119         print('(Removing file chromiumos_test_image.bin.)')
    120       # Remove partially uncompressed file
    121       _ = self._ce.RunCommand(clean_cmd)
    122       # Raise exception for failure to uncompress
    123       raise MissingImage('Cannot uncompress image: %s.' % build_id)
    124 
    125     # Remove compressed image
    126     command = ('cd %s ; rm -f chromiumos_test_image.tar.xz; ' % download_path)
    127     if self.log_level != 'verbose':
    128       self._logger.LogOutput('CMD: %s' % command)
    129       print('(Removing file chromiumos_test_image.tar.xz.)')
    130     # try removing file, its ok to have an error, print if encountered
    131     retval = self._ce.RunCommand(command)
    132     if retval != 0:
    133       print('(Warning: Could not remove file chromiumos_test_image.tar.xz .)')
    134 
    135   def DownloadSingleAutotestFile(self, chromeos_root, build_id,
    136                                  package_file_name):
    137     # Verify if package files exist
    138     status = 0
    139     gs_package_name = ('gs://chromeos-image-archive/%s/%s' %
    140                        (build_id, package_file_name))
    141     gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
    142     if not test_flag.GetTestMode():
    143       cmd = '%s ls %s' % (gsutil_cmd, gs_package_name)
    144       status = self._ce.RunCommand(cmd)
    145     if status != 0:
    146       raise MissingFile('Cannot find autotest package file: %s.' %
    147                         package_file_name)
    148 
    149     if self.log_level == 'average':
    150       self._logger.LogOutput('Preparing to download %s package to local '
    151                              'directory.' % package_file_name)
    152 
    153     # Make sure the directory for downloading the package exists.
    154     download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
    155     package_path = os.path.join(download_path, package_file_name)
    156     if not os.path.exists(download_path):
    157       os.makedirs(download_path)
    158 
    159     # Check to see if the package file has already been downloaded.  If not,
    160     # download it.
    161     if not os.path.exists(package_path):
    162       command = '%s cp %s %s' % (gsutil_cmd, gs_package_name, download_path)
    163 
    164       if self.log_level != 'verbose':
    165         self._logger.LogOutput('CMD: %s' % command)
    166       status = self._ce.RunCommand(command)
    167       if status != 0 or not os.path.exists(package_path):
    168         raise MissingFile('Cannot download package: %s .' % package_path)
    169 
    170   def UncompressSingleAutotestFile(self, chromeos_root, build_id,
    171                                    package_file_name, uncompress_cmd):
    172     # Uncompress file
    173     download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
    174     command = ('cd %s ; %s %s' %
    175                (download_path, uncompress_cmd, package_file_name))
    176 
    177     if self.log_level != 'verbose':
    178       self._logger.LogOutput('CMD: %s' % command)
    179       print('(Uncompressing autotest file %s .)' % package_file_name)
    180     retval = self._ce.RunCommand(command)
    181     if retval != 0:
    182       raise MissingFile('Cannot uncompress file: %s.' % package_file_name)
    183     # Remove uncompressed downloaded file
    184     command = ('cd %s ; rm -f %s' % (download_path, package_file_name))
    185     if self.log_level != 'verbose':
    186       self._logger.LogOutput('CMD: %s' % command)
    187       print('(Removing processed autotest file %s .)' % package_file_name)
    188     # try removing file, its ok to have an error, print if encountered
    189     retval = self._ce.RunCommand(command)
    190     if retval != 0:
    191       print('(Warning: Could not remove file %s .)' % package_file_name)
    192 
    193   def VerifyAutotestFilesExist(self, chromeos_root, build_id, package_file):
    194     # Quickly verify if the files are there
    195     status = 0
    196     gs_package_name = ('gs://chromeos-image-archive/%s/%s' %
    197                        (build_id, package_file))
    198     gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
    199     if not test_flag.GetTestMode():
    200       cmd = '%s ls %s' % (gsutil_cmd, gs_package_name)
    201       if self.log_level != 'verbose':
    202         self._logger.LogOutput('CMD: %s' % cmd)
    203       status = self._ce.RunCommand(cmd)
    204       if status != 0:
    205         print('(Warning: Could not find file %s )' % gs_package_name)
    206         return 1
    207     # Package exists on server
    208     return 0
    209 
    210   def DownloadAutotestFiles(self, chromeos_root, build_id):
    211     # Download autest package files (3 files)
    212     autotest_packages_name = ('autotest_packages.tar')
    213     autotest_server_package_name = ('autotest_server_package.tar.bz2')
    214     autotest_control_files_name = ('control_files.tar')
    215 
    216     download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
    217     # Autotest directory relative path wrt chroot
    218     autotest_rel_path = os.path.join('/tmp', build_id, 'autotest_files')
    219     # Absolute Path to download files
    220     autotest_path = os.path.join(chromeos_root, 'chroot/tmp', build_id,
    221                                  'autotest_files')
    222 
    223     if not os.path.exists(autotest_path):
    224       # Quickly verify if the files are present on server
    225       # If not, just exit with warning
    226       status = self.VerifyAutotestFilesExist(chromeos_root, build_id,
    227                                              autotest_packages_name)
    228       if status != 0:
    229         default_autotest_dir = '~/trunk/src/third_party/autotest/files'
    230         print('(Warning: Could not find autotest packages .)\n'
    231               '(Warning: Defaulting autotest path to %s .' %
    232               default_autotest_dir)
    233         return default_autotest_dir
    234 
    235       # Files exist on server, download and uncompress them
    236       self.DownloadSingleAutotestFile(chromeos_root, build_id,
    237                                       autotest_packages_name)
    238       self.DownloadSingleAutotestFile(chromeos_root, build_id,
    239                                       autotest_server_package_name)
    240       self.DownloadSingleAutotestFile(chromeos_root, build_id,
    241                                       autotest_control_files_name)
    242 
    243       self.UncompressSingleAutotestFile(chromeos_root, build_id,
    244                                         autotest_packages_name, 'tar -xvf ')
    245       self.UncompressSingleAutotestFile(chromeos_root, build_id,
    246                                         autotest_server_package_name,
    247                                         'tar -jxvf ')
    248       self.UncompressSingleAutotestFile(chromeos_root, build_id,
    249                                         autotest_control_files_name,
    250                                         'tar -xvf ')
    251       # Rename created autotest directory to autotest_files
    252       command = ('cd %s ; mv autotest autotest_files' % download_path)
    253       if self.log_level != 'verbose':
    254         self._logger.LogOutput('CMD: %s' % command)
    255         print('(Moving downloaded autotest files to autotest_files)')
    256       retval = self._ce.RunCommand(command)
    257       if retval != 0:
    258         raise MissingFile('Could not create directory autotest_files')
    259 
    260     return autotest_rel_path
    261 
    262   def Run(self, chromeos_root, xbuddy_label, autotest_path):
    263     build_id = self.GetBuildID(chromeos_root, xbuddy_label)
    264     image_name = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz'
    265                   % build_id)
    266 
    267     # Verify that image exists for build_id, before attempting to
    268     # download it.
    269     status = 0
    270     if not test_flag.GetTestMode():
    271       gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
    272       cmd = '%s ls %s' % (gsutil_cmd, image_name)
    273       status = self._ce.RunCommand(cmd)
    274     if status != 0:
    275       raise MissingImage('Cannot find official image: %s.' % image_name)
    276 
    277     image_path = self.DownloadImage(chromeos_root, build_id, image_name)
    278     self.UncompressImage(chromeos_root, build_id)
    279 
    280     if self.log_level != 'quiet':
    281       self._logger.LogOutput('Using image from %s.' % image_path)
    282 
    283     if autotest_path == '':
    284       autotest_path = self.DownloadAutotestFiles(chromeos_root, build_id)
    285 
    286     return image_path, autotest_path
    287