From 7306d8c5139f4c41c2b2334c0a3c8d2432c44c88 Mon Sep 17 00:00:00 2001
From: vishalkadam47 <vishal@jeevops.com>
Date: Sun, 3 Nov 2024 21:34:03 +0530
Subject: [PATCH] feat: Add GPU configuration and Update import path for
 gpu-setup functions

---
 .../settings/servers/gpu-support.tsx          | 219 ++++++++++++++++++
 apps/dokploy/server/api/routers/settings.ts   |   2 +-
 2 files changed, 220 insertions(+), 1 deletion(-)
 create mode 100644 apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx

diff --git a/apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx b/apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx
new file mode 100644
index 00000000..a0ef8d80
--- /dev/null
+++ b/apps/dokploy/components/dashboard/settings/servers/gpu-support.tsx
@@ -0,0 +1,219 @@
+import { Button } from '@/components/ui/button';
+import { useState } from 'react';
+import { api } from '@/utils/api';
+import { toast } from 'sonner';
+import { TRPCClientError } from '@trpc/client';
+import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/components/ui/card';
+import { DialogAction } from '@/components/shared/dialog-action';
+import { AlertBlock } from '@/components/shared/alert-block';
+import { Cpu, CheckCircle2, XCircle, Loader2 } from 'lucide-react';
+
+interface GPUSupportProps {
+  serverId?: string;
+}
+
+export function GPUSupport({ serverId }: GPUSupportProps) {
+  const [isLoading, setIsLoading] = useState(false);
+  const utils = api.useContext();
+
+  const { data: gpuStatus, isLoading: isChecking } = api.settings.checkGPUStatus.useQuery(
+    { serverId },
+    { 
+      enabled: !!serverId,
+      refetchInterval: 5000
+    }
+  );
+
+const setupGPU = api.settings.setupGPU.useMutation({
+  onMutate: () => {
+    setIsLoading(true);
+  },
+  onSuccess: async () => {
+    toast.success('GPU support enabled successfully');
+    setIsLoading(false);
+    
+    await Promise.all([
+      utils.settings.checkGPUStatus.invalidate({ serverId }),
+      utils.server.invalidate()
+    ]);
+  },
+  onError: (error) => {
+    if (error instanceof TRPCClientError) {
+      const errorMessage = error.message;
+      if (errorMessage.includes('permission denied')) {
+        toast.error('Permission denied. Please ensure proper sudo access.');
+      } 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 handleEnableGPU = async () => {
+    if (!serverId) {
+      toast.error('No server selected');
+      return;
+    }
+
+    try {
+      await setupGPU.mutateAsync({ serverId });
+    } catch (error) {
+      // Error handling is done in mutation's onError
+    }
+  };
+
+  return (
+    <CardContent className="p-0">
+      <div className="flex flex-col gap-4">
+        <Card className="bg-background">
+          <CardHeader className="flex flex-row items-center justify-between flex-wrap gap-2">
+            <div className="flex flex-row gap-2 justify-between w-full items-end max-sm:flex-col">
+              <div className="flex flex-col gap-1">
+                <div className="flex items-center gap-2">
+                  <Cpu className="size-5" />
+                  <CardTitle className="text-xl">GPU Configuration</CardTitle>
+                </div>
+                <CardDescription>Configure and monitor GPU support</CardDescription>
+              </div>
+              <DialogAction
+                title="Enable GPU Support?"
+                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 || isChecking}
+                >
+                  {isLoading ? (
+                    'Enabling GPU...'
+                  ) : gpuStatus?.swarmEnabled ? (
+                    'Reconfigure GPU'
+                  ) : (
+                    'Enable GPU'
+                  )}
+                </Button>
+              </DialogAction>
+            </div>
+          </CardHeader>
+
+          <CardContent className="flex flex-col gap-4">
+            <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 Container Runtime is required for GPU support</li>
+                <li>Compatible GPU hardware must be present</li>
+              </ul>
+            </AlertBlock>
+
+            {isChecking ? (
+              <div className="flex items-center justify-center text-muted-foreground py-4">
+                <Loader2 className="mr-2 h-4 w-4 animate-spin" />
+                <span>Checking GPU status...</span>
+              </div>
+            ) : (
+              <div className="grid gap-4">
+                {/* Prerequisites Section */}
+                <div className="border rounded-lg p-4">
+                  <h3 className="text-lg font-semibold mb-1">Prerequisites</h3>
+                  <p className="text-sm text-muted-foreground mb-4">Shows all software checks and available hardware</p>
+                  <div className="grid gap-2.5">
+                    <StatusRow 
+                        label="NVIDIA Driver"
+                        isEnabled={gpuStatus?.driverInstalled}
+                        description={gpuStatus?.driverVersion ? `Installed (v${gpuStatus.driverVersion})` : 'Not Installed'}
+                    />
+                    <StatusRow 
+                        label="GPU Model"
+                        value={gpuStatus?.gpuModel || 'Not Detected'}
+                        showIcon={false}
+                    />
+                    <StatusRow 
+                        label="GPU Memory"
+                        value={gpuStatus?.memoryInfo || 'Not Available'}
+                        showIcon={false}
+                    />
+                    <StatusRow 
+                        label="Available GPUs"
+                        value={gpuStatus?.availableGPUs || 0}
+                        showIcon={false}
+                    />
+                    <StatusRow 
+                        label="CUDA Support"
+                        isEnabled={gpuStatus?.cudaSupport}
+                        description={gpuStatus?.cudaVersion ? `Available (v${gpuStatus.cudaVersion})` : 'Not Available'}
+                    />
+                    <StatusRow 
+                        label="NVIDIA Container Runtime"
+                        isEnabled={gpuStatus?.runtimeInstalled}
+                        description={gpuStatus?.runtimeInstalled ? 'Installed' : 'Not Installed'}
+                    />
+                </div>
+                </div>
+
+                {/* Configuration Status */}
+                <div className="border rounded-lg p-4">
+                  <h3 className="text-lg font-semibold mb-1">Docker Swarm GPU Status</h3>
+                  <p className="text-sm text-muted-foreground mb-4">Shows the configuration state that changes with the Enable GPU</p>
+                  <div className="grid gap-2.5">
+                    <StatusRow 
+                        label="Runtime Configuration"
+                        isEnabled={gpuStatus?.runtimeConfigured}
+                        description={gpuStatus?.runtimeConfigured ? 'Default Runtime' : 'Not Default Runtime'}
+                    />
+                    <StatusRow 
+                        label="Swarm GPU Support"
+                        isEnabled={gpuStatus?.swarmEnabled}
+                        description={gpuStatus?.swarmEnabled 
+                          ? `Enabled (${gpuStatus.gpuResources} GPU${gpuStatus.gpuResources !== 1 ? 's' : ''})` 
+                          : 'Not Enabled'
+                        }
+                    />
+                  </div>
+                </div>
+              </div>
+            )}
+          </CardContent>
+        </Card>
+      </div>
+    </CardContent>
+  );
+}
+
+interface StatusRowProps {
+  label: string;
+  isEnabled?: boolean;
+  description?: string;
+  value?: string | number;
+  showIcon?: boolean;
+}
+
+function StatusRow({ label, isEnabled, description, value, showIcon = true }: StatusRowProps) {
+  return (
+    <div className="flex items-center justify-between">
+      <span className="text-sm">{label}</span>
+      <div className="flex items-center gap-2">
+        {showIcon ? (
+          <>
+            {isEnabled ? (
+              <CheckCircle2 className="size-4 text-green-500" />
+            ) : (
+              <XCircle className="size-4 text-red-500" />
+            )}
+            <span className={`text-sm ${isEnabled ? 'text-green-500' : 'text-red-500'}`}>
+              {description || (isEnabled ? 'Installed' : 'Not Installed')}
+            </span>
+          </>
+        ) : (
+          <span className="text-sm text-muted-foreground">{value}</span>
+        )}
+      </div>
+    </div>
+  );
+}
\ No newline at end of file
diff --git a/apps/dokploy/server/api/routers/settings.ts b/apps/dokploy/server/api/routers/settings.ts
index 13f67126..94167a2e 100644
--- a/apps/dokploy/server/api/routers/settings.ts
+++ b/apps/dokploy/server/api/routers/settings.ts
@@ -55,7 +55,7 @@ import {
 import {
 	checkGPUStatus,
 	setupGPUSupport,
-} from "@dokploy/server/src/utils/gpu-setup";
+} from "@dokploy/server";
 import { generateOpenApiDocument } from "@dokploy/trpc-openapi";
 import { TRPCError } from "@trpc/server";
 import { sql } from "drizzle-orm";