refactor: gpu setup and status checks, extract functions, and improve error handling

apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx

@@ -10,7 +10,7 @@ import {
 } from "@/components/ui/card";
 import { api } from "@/utils/api";
 import { TRPCClientError } from "@trpc/client";
-import { CheckCircle2, Cpu, Loader2, XCircle } from "lucide-react";
+import { CheckCircle2, Cpu, Loader2, RefreshCw, XCircle } from "lucide-react";
 import { useState } from "react";
 import { toast } from "sonner";
 
@@ -20,16 +20,19 @@ interface GPUSupportProps {
 
 export function GPUSupport({ serverId }: GPUSupportProps) {
 	const [isLoading, setIsLoading] = useState(false);
+	const [isRefreshing, setIsRefreshing] = useState(false);
 	const utils = api.useContext();
 
-	const { data: gpuStatus, isLoading: isChecking } =
+	const {
-		api.settings.checkGPUStatus.useQuery(
+		data: gpuStatus,
-			{ serverId },
+		isLoading: isChecking,
-			{
+		refetch,
-				enabled: serverId !== undefined,
+	} = api.settings.checkGPUStatus.useQuery(
-				refetchInterval: 5000,
+		{ serverId },
-			},
+		{
-		);
+			enabled: serverId !== undefined,
+		},
+	);
 
 	const setupGPU = api.settings.setupGPU.useMutation({
 		onMutate: () => {
@@ -41,31 +44,20 @@ export function GPUSupport({ serverId }: GPUSupportProps) {
 			await utils.settings.checkGPUStatus.invalidate({ serverId });
 		},
 		onError: (error) => {
-			if (error instanceof TRPCClientError) {
+			toast.error(
-				const errorMessage = error.message;
+				error.message ||
-				if (
+					"Failed to enable GPU support. Please check server logs.",
-					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.",
-					);
-				} else {
-					toast.error(errorMessage);
-				}
-			} else {
-				toast.error("Failed to enable GPU support. Please check server logs.");
-			}
 			setIsLoading(false);
 		},
 	});
 
+	const handleRefresh = async () => {
+		setIsRefreshing(true);
+		await refetch();
+		setIsRefreshing(false);
+	};
+
 	const handleEnableGPU = async () => {
 		if (serverId === undefined) {
 			toast.error("No server selected");
@@ -94,22 +86,33 @@ export function GPUSupport({ serverId }: GPUSupportProps) {
 									Configure and monitor GPU support
 								</CardDescription>
 							</div>
-							<DialogAction
+							<div className="flex items-center gap-2">
-								title="Enable GPU Support?"
+								<DialogAction
-								description="This will enable GPU support for Docker Swarm on this server. Make sure you have the required hardware and drivers installed."
+									title="Enable GPU Support?"
-								onClick={handleEnableGPU}
+									description="This will enable GPU support for Docker Swarm on this server. Make sure you have the required hardware and drivers installed."
-							>
+									onClick={handleEnableGPU}
-								<Button
-									isLoading={isLoading}
-									disabled={isLoading || serverId === undefined || isChecking}
 								>
-									{isLoading
+									<Button
-										? "Enabling GPU..."
+										isLoading={isLoading}
-										: gpuStatus?.swarmEnabled
+										disabled={isLoading || serverId === undefined || isChecking}
-											? "Reconfigure GPU"
+									>
-											: "Enable GPU"}
+										{isLoading
+											? "Enabling GPU..."
+											: gpuStatus?.swarmEnabled
+												? "Reconfigure GPU"
+												: "Enable GPU"}
+									</Button>
+								</DialogAction>
+								<Button
+									size="icon"
+									onClick={handleRefresh}
+									disabled={isChecking || isRefreshing}
+								>
+									<RefreshCw
+										className={`h-5 w-5 ${isChecking || isRefreshing ? "animate-spin" : ""}`}
+									/>
 								</Button>
-							</DialogAction>
+							</div>
 						</div>
 					</CardHeader>
 
@@ -117,9 +120,17 @@ export function GPUSupport({ serverId }: GPUSupportProps) {
 						<AlertBlock type="info">
 							<div className="font-medium mb-2">System Requirements:</div>
 							<ul className="list-disc list-inside text-sm space-y-1">
-								<li>NVIDIA drivers must be installed on the host system</li>
+								<li>NVIDIA GPU hardware must be physically installed</li>
-								<li>NVIDIA Container Runtime is required for GPU support</li>
+								<li>
-								<li>Compatible GPU hardware must be present</li>
+									NVIDIA drivers must be installed and running (check with
+									nvidia-smi)
+								</li>
+								<li>
+									NVIDIA Container Runtime must be installed
+									(nvidia-container-runtime)
+								</li>
+								<li>User must have sudo/administrative privileges</li>
+								<li>System must support CUDA for GPU acceleration</li>
 							</ul>
 						</AlertBlock>
 

apps/dokploy/server/api/routers/settings.ts

@@ -665,8 +665,8 @@ export const settingsRouter = createTRPCRouter({
 			}),
 		)
 		.mutation(async ({ input }) => {
-			if (IS_CLOUD) {
+			if (IS_CLOUD && !input.serverId) {
-				throw new Error("GPU setup is not available in cloud mode");
+				throw new Error("Select a server to enable the GPU Setup");
 			}
 
 			try {
@@ -684,7 +684,7 @@ export const settingsRouter = createTRPCRouter({
 			}),
 		)
 		.query(async ({ input }) => {
-			if (IS_CLOUD) {
+			if (IS_CLOUD && !input.serverId) {
 				return {
 					driverInstalled: false,
 					driverVersion: undefined,

packages/server/src/utils/gpu-setup.ts

@@ -18,117 +18,21 @@ interface GPUInfo {
 
 export async function checkGPUStatus(serverId?: string): Promise<GPUInfo> {
 	try {
-		// Check NVIDIA Driver
+		const [driverInfo, runtimeInfo, swarmInfo, gpuInfo, cudaInfo] =
-		let driverInstalled = false;
+			await Promise.all([
-		let driverVersion: string | undefined;
+				checkGpuDriver(serverId),
-		let availableGPUs = 0;
+				checkRuntime(serverId),
-
+				checkSwarmResources(serverId),
-		try {
+				checkGpuInfo(serverId),
-			const driverCommand =
+				checkCudaSupport(serverId),
-				"nvidia-smi --query-gpu=driver_version --format=csv,noheader";
+			]);
-			const { stdout: nvidiaSmi } = serverId
-				? await execAsyncRemote(serverId, driverCommand)
-				: await execAsync(driverCommand);
-
-			driverVersion = nvidiaSmi.trim();
-			if (driverVersion) {
-				driverInstalled = true;
-				const countCommand =
-					"nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l";
-				const { stdout: gpuCount } = serverId
-					? await execAsyncRemote(serverId, countCommand)
-					: await execAsync(countCommand);
-
-				availableGPUs = Number.parseInt(gpuCount.trim(), 10);
-			}
-		} catch (error) {
-			console.debug("GPU driver check:", error);
-		}
-
-		// Check Runtime Configuration
-		let runtimeInstalled = false;
-		let runtimeConfigured = false;
-		try {
-			const runtimeCommand = 'docker info --format "{{json .Runtimes}}"';
-			const { stdout: runtimeInfo } = serverId
-				? await execAsyncRemote(serverId, runtimeCommand)
-				: await execAsync(runtimeCommand);
-
-			const runtimes = JSON.parse(runtimeInfo);
-			runtimeInstalled = "nvidia" in runtimes;
-
-			// Check if it's the default runtime
-			const defaultCommand = 'docker info --format "{{.DefaultRuntime}}"';
-			const { stdout: defaultRuntime } = serverId
-				? await execAsyncRemote(serverId, defaultCommand)
-				: await execAsync(defaultCommand);
-
-			runtimeConfigured = defaultRuntime.trim() === "nvidia";
-		} catch (error) {
-			console.debug("Runtime check:", error);
-		}
-
-		// Check Swarm GPU Resources
-		let swarmEnabled = false;
-		let gpuResources = 0;
-
-		try {
-			// Check node resources directly from inspect
-			const nodeCommand =
-				"docker node inspect self --format '{{json .Description.Resources.GenericResources}}'";
-			const { stdout: resources } = serverId
-				? await execAsyncRemote(serverId, nodeCommand)
-				: await execAsync(nodeCommand);
-
-			if (resources && resources !== "null") {
-				const genericResources = JSON.parse(resources);
-				for (const resource of genericResources) {
-					if (
-						resource.DiscreteResourceSpec &&
-						(resource.DiscreteResourceSpec.Kind === "GPU" ||
-							resource.DiscreteResourceSpec.Kind === "gpu")
-					) {
-						gpuResources = resource.DiscreteResourceSpec.Value;
-						swarmEnabled = true;
-						break;
-					}
-				}
-			}
-		} catch (error) {
-			console.debug("Swarm resource check:", error);
-		}
-
-		// Get GPU Model and Memory Info
-		const gpuInfoCommand =
-			"nvidia-smi --query-gpu=gpu_name,memory.total --format=csv,noheader";
-		const { stdout: gpuInfo } = serverId
-			? await execAsyncRemote(serverId, gpuInfoCommand)
-			: await execAsync(gpuInfoCommand);
-
-		const [gpuModel, memoryTotal] = gpuInfo.split(",").map((s) => s.trim());
-
-		// Check CUDA Support
-		const cudaCommand = 'nvidia-smi -q | grep "CUDA Version"';
-		const { stdout: cudaInfo } = serverId
-			? await execAsyncRemote(serverId, cudaCommand)
-			: await execAsync(cudaCommand);
-
-		const cudaMatch = cudaInfo.match(/CUDA Version\s*:\s*([\d\.]+)/);
-		const cudaVersion = cudaMatch ? cudaMatch[1] : undefined;
-		const cudaSupport = !!cudaVersion;
 
 		return {
-			driverInstalled,
+			...driverInfo,
-			driverVersion,
+			...runtimeInfo,
-			runtimeInstalled,
+			...swarmInfo,
-			runtimeConfigured,
+			...gpuInfo,
-			availableGPUs,
+			...cudaInfo,
-			swarmEnabled,
-			gpuResources,
-			gpuModel,
-			memoryInfo: memoryTotal,
-			cudaSupport,
-			cudaVersion,
 		};
 	} catch (error) {
 		console.error("Error in checkGPUStatus:", error);
@@ -148,118 +52,167 @@ export async function checkGPUStatus(serverId?: string): Promise<GPUInfo> {
 	}
 }
 
+const checkGpuDriver = async (serverId?: string) => {
+	let driverVersion: string | undefined;
+	let driverInstalled = false;
+	let availableGPUs = 0;
+
+	try {
+		const driverCommand =
+			"nvidia-smi --query-gpu=driver_version --format=csv,noheader";
+		const { stdout: nvidiaSmi } = serverId
+			? await execAsyncRemote(serverId, driverCommand)
+			: await execAsync(driverCommand);
+
+		driverVersion = nvidiaSmi.trim();
+		if (driverVersion) {
+			driverInstalled = true;
+			const countCommand =
+				"nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l";
+			const { stdout: gpuCount } = serverId
+				? await execAsyncRemote(serverId, countCommand)
+				: await execAsync(countCommand);
+
+			availableGPUs = Number.parseInt(gpuCount.trim(), 10);
+		}
+	} catch (error) {
+		console.debug("GPU driver check:", error);
+	}
+
+	return { driverVersion, driverInstalled, availableGPUs };
+};
+
+const checkRuntime = async (serverId?: string) => {
+	let runtimeInstalled = false;
+	let runtimeConfigured = false;
+
+	try {
+		const runtimeCommand = 'docker info --format "{{json .Runtimes}}"';
+		const { stdout: runtimeInfo } = serverId
+			? await execAsyncRemote(serverId, runtimeCommand)
+			: await execAsync(runtimeCommand);
+
+		const runtimes = JSON.parse(runtimeInfo);
+		runtimeInstalled = "nvidia" in runtimes;
+
+		const defaultCommand = 'docker info --format "{{.DefaultRuntime}}"';
+		const { stdout: defaultRuntime } = serverId
+			? await execAsyncRemote(serverId, defaultCommand)
+			: await execAsync(defaultCommand);
+
+		runtimeConfigured = defaultRuntime.trim() === "nvidia";
+	} catch (error) {
+		console.debug("Runtime check:", error);
+	}
+
+	return { runtimeInstalled, runtimeConfigured };
+};
+
+const checkSwarmResources = async (serverId?: string) => {
+	let swarmEnabled = false;
+	let gpuResources = 0;
+
+	try {
+		const nodeCommand =
+			"docker node inspect self --format '{{json .Description.Resources.GenericResources}}'";
+		const { stdout: resources } = serverId
+			? await execAsyncRemote(serverId, nodeCommand)
+			: await execAsync(nodeCommand);
+
+		if (resources && resources !== "null") {
+			const genericResources = JSON.parse(resources);
+			for (const resource of genericResources) {
+				if (
+					resource.DiscreteResourceSpec &&
+					(resource.DiscreteResourceSpec.Kind === "GPU" ||
+						resource.DiscreteResourceSpec.Kind === "gpu")
+				) {
+					gpuResources = resource.DiscreteResourceSpec.Value;
+					swarmEnabled = true;
+					break;
+				}
+			}
+		}
+	} catch (error) {
+		console.debug("Swarm resource check:", error);
+	}
+
+	return { swarmEnabled, gpuResources };
+};
+
+const checkGpuInfo = async (serverId?: string) => {
+	let gpuModel: string | undefined;
+	let memoryInfo: string | undefined;
+
+	try {
+		const gpuInfoCommand =
+			"nvidia-smi --query-gpu=gpu_name,memory.total --format=csv,noheader";
+		const { stdout: gpuInfo } = serverId
+			? await execAsyncRemote(serverId, gpuInfoCommand)
+			: await execAsync(gpuInfoCommand);
+
+		[gpuModel, memoryInfo] = gpuInfo.split(",").map((s) => s.trim());
+	} catch (error) {
+		console.debug("GPU info check:", error);
+	}
+
+	return { gpuModel, memoryInfo };
+};
+
+const checkCudaSupport = async (serverId?: string) => {
+	let cudaVersion: string | undefined;
+	let cudaSupport = false;
+
+	try {
+		const cudaCommand = 'nvidia-smi -q | grep "CUDA Version"';
+		const { stdout: cudaInfo } = serverId
+			? await execAsyncRemote(serverId, cudaCommand)
+			: await execAsync(cudaCommand);
+
+		const cudaMatch = cudaInfo.match(/CUDA Version\s*:\s*([\d\.]+)/);
+		cudaVersion = cudaMatch ? cudaMatch[1] : undefined;
+		cudaSupport = !!cudaVersion;
+	} catch (error) {
+		console.debug("CUDA support check:", error);
+	}
+
+	return { cudaVersion, cudaSupport };
+};
+
 export async function setupGPUSupport(serverId?: string): Promise<void> {
 	try {
-		// 1. Check current GPU status first
+		// 1. Initial status check and validation
 		const initialStatus = await checkGPUStatus(serverId);
+		const shouldContinue = await validatePrerequisites(initialStatus);
+		if (!shouldContinue) return;
 
-		// If GPU is already configured, just verify and return quickly
+		// 2. Get node ID
-		if (
+		const nodeId = await getNodeId(serverId);
-			initialStatus.swarmEnabled &&
-			initialStatus.runtimeConfigured &&
-			initialStatus.driverInstalled
-		) {
-			console.log("GPU already configured, skipping setup");
-			return;
-		}
 
-		// 2. Verify GPU prerequisites
+		// 3. Create daemon configuration
-		if (!initialStatus.driverInstalled || !initialStatus.runtimeInstalled) {
+		const daemonConfig = createDaemonConfig(initialStatus.availableGPUs);
-			throw new Error(
-				"NVIDIA drivers or runtime not installed. Please install them first.",
-			);
-		}
 
-		// Get the node ID
+		// 4. Setup server based on environment
-		const nodeIdCommand = 'docker info --format "{{.Swarm.NodeID}}"';
-		const { stdout: nodeId } = serverId
-			? await execAsyncRemote(serverId, nodeIdCommand)
-			: await execAsync(nodeIdCommand);
-
-		if (!nodeId.trim()) {
-			throw new Error("Setup Server before enabling GPU support");
-		}
-
-		// 3. Configure NVIDIA runtime in daemon.json
-		const daemonConfig = {
-			runtimes: {
-				nvidia: {
-					path: "nvidia-container-runtime",
-					runtimeArgs: [],
-				},
-			},
-			"default-runtime": "nvidia",
-			"node-generic-resources": [`GPU=${initialStatus.availableGPUs}`],
-		};
-
-		// Different commands for local and remote setup
 		if (serverId) {
-			// Remote server setup (using sudo)
+			await setupRemoteServer(serverId, daemonConfig);
-			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)
+			await setupLocalServer(daemonConfig);
-			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);
 		}
 
-		// 4. Reduced wait time for Docker restart
+		// 5. Wait for Docker restart
-		await new Promise((resolve) => setTimeout(resolve, 10000));
+		await sleep(10000);
 
-		// 5. Add GPU label to the node
+		// 6. Add GPU label
-		const labelCommand = `docker node update --label-add gpu=true ${nodeId.trim()}`;
+		await addGpuLabel(nodeId, serverId);
-		if (serverId) {
-			await execAsyncRemote(serverId, labelCommand);
-		} else {
-			await execAsync(labelCommand);
-		}
 
-		// 6. Quick final verification
+		// 7. Final verification
-		await new Promise((resolve) => setTimeout(resolve, 5000));
+		await sleep(5000);
-		const finalStatus = await checkGPUStatus(serverId);
+		const finalStatus = await verifySetup(nodeId, serverId);
-
-		if (!finalStatus.swarmEnabled) {
-			const diagnosticCommands = [
-				`docker node inspect ${nodeId.trim()}`,
-				'nvidia-smi -a | grep "GPU UUID"',
-				"cat /etc/docker/daemon.json",
-				"cat /etc/nvidia-container-runtime/config.toml",
-			].join(" && ");
-
-			const { stdout: diagnostics } = serverId
-				? await execAsyncRemote(serverId, diagnosticCommands)
-				: await execAsync(diagnosticCommands);
-
-			console.error("Diagnostic Information:", diagnostics);
-			throw new Error("GPU support not detected in swarm after setup");
-		}
 
 		console.log("GPU setup completed successfully:", {
 			availableGPUs: initialStatus.availableGPUs,
 			driverVersion: initialStatus.driverVersion,
-			nodeId: nodeId.trim(),
+			nodeId,
 		});
 	} catch (error) {
 		console.error("GPU Setup Error:", error);
@@ -274,3 +227,113 @@ export async function setupGPUSupport(serverId?: string): Promise<void> {
 		throw error;
 	}
 }
+
+const validatePrerequisites = async (initialStatus: GPUInfo) => {
+	if (!initialStatus.driverInstalled) {
+		throw new Error(
+			"NVIDIA drivers not installed. Please install appropriate NVIDIA drivers first.",
+		);
+	}
+
+	if (!initialStatus.runtimeInstalled) {
+		throw new Error(
+			"NVIDIA Container Runtime not installed. Please install nvidia-container-runtime first.",
+		);
+	}
+
+	if (initialStatus.swarmEnabled && initialStatus.runtimeConfigured) {
+		console.log("GPU already configured, skipping setup");
+		return false;
+	}
+
+	return true;
+};
+
+const getNodeId = async (serverId?: string) => {
+	const nodeIdCommand = 'docker info --format "{{.Swarm.NodeID}}"';
+	const { stdout: nodeId } = serverId
+		? await execAsyncRemote(serverId, nodeIdCommand)
+		: await execAsync(nodeIdCommand);
+
+	const trimmedNodeId = nodeId.trim();
+	if (!trimmedNodeId) {
+		throw new Error("Setup Server before enabling GPU support");
+	}
+
+	return trimmedNodeId;
+};
+
+const createDaemonConfig = (availableGPUs: number) => ({
+	runtimes: {
+		nvidia: {
+			path: "nvidia-container-runtime",
+			runtimeArgs: [],
+		},
+	},
+	"default-runtime": "nvidia",
+	"node-generic-resources": [`GPU=${availableGPUs}`],
+});
+
+const setupRemoteServer = async (serverId: string, daemonConfig: any) => {
+	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);
+};
+
+const setupLocalServer = async (daemonConfig: any) => {
+	const configFile = `/tmp/docker-daemon-${Date.now()}.json`;
+	await fs.writeFile(configFile, JSON.stringify(daemonConfig, null, 2));
+
+	const setupCommands = [
+		`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}`,
+	].join(" && ");
+
+	await execAsync(setupCommands);
+};
+
+const addGpuLabel = async (nodeId: string, serverId?: string) => {
+	const labelCommand = `docker node update --label-add gpu=true ${nodeId}`;
+	if (serverId) {
+		await execAsyncRemote(serverId, labelCommand);
+	} else {
+		await execAsync(labelCommand);
+	}
+};
+
+const verifySetup = async (nodeId: string, serverId?: string) => {
+	const finalStatus = await checkGPUStatus(serverId);
+
+	if (!finalStatus.swarmEnabled) {
+		const diagnosticCommands = [
+			`docker node inspect ${nodeId}`,
+			'nvidia-smi -a | grep "GPU UUID"',
+			"cat /etc/docker/daemon.json",
+			"cat /etc/nvidia-container-runtime/config.toml",
+		].join(" && ");
+
+		const { stdout: diagnostics } = serverId
+			? await execAsyncRemote(serverId, diagnosticCommands)
+			: await execAsync(diagnosticCommands);
+
+		console.error("Diagnostic Information:", diagnostics);
+		throw new Error("GPU support not detected in swarm after setup");
+	}
+
+	return finalStatus;
+};
+
+const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));