1#!/usr/bin/python 2# Copyright (c) 2014 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6"""Manage vms through vagrant. 7 8The intent of this interface is to provde a layer of abstraction 9between the box providers and the creation of a lab cluster. To switch to a 10different provider: 11 12* Create a VagrantFile template and specify _template in the subclass 13 Eg: GCE VagrantFiles need a :google section 14* Override vagrant_cmd to massage parameters 15 Eg: vagrant up => vagrant up --provider=google 16 17Note that the second is optional because most providers honor 18`VAGRANT_DEFAULT_PROVIDER` directly in the template. 19""" 20 21 22import logging 23import subprocess 24import sys 25import os 26 27import common 28from autotest_lib.site_utils.lib import infra 29 30 31class VagrantCmdError(Exception): 32 """Raised when a vagrant command fails.""" 33 34 35# TODO: We don't really need to setup everythig in the same VAGRANT_DIR. 36# However managing vms becomes a headache once the VagrantFile and its 37# related dot files are removed, as one has to resort to directly 38# querying the box provider. Always running the cluster from the same 39# directory simplifies vm lifecycle management. 40VAGRANT_DIR = os.path.abspath(os.path.join(__file__, os.pardir)) 41VAGRANT_VERSION = '1.6.0' 42 43 44def format_msg(msg): 45 """Format the give message. 46 47 @param msg: A message to format out to stdout. 48 """ 49 print '\n{:^20s}%s'.format('') % msg 50 51 52class VagrantProvisioner(object): 53 """Provisiong vms with vagrant.""" 54 55 # A path to a Vagrantfile template specific to the vm provider, specified 56 # in the child class. 57 _template = None 58 _box_name = 'base' 59 60 61 @classmethod 62 def vagrant_cmd(cls, cmd, stream_output=False): 63 """Execute a vagrant command in VAGRANT_DIR. 64 65 @param cmd: The command to execute. 66 @param stream_output: If True, stream the output of `cmd`. 67 Waits for `cmd` to finish and returns a string with the 68 output if false. 69 """ 70 with infra.chdir(VAGRANT_DIR): 71 try: 72 return infra.execute_command( 73 'localhost', 74 'vagrant %s' % cmd, stream_output=stream_output) 75 except subprocess.CalledProcessError as e: 76 raise VagrantCmdError( 77 'Command "vagrant %s" failed with %s' % (cmd, e)) 78 79 80 def _check_vagrant(self): 81 """Check Vagrant.""" 82 83 # TODO: Automate the installation of vagrant. 84 try: 85 version = int(self.vagrant_cmd('--version').rstrip('\n').rsplit( 86 ' ')[-1].replace('.', '')) 87 except VagrantCmdError: 88 logging.error( 89 'Looks like you don\'t have vagrant. Please run: \n' 90 '`apt-get install virtualbox vagrant`. This assumes you ' 91 'are on Trusty; There is a TODO to automate installation.') 92 sys.exit(1) 93 except TypeError as e: 94 logging.warning('The format of the vagrant version string seems to ' 95 'have changed, assuming you have a version > %s.', 96 VAGRANT_VERSION) 97 return 98 if version < int(VAGRANT_VERSION.replace('.', '')): 99 logging.error('Please upgrade vagrant to a version > %s by ' 100 'downloading a deb file from ' 101 'https://www.vagrantup.com/downloads and installing ' 102 'it with dpkg -i file.deb', VAGRANT_VERSION) 103 sys.exit(1) 104 105 106 def __init__(self, puppet_path): 107 """Initialize a vagrant provisioner. 108 109 @param puppet_path: Since vagrant uses puppet to provision machines, 110 this is the location of puppet modules for various server roles. 111 """ 112 self._check_vagrant() 113 self.puppet_path = puppet_path 114 115 116 def register_box(self, source, name=_box_name): 117 """Register a box with vagrant. 118 119 Eg: vagrant box add core_cluster chromeos_lab_core_cluster.box 120 121 @param source: A path to the box, typically a file path on localhost. 122 @param name: A name to register the box under. 123 """ 124 if name in self.vagrant_cmd('box list'): 125 logging.warning("Name %s already in registry, will reuse.", name) 126 return 127 logging.info('Adding a new box from %s under name: %s', source, name) 128 self.vagrant_cmd('box add %s %s' % (name, source)) 129 130 131 def unregister_box(self, name): 132 """Unregister a box. 133 134 Eg: vagrant box remove core_cluster. 135 136 @param name: The name of the box as it appears in `vagrant box list` 137 """ 138 if name not in self.vagrant_cmd('box list'): 139 logging.warning("Name %s not in registry.", name) 140 return 141 logging.info('Removing box %s', name) 142 self.vagrant_cmd('box remove %s' % name) 143 144 145 def create_vagrant_file(self, **kwargs): 146 """Create a vagrant file. 147 148 Read the template, apply kwargs and the puppet_path so vagrant can find 149 server provisioning rules, and write it back out as the VagrantFile. 150 151 @param kwargs: Extra args needed to convert a template 152 to a real VagrantFile. 153 """ 154 vagrant_file = os.path.join(VAGRANT_DIR, 'Vagrantfile') 155 kwargs.update({ 156 'manifest_path': os.path.join(self.puppet_path, 'manifests'), 157 'module_path': os.path.join(self.puppet_path, 'modules'), 158 }) 159 vagrant_template = '' 160 with open(self._template, 'r') as template: 161 vagrant_template = template.read() 162 with open(vagrant_file, 'w') as vagrantfile: 163 vagrantfile.write(vagrant_template % kwargs) 164 165 166 # TODO: This is a leaky abstraction, since it isn't really clear 167 # what the kwargs are. It's the best we can do, because the kwargs 168 # really need to match the VagrantFile. We leave parsing the VagrantFile 169 # for the right args upto the caller. 170 def initialize_vagrant(self, **kwargs): 171 """Initialize vagrant. 172 173 @param kwargs: The kwargs to pass to the VagrantFile. 174 Eg: { 175 'shard1': 'stumpyshard', 176 'shard1_port': 8002, 177 'shard1_shadow_config_hostname': 'localhost:8002', 178 } 179 @return: True if vagrant was initialized, False if the cwd already 180 contains a vagrant environment. 181 """ 182 # TODO: Split this out. There are cases where we will need to 183 # reinitialize (by destroying all vms and recreating the VagrantFile) 184 # that we cannot do without manual intervention right now. 185 try: 186 self.vagrant_cmd('status') 187 logging.info('Vagrant already initialized in %s', VAGRANT_DIR) 188 return False 189 except VagrantCmdError: 190 logging.info('Initializing vagrant in %s', VAGRANT_DIR) 191 self.create_vagrant_file(**kwargs) 192 return True 193 194 195 def provision(self, force=False): 196 """Provision vms according to the vagrant file. 197 198 @param force: If True, vms in the VAGRANT_DIR will be destroyed and 199 reprovisioned. 200 """ 201 if force: 202 logging.info('Destroying vagrant setup.') 203 try: 204 self.vagrant_cmd('destroy --force', stream_output=True) 205 except VagrantCmdError: 206 pass 207 format_msg('Starting vms. This should take no longer than 5 minutes') 208 self.vagrant_cmd('up', stream_output=True) 209 210 211class VirtualBox(VagrantProvisioner): 212 """A VirtualBoxProvisioner.""" 213 214 _template = os.path.join(VAGRANT_DIR, 'ClusterTemplate') 215