diff --git a/subiquity/server/controllers/drivers.py b/subiquity/server/controllers/drivers.py index d3ff956d..4da69cee 100644 --- a/subiquity/server/controllers/drivers.py +++ b/subiquity/server/controllers/drivers.py @@ -15,7 +15,6 @@ import asyncio import logging -import subprocess from typing import List, Optional from subiquitycore.context import with_context @@ -23,8 +22,12 @@ from subiquitycore.context import with_context from subiquity.common.apidef import API from subiquity.common.types import DriversResponse from subiquity.server.controller import SubiquityController -from subiquity.server.curtin import run_curtin_command from subiquity.server.types import InstallerChannels +from subiquity.server.ubuntu_drivers import ( + CommandNotFoundError, + UbuntuDriversInterface, + get_ubuntu_drivers_interface, + ) log = logging.getLogger('subiquity.server.controllers.drivers') @@ -46,6 +49,11 @@ class DriversController(SubiquityController): drivers: Optional[List[str]] = None + def __init__(self, app) -> None: + super().__init__(app) + self.ubuntu_drivers: UbuntuDriversInterface = \ + get_ubuntu_drivers_interface(app) + def make_autoinstall(self): return { "install": self.model.do_install, @@ -67,49 +75,17 @@ class DriversController(SubiquityController): with context.child("wait_apt"): await self._wait_apt.wait() apt = self.app.controllers.Mirror.apt_configurer - # TODO make sure --recommended is a supported option - cmd = ['ubuntu-drivers', 'list', '--recommended'] - server: bool = self.app.base_model.source.current.variant == "server" - if server: - cmd.append('--gpgpu') - if self.app.opts.dry_run: - if 'has-drivers' in self.app.debug_flags: - if server: - self.drivers = ["nvidia-driver-470-server"] - else: - self.drivers = ["nvidia-driver-510"] - return - elif 'run-drivers' in self.app.debug_flags: - pass - else: - self.drivers = [] - await self.configured() - return async with apt.overlay() as d: try: - await self.app.command_runner.run( - ['chroot', d.mountpoint, - 'sh', '-c', - "command -v ubuntu-drivers"]) - except subprocess.CalledProcessError: + # Make sure ubuntu-drivers is available. + self.ubuntu_drivers.ensure_cmd_exists(d.mountpoint) + except CommandNotFoundError: self.drivers = [] - await self.configured() - return - result = await run_curtin_command( - self.app, context, "in-target", "-t", d.mountpoint, - "--", *cmd, capture=True) - # Drivers are listed one per line, but each is followed by a - # linux-modules-* package (which we are not interested in) ; e.g.: - # $ ubuntu-drivers list --recommended - # nvidia-driver-470 linux-modules-nvidia-470-generic-hwe-20.04 - self.drivers = [] - # Currently we have no way to specify universal_newlines=True or - # encoding="utf-8" to run_curtin_command. - stdout = result.stdout.decode("utf-8") - for line in [x.strip() for x in stdout.split("\n")]: - if not line: - continue - self.drivers.append(line.split(" ", maxsplit=1)[0]) + else: + self.drivers = await self.ubuntu_drivers.list_drivers( + root_dir=d.mountpoint, + context=context) + log.debug("Available drivers to install: %s", self.drivers) if not self.drivers: await self.configured() diff --git a/subiquity/server/controllers/install.py b/subiquity/server/controllers/install.py index 387a5c10..46ed4463 100644 --- a/subiquity/server/controllers/install.py +++ b/subiquity/server/controllers/install.py @@ -200,12 +200,9 @@ class InstallController(SubiquityController): with context.child( "ubuntu-drivers-install", "installing third-party drivers") as child: - cmd = ["ubuntu-drivers", "install"] - if self.model.source.current.variant == 'server': - cmd.append('--gpgpu') - await run_curtin_command( - self.app, child, "in-target", "-t", self.tpath(), - "--", *cmd) + ubuntu_drivers = self.app.controllers.Drivers.ubuntu_drivers + await ubuntu_drivers.install_drivers(root_dir=self.tpath(), + context=child) if self.model.network.has_network: self.app.update_state(ApplicationState.UU_RUNNING) diff --git a/subiquity/server/runner.py b/subiquity/server/runner.py index e75b7c0a..f8e9f821 100644 --- a/subiquity/server/runner.py +++ b/subiquity/server/runner.py @@ -57,12 +57,6 @@ class DryRunCommandRunner(LoggedCommandRunner): async def start(self, cmd, *, capture=False): if 'scripts/replay-curtin-log.py' in cmd: delay = 0 - elif cmd[-4:] == ['ubuntu-drivers', 'list', '--recommended', '--gpgpu']: - cmd = cmd[-4:] - delay = 0 - elif cmd[-3:] == ['ubuntu-drivers', 'list', '--recommended']: - cmd = cmd[-3:] - delay = 0 else: cmd = ['echo', 'not running:'] + cmd if 'unattended-upgrades' in cmd: diff --git a/subiquity/server/ubuntu_drivers.py b/subiquity/server/ubuntu_drivers.py new file mode 100644 index 00000000..aa5728bb --- /dev/null +++ b/subiquity/server/ubuntu_drivers.py @@ -0,0 +1,162 @@ +# Copyright 2022 Canonical, Ltd. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +""" Module that defines helpers to use the ubuntu-drivers command. """ + +from abc import ABC, abstractmethod +import logging +import subprocess +from typing import List, Type + +from subiquity.server.curtin import run_curtin_command +from subiquitycore.utils import arun_command + + +log = logging.getLogger("subiquity.server.ubuntu_drivers") + + +class CommandNotFoundError(Exception): + """ Exception to be raised when the ubuntu-drivers command is not + available. + """ + + +class UbuntuDriversInterface(ABC): + def __init__(self, app, gpgpu: bool) -> None: + self.app = app + + self.list_drivers_cmd = [ + "ubuntu-drivers", "list", + "--recommended", + ] + self.install_drivers_cmd = [ + "ubuntu-drivers", "install", + ] + if gpgpu: + self.list_drivers_cmd.append("--gpgpu") + self.install_drivers_cmd.append("--gpgpu") + + @abstractmethod + async def ensure_cmd_exists(self, root_dir: str) -> None: + pass + + @abstractmethod + async def list_drivers(self, root_dir: str, context) -> List[str]: + pass + + async def install_drivers(self, root_dir: str, context) -> None: + await run_curtin_command( + self.app, context, + "in-target", "-t", root_dir, "--", *self.install_drivers_cmd) + + def _drivers_from_output(self, output: str) -> List[str]: + """ Parse the output of ubuntu-drivers list --recommended and return a + list of drivers. """ + drivers: List[str] = [] + # Drivers are listed one per line, but each driver is followed by a + # linux-modules-* package (which we are not interested in showing). + # e.g.,: + # $ ubuntu-drivers list --recommended + # nvidia-driver-470 linux-modules-nvidia-470-generic-hwe-20.04 + for line in [x.strip() for x in output.split("\n")]: + if not line: + continue + drivers.append(line.split(" ", maxsplit=1)[0]) + + return drivers + + +class UbuntuDriversClientInterface(UbuntuDriversInterface): + """ UbuntuDrivers interface that uses the ubuntu-drivers command from the + specified root directory. """ + + async def ensure_cmd_exists(self, root_dir: str) -> None: + # TODO This does not tell us if the "--recommended" option is + # available. + try: + await self.app.command_runner.run( + ['chroot', root_dir, + 'sh', '-c', + "command -v ubuntu-drivers"]) + except subprocess.CalledProcessError: + raise CommandNotFoundError( + f"Command ubuntu-drivers is not available in {root_dir}") + + async def list_drivers(self, root_dir: str, context) -> List[str]: + result = await run_curtin_command( + self.app, context, + "in-target", "-t", root_dir, "--", *self.list_drivers_cmd, + capture=True) + # Currently we have no way to specify universal_newlines=True or + # encoding="utf-8" to run_curtin_command so we need to decode the + # output. + return self._drivers_from_output(result.stdout.decode("utf-8")) + + +class UbuntuDriversHasDriversInterface(UbuntuDriversInterface): + """ A dry-run implementation of ubuntu-drivers that returns a hard-coded + list of drivers. """ + gpgpu_drivers: List[str] = ["nvidia-driver-470-server"] + not_gpgpu_drivers: List[str] = ["nvidia-driver-510"] + + def __init__(self, app, gpgpu: bool) -> None: + super().__init__(app, gpgpu) + self.drivers = self.gpgpu_drivers if gpgpu else self.not_gpgpu_drivers + + async def ensure_cmd_exists(self, root_dir: str) -> None: + pass + + async def list_drivers(self, root_dir: str, context) -> List[str]: + return self.drivers + + +class UbuntuDriversNoDriversInterface(UbuntuDriversHasDriversInterface): + """ A dry-run implementation of ubuntu-drivers that returns a hard-coded + empty list of drivers. """ + + gpgpu_drivers: List[str] = [] + not_gpgpu_drivers: List[str] = [] + + +class UbuntuDriversRunDriversInterface(UbuntuDriversInterface): + """ A dry-run implementation of ubuntu-drivers that actually runs the + ubuntu-drivers command but locally. """ + async def ensure_cmd_exists(self, root_dir: str) -> None: + # TODO This does not tell us if the "--recommended" option is + # available. + try: + await arun_command(["command", "-v", "ubuntu-drivers"]) + except subprocess.CalledProcessError: + raise CommandNotFoundError( + "Command ubuntu-drivers is not available in this system") + + async def list_drivers(self, root_dir: str, context) -> List[str]: + # We run the command locally - ignoring the root_dir. + result = await arun_command(self.list_drivers_cmd) + return self._drivers_from_output(result.stdout) + + +def get_ubuntu_drivers_interface(app) -> UbuntuDriversInterface: + is_server = app.base_model.source.current.variant == "server" + cls: Type[UbuntuDriversInterface] = UbuntuDriversClientInterface + if app.opts.dry_run: + if 'has-drivers' in app.debug_flags: + cls = UbuntuDriversHasDriversInterface + elif 'run-drivers' in app.debug_flags: + cls = UbuntuDriversRunDriversInterface + else: + cls = UbuntuDriversNoDriversInterface + + return cls(app, gpgpu=is_server)