Add the driver-package config option (#1)

Add the driver-package option to set
the package to be installed.

config.yaml

@@ -0,0 +1,6 @@
+options:
+  driver-package:
+    type: string
+    default: "nvidia-driver-latest-dkms"
+    description: |
+      Driver package to be installed.

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,

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()