Add ubuntu_drivers config module

The ubuntu_drivers config module enables usage of the 'ubuntu-drivers'
command. At this point it only serves as a way of installing NVIDIA
drivers for general purpose graphics processing unit (GPGPU)
functionality.

Also, a small usability improvement to get_cfg_by_path to allow it to
take a string for the key path
  "toplevel/second/mykey"
in addition to the original:
  ("toplevel", "second", "mykey")

6 files changed, 306 insertions, 0 deletions

diff --git a/cloudinit/config/cc_ubuntu_drivers.py b/cloudinit/config/cc_ubuntu_drivers.py
new file mode 100644
index 00000000..91feb603
--- /dev/null
+++ b/cloudinit/config/cc_ubuntu_drivers.py
@@ -0,0 +1,112 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+"""Ubuntu Drivers: Interact with third party drivers in Ubuntu."""
+
+from textwrap import dedent
+
+from cloudinit.config.schema import (
+    get_schema_doc, validate_cloudconfig_schema)
+from cloudinit import log as logging
+from cloudinit.settings import PER_INSTANCE
+from cloudinit import type_utils
+from cloudinit import util
+
+LOG = logging.getLogger(__name__)
+
+frequency = PER_INSTANCE
+distros = ['ubuntu']
+schema = {
+    'id': 'cc_ubuntu_drivers',
+    'name': 'Ubuntu Drivers',
+    'title': 'Interact with third party drivers in Ubuntu.',
+    'description': dedent("""\
+        This module interacts with the 'ubuntu-drivers' command to install
+        third party driver packages."""),
+    'distros': distros,
+    'examples': [dedent("""\
+        drivers:
+          nvidia:
+            license-accepted: true
+        """)],
+    'frequency': frequency,
+    'type': 'object',
+    'properties': {
+        'drivers': {
+            'type': 'object',
+            'additionalProperties': False,
+            'properties': {
+                'nvidia': {
+                    'type': 'object',
+                    'additionalProperties': False,
+                    'required': ['license-accepted'],
+                    'properties': {
+                        'license-accepted': {
+                            'type': 'boolean',
+                            'description': ("Do you accept the NVIDIA driver"
+                                            " license?"),
+                        },
+                        'version': {
+                            'type': 'string',
+                            'description': (
+                                'The version of the driver to install (e.g.'
+                                ' "390", "410"). Defaults to the latest'
+                                ' version.'),
+                        },
+                    },
+                },
+            },
+        },
+    },
+}
+OLD_UBUNTU_DRIVERS_STDERR_NEEDLE = (
+    "ubuntu-drivers: error: argument <command>: invalid choice: 'install'")
+
+__doc__ = get_schema_doc(schema)  # Supplement python help()
+
+
+def install_drivers(cfg, pkg_install_func):
+    if not isinstance(cfg, dict):
+        raise TypeError(
+            "'drivers' config expected dict, found '%s': %s" %
+            (type_utils.obj_name(cfg), cfg))
+
+    cfgpath = 'nvidia/license-accepted'
+    # Call translate_bool to ensure that we treat string values like "yes" as
+    # acceptance and _don't_ treat string values like "nah" as acceptance
+    # because they're True-ish
+    nv_acc = util.translate_bool(util.get_cfg_by_path(cfg, cfgpath))
+    if not nv_acc:
+        LOG.debug("Not installing NVIDIA drivers. %s=%s", cfgpath, nv_acc)
+        return
+
+    if not util.which('ubuntu-drivers'):
+        LOG.debug("'ubuntu-drivers' command not available.  "
+                  "Installing ubuntu-drivers-common")
+        pkg_install_func(['ubuntu-drivers-common'])
+
+    driver_arg = 'nvidia'
+    version_cfg = util.get_cfg_by_path(cfg, 'nvidia/version')
+    if version_cfg:
+        driver_arg += ':{}'.format(version_cfg)
+
+    LOG.debug("Installing NVIDIA drivers (%s=%s, version=%s)",
+              cfgpath, nv_acc, version_cfg if version_cfg else 'latest')
+
+    try:
+        util.subp(['ubuntu-drivers', 'install', '--gpgpu', driver_arg])
+    except util.ProcessExecutionError as exc:
+        if OLD_UBUNTU_DRIVERS_STDERR_NEEDLE in exc.stderr:
+            LOG.warning('the available version of ubuntu-drivers is'
+                        ' too old to perform requested driver installation')
+        elif 'No drivers found for installation.' in exc.stdout:
+            LOG.warning('ubuntu-drivers found no drivers for installation')
+        raise
+
+
+def handle(name, cfg, cloud, log, _args):
+    if "drivers" not in cfg:
+        log.debug("Skipping module named %s, no 'drivers' key in config", name)
+        return
+
+    validate_cloudconfig_schema(cfg, schema)
+    install_drivers(cfg['drivers'], cloud.distro.install_packages)
diff --git a/cloudinit/config/tests/test_ubuntu_drivers.py b/cloudinit/config/tests/test_ubuntu_drivers.py
new file mode 100644
index 00000000..efba4ce7
--- /dev/null
+++ b/cloudinit/config/tests/test_ubuntu_drivers.py
@@ -0,0 +1,174 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+import copy
+
+from cloudinit.tests.helpers import CiTestCase, skipUnlessJsonSchema, mock
+from cloudinit.config.schema import (
+    SchemaValidationError, validate_cloudconfig_schema)
+from cloudinit.config import cc_ubuntu_drivers as drivers
+from cloudinit.util import ProcessExecutionError
+
+MPATH = "cloudinit.config.cc_ubuntu_drivers."
+OLD_UBUNTU_DRIVERS_ERROR_STDERR = (
+    "ubuntu-drivers: error: argument <command>: invalid choice: 'install' "
+    "(choose from 'list', 'autoinstall', 'devices', 'debug')\n")
+
+
+class TestUbuntuDrivers(CiTestCase):
+    cfg_accepted = {'drivers': {'nvidia': {'license-accepted': True}}}
+    install_gpgpu = ['ubuntu-drivers', 'install', '--gpgpu', 'nvidia']
+
+    with_logs = True
+
+    @skipUnlessJsonSchema()
+    def test_schema_requires_boolean_for_license_accepted(self):
+        with self.assertRaisesRegex(
+                SchemaValidationError, ".*license-accepted.*TRUE.*boolean"):
+            validate_cloudconfig_schema(
+                {'drivers': {'nvidia': {'license-accepted': "TRUE"}}},
+                schema=drivers.schema, strict=True)
+
+    @mock.patch(MPATH + "util.subp", return_value=('', ''))
+    @mock.patch(MPATH + "util.which", return_value=False)
+    def _assert_happy_path_taken(self, config, m_which, m_subp):
+        """Positive path test through handle. Package should be installed."""
+        myCloud = mock.MagicMock()
+        drivers.handle('ubuntu_drivers', config, myCloud, None, None)
+        self.assertEqual([mock.call(['ubuntu-drivers-common'])],
+                         myCloud.distro.install_packages.call_args_list)
+        self.assertEqual([mock.call(self.install_gpgpu)],
+                         m_subp.call_args_list)
+
+    def test_handle_does_package_install(self):
+        self._assert_happy_path_taken(self.cfg_accepted)
+
+    def test_trueish_strings_are_considered_approval(self):
+        for true_value in ['yes', 'true', 'on', '1']:
+            new_config = copy.deepcopy(self.cfg_accepted)
+            new_config['drivers']['nvidia']['license-accepted'] = true_value
+            self._assert_happy_path_taken(new_config)
+
+    @mock.patch(MPATH + "util.subp", side_effect=ProcessExecutionError(
+        stdout='No drivers found for installation.\n', exit_code=1))
+    @mock.patch(MPATH + "util.which", return_value=False)
+    def test_handle_raises_error_if_no_drivers_found(self, m_which, m_subp):
+        """If ubuntu-drivers doesn't install any drivers, raise an error."""
+        myCloud = mock.MagicMock()
+        with self.assertRaises(Exception):
+            drivers.handle(
+                'ubuntu_drivers', self.cfg_accepted, myCloud, None, None)
+        self.assertEqual([mock.call(['ubuntu-drivers-common'])],
+                         myCloud.distro.install_packages.call_args_list)
+        self.assertEqual([mock.call(self.install_gpgpu)],
+                         m_subp.call_args_list)
+        self.assertIn('ubuntu-drivers found no drivers for installation',
+                      self.logs.getvalue())
+
+    @mock.patch(MPATH + "util.subp", return_value=('', ''))
+    @mock.patch(MPATH + "util.which", return_value=False)
+    def _assert_inert_with_config(self, config, m_which, m_subp):
+        """Helper to reduce repetition when testing negative cases"""
+        myCloud = mock.MagicMock()
+        drivers.handle('ubuntu_drivers', config, myCloud, None, None)
+        self.assertEqual(0, myCloud.distro.install_packages.call_count)
+        self.assertEqual(0, m_subp.call_count)
+
+    def test_handle_inert_if_license_not_accepted(self):
+        """Ensure we don't do anything if the license is rejected."""
+        self._assert_inert_with_config(
+            {'drivers': {'nvidia': {'license-accepted': False}}})
+
+    def test_handle_inert_if_garbage_in_license_field(self):
+        """Ensure we don't do anything if unknown text is in license field."""
+        self._assert_inert_with_config(
+            {'drivers': {'nvidia': {'license-accepted': 'garbage'}}})
+
+    def test_handle_inert_if_no_license_key(self):
+        """Ensure we don't do anything if no license key."""
+        self._assert_inert_with_config({'drivers': {'nvidia': {}}})
+
+    def test_handle_inert_if_no_nvidia_key(self):
+        """Ensure we don't do anything if other license accepted."""
+        self._assert_inert_with_config(
+            {'drivers': {'acme': {'license-accepted': True}}})
+
+    def test_handle_inert_if_string_given(self):
+        """Ensure we don't do anything if string refusal given."""
+        for false_value in ['no', 'false', 'off', '0']:
+            self._assert_inert_with_config(
+                {'drivers': {'nvidia': {'license-accepted': false_value}}})
+
+    @mock.patch(MPATH + "install_drivers")
+    def test_handle_no_drivers_does_nothing(self, m_install_drivers):
+        """If no 'drivers' key in the config, nothing should be done."""
+        myCloud = mock.MagicMock()
+        myLog = mock.MagicMock()
+        drivers.handle('ubuntu_drivers', {'foo': 'bzr'}, myCloud, myLog, None)
+        self.assertIn('Skipping module named',
+                      myLog.debug.call_args_list[0][0][0])
+        self.assertEqual(0, m_install_drivers.call_count)
+
+    @mock.patch(MPATH + "util.subp", return_value=('', ''))
+    @mock.patch(MPATH + "util.which", return_value=True)
+    def test_install_drivers_no_install_if_present(self, m_which, m_subp):
+        """If 'ubuntu-drivers' is present, no package install should occur."""
+        pkg_install = mock.MagicMock()
+        drivers.install_drivers(self.cfg_accepted['drivers'],
+                                pkg_install_func=pkg_install)
+        self.assertEqual(0, pkg_install.call_count)
+        self.assertEqual([mock.call('ubuntu-drivers')],
+                         m_which.call_args_list)
+        self.assertEqual([mock.call(self.install_gpgpu)],
+                         m_subp.call_args_list)
+
+    def test_install_drivers_rejects_invalid_config(self):
+        """install_drivers should raise TypeError if not given a config dict"""
+        pkg_install = mock.MagicMock()
+        with self.assertRaisesRegex(TypeError, ".*expected dict.*"):
+            drivers.install_drivers("mystring", pkg_install_func=pkg_install)
+        self.assertEqual(0, pkg_install.call_count)
+
+    @mock.patch(MPATH + "util.subp", side_effect=ProcessExecutionError(
+        stderr=OLD_UBUNTU_DRIVERS_ERROR_STDERR, exit_code=2))
+    @mock.patch(MPATH + "util.which", return_value=False)
+    def test_install_drivers_handles_old_ubuntu_drivers_gracefully(
+            self, m_which, m_subp):
+        """Older ubuntu-drivers versions should emit message and raise error"""
+        myCloud = mock.MagicMock()
+        with self.assertRaises(Exception):
+            drivers.handle(
+                'ubuntu_drivers', self.cfg_accepted, myCloud, None, None)
+        self.assertEqual([mock.call(['ubuntu-drivers-common'])],
+                         myCloud.distro.install_packages.call_args_list)
+        self.assertEqual([mock.call(self.install_gpgpu)],
+                         m_subp.call_args_list)
+        self.assertIn('WARNING: the available version of ubuntu-drivers is'
+                      ' too old to perform requested driver installation',
+                      self.logs.getvalue())
+
+
+# Sub-class TestUbuntuDrivers to run the same test cases, but with a version
+class TestUbuntuDriversWithVersion(TestUbuntuDrivers):
+    cfg_accepted = {
+        'drivers': {'nvidia': {'license-accepted': True, 'version': '123'}}}
+    install_gpgpu = ['ubuntu-drivers', 'install', '--gpgpu', 'nvidia:123']
+
+    @mock.patch(MPATH + "util.subp", return_value=('', ''))
+    @mock.patch(MPATH + "util.which", return_value=False)
+    def test_version_none_uses_latest(self, m_which, m_subp):
+        myCloud = mock.MagicMock()
+        version_none_cfg = {
+            'drivers': {'nvidia': {'license-accepted': True, 'version': None}}}
+        drivers.handle(
+            'ubuntu_drivers', version_none_cfg, myCloud, None, None)
+        self.assertEqual(
+            [mock.call(['ubuntu-drivers', 'install', '--gpgpu', 'nvidia'])],
+            m_subp.call_args_list)
+
+    def test_specifying_a_version_doesnt_override_license_acceptance(self):
+        self._assert_inert_with_config({
+            'drivers': {'nvidia': {'license-accepted': False,
+                                   'version': '123'}}
+        })
+
+# vi: ts=4 expandtab
diff --git a/cloudinit/util.py b/cloudinit/util.py
index a192091f..385f231c 100644
--- a/cloudinit/util.py
+++ b/cloudinit/util.py
@@ -703,6 +703,21 @@ def get_cfg_option_list(yobj, key, default=None):
 # get a cfg entry by its path array
 # for f['a']['b']: get_cfg_by_path(mycfg,('a','b'))
 def get_cfg_by_path(yobj, keyp, default=None):
+    """Return the value of the item at path C{keyp} in C{yobj}.
+
+    example:
+      get_cfg_by_path({'a': {'b': {'num': 4}}}, 'a/b/num') == 4
+      get_cfg_by_path({'a': {'b': {'num': 4}}}, 'c/d') == None
+
+    @param yobj: A dictionary.
+    @param keyp: A path inside yobj.  it can be a '/' delimited string,
+                 or an iterable.
+    @param default: The default to return if the path does not exist.
+    @return: The value of the item at keyp."
+        is not found."""
+
+    if isinstance(keyp, six.string_types):
+        keyp = keyp.split("/")
     cur = yobj
     for tok in keyp:
         if tok not in cur:
diff --git a/config/cloud.cfg.tmpl b/config/cloud.cfg.tmpl
index 7513176b..25db43e0 100644
--- a/config/cloud.cfg.tmpl
+++ b/config/cloud.cfg.tmpl
@@ -112,6 +112,9 @@ cloud_final_modules:
  - landscape
  - lxd
 {% endif %}
+{% if variant in ["ubuntu", "unknown"] %}
+ - ubuntu-drivers
+{% endif %}
 {% if variant not in ["freebsd"] %}
  - puppet
  - chef
diff --git a/doc/rtd/topics/modules.rst b/doc/rtd/topics/modules.rst
index d9720f6a..3dcdd3bc 100644
--- a/doc/rtd/topics/modules.rst
+++ b/doc/rtd/topics/modules.rst
@@ -54,6 +54,7 @@ Modules
 .. automodule:: cloudinit.config.cc_ssh_import_id
 .. automodule:: cloudinit.config.cc_timezone
 .. automodule:: cloudinit.config.cc_ubuntu_advantage
+.. automodule:: cloudinit.config.cc_ubuntu_drivers
 .. automodule:: cloudinit.config.cc_update_etc_hosts
 .. automodule:: cloudinit.config.cc_update_hostname
 .. automodule:: cloudinit.config.cc_users_groups
diff --git a/tests/unittests/test_handler/test_schema.py b/tests/unittests/test_handler/test_schema.py
index 1bad07f6..e69a47a9 100644
--- a/tests/unittests/test_handler/test_schema.py
+++ b/tests/unittests/test_handler/test_schema.py
@@ -28,6 +28,7 @@ class GetSchemaTest(CiTestCase):
                 'cc_runcmd',
                 'cc_snap',
                 'cc_ubuntu_advantage',
+                'cc_ubuntu_drivers',
                 'cc_zypper_add_repo'
             ],
             [subschema['id'] for subschema in schema['allOf']])