Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add nouveau check and modprobe to install #359

Merged
merged 2 commits into from
Dec 3, 2024
Merged

Add nouveau check and modprobe to install #359

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ea5c789
Add nouveau check and modprobe
Deezzir Nov 29, 2024
e66d60e
Add unit tests for the new logic
Deezzir Dec 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion src/hw_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,13 +277,20 @@ def install(self) -> None:
logger.info("NVIDIA driver already installed in the machine")
return

with open("/proc/modules", encoding="utf-8") as modules:
if "nouveau" in modules.read():
logger.error("Nouveau driver is loaded. Unload it before installing NVIDIA driver")
raise ResourceInstallationError(self._name)

logger.info("Installing NVIDIA driver")
apt.add_package("ubuntu-drivers-common", update_cache=True)

try:
# This can be changed to check_call and not rely in the output if this is fixed
# https://github.com/canonical/ubuntu-drivers-common/issues/106
result = subprocess.check_output("ubuntu-drivers install --gpgpu".split(), text=True)
# https://bugs.launchpad.net/ubuntu/+source/ubuntu-drivers-common/+bug/2090502
result = subprocess.check_output("ubuntu-drivers --gpgpu install".split(), text=True)
subprocess.check_call("modprobe nvidia".split())

except subprocess.CalledProcessError as err:
logger.error("Failed to install the NVIDIA driver: %s", err)
Expand Down
19 changes: 17 additions & 2 deletions tests/unit/test_hw_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1219,7 +1219,7 @@ def test_dcgm_create_custom_metrics_copy_fail(


def test_nvidia_driver_strategy_install_success(
mock_path, mock_check_output, mock_apt_lib, nvidia_driver_strategy
mock_path, mock_check_output, mock_apt_lib, mock_check_call, nvidia_driver_strategy
):
nvidia_version = mock.MagicMock()
nvidia_version.exists.return_value = False
Expand All @@ -1228,7 +1228,8 @@ def test_nvidia_driver_strategy_install_success(
nvidia_driver_strategy.install()

mock_apt_lib.add_package.assert_called_once_with("ubuntu-drivers-common", update_cache=True)
mock_check_output.assert_called_once_with("ubuntu-drivers install --gpgpu".split(), text=True)
mock_check_output.assert_called_once_with("ubuntu-drivers --gpgpu install".split(), text=True)
mock_check_call.assert_called_once_with("modprobe nvidia".split())


def test_install_nvidia_drivers_already_installed(
Expand All @@ -1244,6 +1245,20 @@ def test_install_nvidia_drivers_already_installed(
mock_check_output.assert_not_called()


def test_install_nvidia_drivers_nouveau_installed(mock_path, nvidia_driver_strategy, mock_apt_lib):
nvidia_version = mock.MagicMock()
nvidia_version.exists.return_value = False
mock_path.return_value = nvidia_version
mocked_open = mock.mock_open(read_data="nouveau")

with mock.patch("builtins.open", mocked_open):
with pytest.raises(ResourceInstallationError):
nvidia_driver_strategy.install()

mock_apt_lib.add_package.assert_not_called()
mocked_open.assert_called_once_with("/proc/modules", encoding="utf-8")


def test_install_nvidia_drivers_subprocess_exception(
mock_path, mock_check_output, mock_apt_lib, nvidia_driver_strategy
):
Expand Down
Loading