fix: no need to manually load the nvidia drivers after reboot (#47)

Author: devalpatel67

application/service/akash_cluster_service.py

@@ -181,20 +181,20 @@ def _create_k3s_tasks(self, nodes):
                     )
                 )
 
-        # for i, node in reversed(list(enumerate(nodes))):
-        #     if node.install_gpu_drivers:
-        #         node_type = "main_node" if i == 0 else "worker_node"
-        #         k3s_tasks.append(
-        #             Task(
-        #                 str(uuid4()),
-        #                 f"restart_node_{node.hostname}",
-        #                 f"Restart node {node.hostname}",
-        #                 self.k3s_service._reboot_node,
-        #                 ssh_client,
-        #                 node,
-        #                 node_type,
-        #             )
-        #         )
+        for i, node in reversed(list(enumerate(nodes))):
+            if node.install_gpu_drivers:
+                node_type = "main_node" if i == 0 else "worker_node"
+                k3s_tasks.append(
+                    Task(
+                        str(uuid4()),
+                        f"restart_node_{node.hostname}",
+                        f"Restart node {node.hostname}",
+                        self.k3s_service._reboot_node,
+                        ssh_client,
+                        node,
+                        node_type,
+                    )
+                )
 
         return k3s_tasks
 

application/service/k3s_service.py

@@ -594,8 +594,6 @@ def _install_nvidia_drivers(
             "apt update",
             "apt-get install build-essential dkms linux-headers-$(uname -r) -y",
             "apt-get install nvidia-driver-570 -y",
-            "modprobe nvidia",
-            "nvidia-smi"
         ]
 
         nvidia_5090_commands = [