From 2e6d9c34c0bdc61d9c4d8fb759f2ea9fddcbe654 Mon Sep 17 00:00:00 2001
From: vishalkadam47 <vishal@jeevops.com>
Date: Thu, 7 Nov 2024 02:52:41 +0530
Subject: [PATCH] feat: add dokploy server gpu setup

---
 .../servers/actions/show-dokploy-actions.tsx  |  2 +
 .../settings/servers/gpu-support-modal.tsx    | 36 +++++++++++++++
 .../settings/servers/gpu-support.tsx          | 26 ++++++-----
 apps/dokploy/server/api/routers/settings.ts   | 37 +++++++--------
 packages/server/src/utils/gpu-setup.ts        | 45 +++++++++++++++----
 5 files changed, 104 insertions(+), 42 deletions(-)
 create mode 100644 apps/dokploy/components/dashboard/settings/servers/gpu-support-modal.tsx

diff --git a/apps/dokploy/components/dashboard/settings/servers/actions/show-dokploy-actions.tsx b/apps/dokploy/components/dashboard/settings/servers/actions/show-dokploy-actions.tsx
index 49f6772b..9b12af84 100644
--- a/apps/dokploy/components/dashboard/settings/servers/actions/show-dokploy-actions.tsx
+++ b/apps/dokploy/components/dashboard/settings/servers/actions/show-dokploy-actions.tsx
@@ -13,6 +13,7 @@ import {
 import { api } from "@/utils/api";
 import { toast } from "sonner";
 import { ShowModalLogs } from "../../web-server/show-modal-logs";
+import { GPUSupportModal } from "../gpu-support-modal";
 
 export const ShowDokployActions = () => {
 	const { mutateAsync: reloadServer, isLoading } =
@@ -45,6 +46,7 @@ export const ShowDokployActions = () => {
 					<ShowModalLogs appName="dokploy">
 						<span>Watch logs</span>
 					</ShowModalLogs>
+					<GPUSupportModal />
 				</DropdownMenuGroup>
 			</DropdownMenuContent>
 		</DropdownMenu>
diff --git a/apps/dokploy/components/dashboard/settings/servers/gpu-support-modal.tsx b/apps/dokploy/components/dashboard/settings/servers/gpu-support-modal.tsx
new file mode 100644
index 00000000..9cf858cd
--- /dev/null
+++ b/apps/dokploy/components/dashboard/settings/servers/gpu-support-modal.tsx
@@ -0,0 +1,36 @@
+import {
+	Dialog,
+	DialogContent,
+	DialogHeader,
+	DialogTitle,
+	DialogTrigger,
+} from "@/components/ui/dialog";
+import { DropdownMenuItem } from "@/components/ui/dropdown-menu";
+import { useState } from "react";
+import { GPUSupport } from "./gpu-support";
+
+export const GPUSupportModal = () => {
+	const [isOpen, setIsOpen] = useState(false);
+
+	return (
+		<Dialog open={isOpen} onOpenChange={setIsOpen}>
+			<DialogTrigger asChild>
+				<DropdownMenuItem
+					className="w-full cursor-pointer"
+					onSelect={(e) => e.preventDefault()}
+				>
+					<span>GPU Setup</span>
+				</DropdownMenuItem>
+			</DialogTrigger>
+			<DialogContent className="sm:max-w-4xl overflow-y-auto max-h-screen">
+				<DialogHeader>
+					<DialogTitle className="flex items-center gap-2">
+						Dokploy Server GPU Setup
+					</DialogTitle>
+				</DialogHeader>
+
+				<GPUSupport serverId="" />
+			</DialogContent>
+		</Dialog>
+	);
+};
diff --git a/apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx b/apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx
index ae931a3a..d0c178c4 100644
--- a/apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx
+++ b/apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx
@@ -26,7 +26,7 @@ export function GPUSupport({ serverId }: GPUSupportProps) {
 		api.settings.checkGPUStatus.useQuery(
 			{ serverId },
 			{
-				enabled: !!serverId,
+				enabled: serverId !== undefined,
 				refetchInterval: 5000,
 			},
 		);
@@ -38,17 +38,20 @@ export function GPUSupport({ serverId }: GPUSupportProps) {
 		onSuccess: async () => {
 			toast.success("GPU support enabled successfully");
 			setIsLoading(false);
-
-			await Promise.all([
-				utils.settings.checkGPUStatus.invalidate({ serverId }),
-				utils.server.invalidate(),
-			]);
+			await utils.settings.checkGPUStatus.invalidate({ serverId });
 		},
 		onError: (error) => {
 			if (error instanceof TRPCClientError) {
 				const errorMessage = error.message;
-				if (errorMessage.includes("permission denied")) {
-					toast.error("Permission denied. Please ensure proper sudo access.");
+				if (
+					errorMessage.includes(
+						"Permission denied. Please ensure proper sudo access.",
+					) ||
+					errorMessage.includes("sudo access required")
+				) {
+					toast.error(
+						"Administrator privileges required. Please enter your password when prompted.",
+					);
 				} else if (errorMessage.includes("Failed to configure GPU")) {
 					toast.error(
 						"GPU configuration failed. Please check system requirements.",
@@ -59,13 +62,12 @@ export function GPUSupport({ serverId }: GPUSupportProps) {
 			} else {
 				toast.error("Failed to enable GPU support. Please check server logs.");
 			}
-
 			setIsLoading(false);
 		},
 	});
 
 	const handleEnableGPU = async () => {
-		if (!serverId) {
+		if (serverId === undefined) {
 			toast.error("No server selected");
 			return;
 		}
@@ -99,7 +101,7 @@ export function GPUSupport({ serverId }: GPUSupportProps) {
 							>
 								<Button
 									isLoading={isLoading}
-									disabled={isLoading || !serverId || isChecking}
+									disabled={isLoading || serverId === undefined || isChecking}
 								>
 									{isLoading
 										? "Enabling GPU..."
@@ -227,7 +229,7 @@ interface StatusRowProps {
 	showIcon?: boolean;
 }
 
-function StatusRow({
+export function StatusRow({
 	label,
 	isEnabled,
 	description,
diff --git a/apps/dokploy/server/api/routers/settings.ts b/apps/dokploy/server/api/routers/settings.ts
index 608a5028..a700c82a 100644
--- a/apps/dokploy/server/api/routers/settings.ts
+++ b/apps/dokploy/server/api/routers/settings.ts
@@ -661,33 +661,20 @@ export const settingsRouter = createTRPCRouter({
 	setupGPU: adminProcedure
 		.input(
 			z.object({
-				serverId: z.string(),
+				serverId: z.string().optional(),
 			}),
 		)
 		.mutation(async ({ input }) => {
+			if (IS_CLOUD) {
+				throw new Error("GPU setup is not available in cloud mode");
+			}
+
 			try {
-				if (IS_CLOUD) {
-					return { success: true };
-				}
-
-				if (!input.serverId) {
-					throw new TRPCError({
-						code: "BAD_REQUEST",
-						message: "Server ID is required",
-					});
-				}
-
 				await setupGPUSupport(input.serverId);
 				return { success: true };
 			} catch (error) {
-				throw new TRPCError({
-					code: "INTERNAL_SERVER_ERROR",
-					message:
-						error instanceof Error
-							? error.message
-							: "Failed to enable GPU support",
-					cause: error,
-				});
+				console.error("GPU Setup Error:", error);
+				throw error;
 			}
 		}),
 	checkGPUStatus: adminProcedure
@@ -712,7 +699,15 @@ export const settingsRouter = createTRPCRouter({
 					gpuResources: 0,
 				};
 			}
-			return await checkGPUStatus(input.serverId);
+
+			try {
+				const status = await checkGPUStatus(input.serverId || "");
+				console.log("GPU Status Check Result:", status);
+				return status;
+			} catch (error) {
+				console.error("GPU Status Check Error:", error);
+				throw new Error("Failed to check GPU status");
+			}
 		}),
 });
 // {
diff --git a/packages/server/src/utils/gpu-setup.ts b/packages/server/src/utils/gpu-setup.ts
index f1936bf6..ecdb3e2b 100644
--- a/packages/server/src/utils/gpu-setup.ts
+++ b/packages/server/src/utils/gpu-setup.ts
@@ -1,3 +1,4 @@
+import * as fs from "node:fs/promises";
 import { execAsync } from "../utils/process/execAsync";
 import { execAsyncRemote } from "../utils/process/execAsync";
 
@@ -191,18 +192,36 @@ export async function setupGPUSupport(serverId?: string): Promise<void> {
 			"node-generic-resources": [`GPU=${initialStatus.availableGPUs}`],
 		};
 
-		const setupCommands = [
-			"sudo -n true",
-			`echo '${JSON.stringify(daemonConfig, null, 2)}' | sudo tee /etc/docker/daemon.json`,
-			"sudo mkdir -p /etc/nvidia-container-runtime",
-			'echo "swarm-resource = \\"DOCKER_RESOURCE_GPU\\"" | sudo tee -a /etc/nvidia-container-runtime/config.toml',
-			"sudo systemctl daemon-reload",
-			"sudo systemctl restart docker",
-		].join(" && ");
-
+		// Different commands for local and remote setup
 		if (serverId) {
+			// Remote server setup (using sudo)
+			const setupCommands = [
+				"sudo -n true",
+				`echo '${JSON.stringify(daemonConfig, null, 2)}' | sudo tee /etc/docker/daemon.json`,
+				"sudo mkdir -p /etc/nvidia-container-runtime",
+				'echo "swarm-resource = \\"DOCKER_RESOURCE_GPU\\"" | sudo tee -a /etc/nvidia-container-runtime/config.toml',
+				"sudo systemctl daemon-reload",
+				"sudo systemctl restart docker",
+			].join(" && ");
+
 			await execAsyncRemote(serverId, setupCommands);
 		} else {
+			// Local server setup (using pkexec for GUI password prompt)
+			const configFile = `/tmp/docker-daemon-${Date.now()}.json`;
+			await fs.writeFile(configFile, JSON.stringify(daemonConfig, null, 2));
+
+			const setupCommands = [
+				// Use pkexec for GUI password prompt
+				`pkexec sh -c '
+					cp ${configFile} /etc/docker/daemon.json && 
+					mkdir -p /etc/nvidia-container-runtime && 
+					echo "swarm-resource = \\"DOCKER_RESOURCE_GPU\\"" >> /etc/nvidia-container-runtime/config.toml && 
+					systemctl daemon-reload && 
+					systemctl restart docker
+				'`,
+				`rm ${configFile}`, // Clean up temp file
+			].join(" && ");
+
 			await execAsync(setupCommands);
 		}
 
@@ -244,6 +263,14 @@ export async function setupGPUSupport(serverId?: string): Promise<void> {
 		});
 	} catch (error) {
 		console.error("GPU Setup Error:", error);
+		if (
+			error instanceof Error &&
+			error.message.includes("password is required")
+		) {
+			throw new Error(
+				"Sudo access required. Please run with appropriate permissions.",
+			);
+		}
 		throw error;
 	}
 }