1 files changed, 112 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)