puppylab/vm_manager.py

#!/usr/bin/python
# Copyright (c) 2014 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.

"""Manage vms through vagrant.

The intent of this interface is to provde a layer of abstraction
between the box providers and the creation of a lab cluster. To switch to a
different provider:

* Create a VagrantFile template and specify _template in the subclass
  Eg: GCE VagrantFiles need a :google section
* Override vagrant_cmd to massage parameters
  Eg: vagrant up => vagrant up --provider=google

Note that the second is optional because most providers honor
`VAGRANT_DEFAULT_PROVIDER` directly in the template.
"""


import logging
import subprocess
import sys
import os

import common
from autotest_lib.site_utils.lib import infra


class VagrantCmdError(Exception):
    """Raised when a vagrant command fails."""


# TODO: We don't really need to setup everythig in the same VAGRANT_DIR.
# However managing vms becomes a headache once the VagrantFile and its
# related dot files are removed, as one has to resort to directly
# querying the box provider. Always running the cluster from the same
# directory simplifies vm lifecycle management.
VAGRANT_DIR = os.path.abspath(os.path.join(__file__, os.pardir))

class VagrantProvisioner(object):
    """Provisiong vms with vagrant."""

    # A path to a Vagrantfile template specific to the vm provider, specified
    # in the child class.
    _template = None
    _box_name = 'base'


    @classmethod
    def vagrant_cmd(cls, cmd, stream_output=False):
        """Execute a vagrant command in VAGRANT_DIR.

        @param cmd: The command to execute.
        @param stream_output: If True, stream the output of `cmd`.
                Waits for `cmd` to finish and returns a string with the
                output if false.
        """
        with infra.chdir(VAGRANT_DIR):
            try:
                return infra.execute_command(
                        'localhost',
                        'vagrant %s' % cmd, stream_output=stream_output)
            except subprocess.CalledProcessError as e:
                raise VagrantCmdError(
                        'Command "vagrant %s" failed with %s' % (cmd, e))


    def _check_vagrant(self):
        """Check Vagrant."""

        # TODO: Automate the installation of vagrant.
        try:
            self.vagrant_cmd('--version')
        except VagrantCmdError:
            logging.error(
                    'Looks like you don\'t have vagrant. Please run: \n'
                    '`apt-get install virtualbox vagrant`. This assumes you '
                    'are on Trusty; There is a TODO to automate installation.')
            sys.exit(1)


    def __init__(self, puppet_path):
        """Initialize a vagrant provisioner.

        @param puppet_path: Since vagrant uses puppet to provision machines,
                this is the location of puppet modules for various server roles.
        """
        self._check_vagrant()
        self.puppet_path = puppet_path


    def register_box(self, source, name=_box_name):
        """Register a box with vagrant.

        Eg: vagrant box add core_cluster chromeos_lab_core_cluster.box

        @param source: A path to the box, typically a file path on localhost.
        @param name: A name to register the box under.
        """
        if name in self.vagrant_cmd('box list'):
            logging.warning("Name %s already in registry, will reuse.", name)
            return
        logging.info('Adding a new box from %s under name: %s', source, name)
        self.vagrant_cmd('box add %s %s' % (name, source))


    def unregister_box(self, name):
        """Unregister a box.

        Eg: vagrant box remove core_cluster.

        @param name: The name of the box as it appears in `vagrant box list`
        """
        if name not in self.vagrant_cmd('box list'):
            logging.warning("Name %s not in registry.", name)
            return
        logging.info('Removing box %s', name)
        self.vagrant_cmd('box remove %s' % name)


    def create_vagrant_file(self, **kwargs):
        """Create a vagrant file.

        Read the template, apply kwargs and the puppet_path so vagrant can find
        server provisioning rules, and write it back out as the VagrantFile.

        @param kwargs: Extra args needed to convert a template
                to a real VagrantFile.
        """
        vagrant_file = os.path.join(VAGRANT_DIR, 'Vagrantfile')
        kwargs.update({
            'manifest_path': os.path.join(self.puppet_path, 'manifests'),
            'module_path': os.path.join(self.puppet_path, 'modules'),
        })
        vagrant_template = ''
        with open(self._template, 'r') as template:
            vagrant_template = template.read()
        with open(vagrant_file, 'w') as vagrantfile:
            vagrantfile.write(vagrant_template % kwargs)


    # TODO: This is a leaky abstraction, since it isn't really clear
    # what the kwargs are. It's the best we can do, because the kwargs
    # really need to match the VagrantFile. We leave parsing the VagrantFile
    # for the right args upto the caller.
    def initialize_vagrant(self, **kwargs):
        """Initialize vagrant.

        @param kwargs: The kwargs to pass to the VagrantFile.
            Eg: {
                'shard1': 'lumpyshard',
                'shard1_port': 8002,
                'shard1_shadow_config_hostname': 'localhost:8002',
            }
        @return: True if vagrant was initialized, False if the cwd already
                 contains a vagrant environment.
        """
        # TODO: Split this out. There are cases where we will need to
        # reinitialize (by destroying all vms and recreating the VagrantFile)
        # that we cannot do without manual intervention right now.
        try:
            self.vagrant_cmd('status')
            logging.info('Vagrant already initialized in %s', VAGRANT_DIR)
            return False
        except VagrantCmdError:
            logging.info('Initializing vagrant in %s', VAGRANT_DIR)
            self.create_vagrant_file(**kwargs)
            return True


    def provision(self, force=False):
        """Provision vms according to the vagrant file.

        @param force: If True, vms in the VAGRANT_DIR will be destroyed and
                reprovisioned.
        """
        if force:
            logging.info('Destroying vagrant setup.')
            try:
                self.vagrant_cmd('destroy --force', stream_output=True)
            except VagrantCmdError:
                pass
        logging.info('Starting vms. This should take no longer than 20 minutes '
                     'the first time, and no longer than 5 subsequently.')
        self.vagrant_cmd('up', stream_output=True)


class VirtualBox(VagrantProvisioner):
    """A VirtualBoxProvisioner."""

    _template = os.path.join(VAGRANT_DIR, 'ClusterTemplate')