From 6e19e7449e49a1823409d0531f02e941564ec798 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sat, 22 Feb 2025 14:14:23 -0700
Subject: [PATCH] Fixed some issues with gpu info refreshing

---
 ui/src/components/JobOverview.tsx | 41 ++++++++++----------
 ui/src/hooks/useGPUInfo.tsx       | 63 ++++++++++++++++++++-----------
 2 files changed, 60 insertions(+), 44 deletions(-)

diff --git a/ui/src/components/JobOverview.tsx b/ui/src/components/JobOverview.tsx
index 61c48b84..92502491 100644
--- a/ui/src/components/JobOverview.tsx
+++ b/ui/src/components/JobOverview.tsx
@@ -3,26 +3,34 @@ import useGPUInfo from '@/hooks/useGPUInfo';
 import GPUWidget from '@/components/GPUWidget';
 import { getJobConfig, getTotalSteps } from '@/utils/jobs';
 import { Cpu, HardDrive, Info } from 'lucide-react';
+import { useMemo } from 'react';
 
 interface JobOverviewProps {
   job: Job;
 }
 
 export default function JobOverview({ job }: JobOverviewProps) {
-  const gpuIds = job.gpu_ids.split(',').map(id => parseInt(id));
-  const { gpuList, isGPUInfoLoaded } = useGPUInfo(gpuIds);
+  const gpuIds = useMemo(() => job.gpu_ids.split(',').map(id => parseInt(id)), [job.gpu_ids]);
+
+  const { gpuList, isGPUInfoLoaded } = useGPUInfo(gpuIds, 5000);
   const totalSteps = getTotalSteps(job);
   const progress = (job.step / totalSteps) * 100;
   const isStopping = job.stop && job.status === 'running';
 
   const getStatusColor = (status: string) => {
     switch (status.toLowerCase()) {
-      case 'running': return 'bg-emerald-500/10 text-emerald-500';
-      case 'stopping': return 'bg-amber-500/10 text-amber-500';
-      case 'stopped': return 'bg-gray-500/10 text-gray-400';
-      case 'completed': return 'bg-blue-500/10 text-blue-500';
-      case 'error': return 'bg-rose-500/10 text-rose-500';
-      default: return 'bg-gray-500/10 text-gray-400';
+      case 'running':
+        return 'bg-emerald-500/10 text-emerald-500';
+      case 'stopping':
+        return 'bg-amber-500/10 text-amber-500';
+      case 'stopped':
+        return 'bg-gray-500/10 text-gray-400';
+      case 'completed':
+        return 'bg-blue-500/10 text-blue-500';
+      case 'error':
+        return 'bg-rose-500/10 text-rose-500';
+      default:
+        return 'bg-gray-500/10 text-gray-400';
     }
   };
 
@@ -37,9 +45,7 @@ export default function JobOverview({ job }: JobOverviewProps) {
       <div className="col-span-2 bg-gray-900 rounded-xl shadow-lg overflow-hidden border border-gray-800">
         <div className="bg-gray-800 px-4 py-3 flex items-center justify-between">
           <h2 className="font-semibold text-gray-100">Job Details</h2>
-          <span className={`px-3 py-1 rounded-full text-sm ${getStatusColor(job.status)}`}>
-            {job.status}
-          </span>
+          <span className={`px-3 py-1 rounded-full text-sm ${getStatusColor(job.status)}`}>{job.status}</span>
         </div>
 
         <div className="p-4 space-y-6">
@@ -52,10 +58,7 @@ export default function JobOverview({ job }: JobOverviewProps) {
               </span>
             </div>
             <div className="w-full bg-gray-800 rounded-full h-2">
-              <div
-                className="h-2 rounded-full bg-blue-500 transition-all"
-                style={{ width: `${progress}%` }}
-              />
+              <div className="h-2 rounded-full bg-blue-500 transition-all" style={{ width: `${progress}%` }} />
             </div>
           </div>
 
@@ -89,11 +92,7 @@ export default function JobOverview({ job }: JobOverviewProps) {
       </div>
 
       {/* GPU Widget Panel */}
-      <div className="col-span-1">
-        {isGPUInfoLoaded && gpuList.length > 0 && (
-          <GPUWidget gpu={gpuList[0]} />
-        )}
-      </div>
+      <div className="col-span-1">{isGPUInfoLoaded && gpuList.length > 0 && <GPUWidget gpu={gpuList[0]} />}</div>
     </div>
   );
-}
\ No newline at end of file
+}
diff --git a/ui/src/hooks/useGPUInfo.tsx b/ui/src/hooks/useGPUInfo.tsx
index f8474701..5f2eda38 100644
--- a/ui/src/hooks/useGPUInfo.tsx
+++ b/ui/src/hooks/useGPUInfo.tsx
@@ -3,35 +3,52 @@
 import { GPUApiResponse, GpuInfo } from '@/types';
 import { useEffect, useState } from 'react';
 
-export default function useGPUInfo(gpuIds: null | number[] = null) {
+export default function useGPUInfo(gpuIds: null | number[] = null, reloadInterval: null | number = null) {
   const [gpuList, setGpuList] = useState<GpuInfo[]>([]);
   const [isGPUInfoLoaded, setIsLoaded] = useState(false);
-  useEffect(() => {
-    const fetchGpuInfo = async () => {
-      try {
-        const response = await fetch('/api/gpu');
+  const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
 
-        if (!response.ok) {
-          throw new Error(`HTTP error! Status: ${response.status}`);
-        }
+  const fetchGpuInfo = async () => {
+    setStatus('loading');
+    try {
+      const response = await fetch('/api/gpu');
 
-        const data: GPUApiResponse = await response.json();
-        let gpus = data.gpus.sort((a, b) => a.index - b.index);
-        if (gpuIds) {
-          gpus = gpus.filter(gpu => gpuIds.includes(gpu.index));
-        }
-
-        setGpuList(gpus);
-      } catch (err) {
-        console.log(`Failed to fetch GPU data: ${err instanceof Error ? err.message : String(err)}`);
-      } finally {
-        setIsLoaded(true);
+      if (!response.ok) {
+        throw new Error(`HTTP error! Status: ${response.status}`);
       }
-    };
 
+      const data: GPUApiResponse = await response.json();
+      let gpus = data.gpus.sort((a, b) => a.index - b.index);
+      if (gpuIds) {
+        gpus = gpus.filter(gpu => gpuIds.includes(gpu.index));
+      }
+
+      setGpuList(gpus);
+      setStatus('success');
+    } catch (err) {
+      console.error(`Failed to fetch GPU data: ${err instanceof Error ? err.message : String(err)}`);
+      setStatus('error');
+    } finally {
+      setIsLoaded(true);
+    }
+  };
+
+  useEffect(() => {
     // Fetch immediately on component mount
     fetchGpuInfo();
-  }, []);
 
-  return { gpuList, setGpuList, isGPUInfoLoaded };
-}
+    // Set up interval if specified
+    if (reloadInterval) {
+      const interval = setInterval(() => {
+        fetchGpuInfo();
+      }, reloadInterval);
+
+      // Cleanup interval on unmount
+      return () => {
+        clearInterval(interval);
+      };
+    }
+  }, [gpuIds, reloadInterval]); // Added dependencies
+
+  return { gpuList, setGpuList, isGPUInfoLoaded, status, refreshGpuInfo: fetchGpuInfo };
+}
\ No newline at end of file