diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..a6d22cb --- /dev/null +++ b/config.yaml @@ -0,0 +1,6 @@ +options: + driver-package: + type: string + default: "nvidia-driver-latest-dkms" + description: | + Driver package to be installed. \ No newline at end of file diff --git a/src/charm.py b/src/charm.py index b3f7840..266a9a9 100755 --- a/src/charm.py +++ b/src/charm.py @@ -17,7 +17,9 @@ def __init__(self, *args): """Initialize the charm.""" super().__init__(*args) - self._nvidia_ops_manager = NvidiaOpsManager() + driver_package = self.config.get("driver-package") + + self._nvidia_ops_manager = NvidiaOpsManager(driver_package) event_handler_bindings = { self.on.install: self._on_install, diff --git a/src/nvidia_ops_manager.py b/src/nvidia_ops_manager.py index 04848c0..01d73e9 100644 --- a/src/nvidia_ops_manager.py +++ b/src/nvidia_ops_manager.py @@ -18,7 +18,7 @@ def __init__(self, message: str): class NvidiaOpsManager: """NvidiaOpsManager.""" - def __init__(self): + def __init__(self, driver_package): """Initialize class level variables.""" self.PACKAGE_DEPS = [ "tar", @@ -39,7 +39,7 @@ def __init__(self): self.EPEL_RELEASE_REPO = ( "https://dl.fedoraproject.org/pub/epel/epel-release-latest-7.noarch.rpm" ) - self.NVIDIA_DRIVER_PACKAGE = "nvidia-driver-latest-dkms" + self.NVIDIA_DRIVER_PACKAGE = driver_package self.NVIDIA_DRIVER_REPO_FILEPATH = Path("/etc/yum.repos.d/cuda-rhel7.repo") @property @@ -89,7 +89,7 @@ def install(self) -> None: req = requests.get(self._nvidia_developer_repo) except requests.exceptions.HTTPError: raise NvidiaDriverOpsError( - "Error getting nvidia_developer_repository from {self._nvidia_developer_repo}." + f"Error getting nvidia_developer_repository from {self._nvidia_developer_repo}." ) self.NVIDIA_DRIVER_REPO_FILEPATH.write_text(req.text) # Add the devel kernel and kernel headers. @@ -102,19 +102,19 @@ def install(self) -> None: run(["yum", "clean", "expire-cache"]) except CalledProcessError: raise NvidiaDriverOpsError("Error flushing the cache.") - # Install nvidia-driver-latest-dkms.. + # Install nvidia-driver package.. try: run(["yum", "install", "-y", self.NVIDIA_DRIVER_PACKAGE]) except CalledProcessError: - raise NvidiaDriverOpsError("Error installing nvidia-driver-latest-dkms.") + raise NvidiaDriverOpsError("Error installing nvidia drivers.") def remove(self) -> None: """Remove nvidia drivers from the system.""" - # Remove nvidia-driver-latest-dkms.. + # Remove nvidia-driver package.. try: run(["yum", "erase", "-y", self.NVIDIA_DRIVER_PACKAGE]) except CalledProcessError: - raise NvidiaDriverOpsError("Error removing nvidia-driver-latest-dkms from the system.") + raise NvidiaDriverOpsError("Error removing nvidia drivers from the system.") # Remove the drivers repo if self.NVIDIA_DRIVER_REPO_FILEPATH.exists(): self.NVIDIA_DRIVER_REPO_FILEPATH.unlink()