From a6d46ad9ae664e68620b0f9afb9f09a0c4f6919d Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sat, 22 Feb 2025 13:54:06 -0700
Subject: [PATCH] Cleanup of job page

---
 ui/src/app/jobs/new/page.tsx      |   4 +-
 ui/src/components/GPUWidget.tsx   |   1 -
 ui/src/components/JobOverview.tsx | 100 ++++++++++++++++++++++++++----
 ui/src/hooks/useGPUInfo.tsx       |  13 ++--
 4 files changed, 99 insertions(+), 19 deletions(-)

diff --git a/ui/src/app/jobs/new/page.tsx b/ui/src/app/jobs/new/page.tsx
index 0633f31e..527334ed 100644
--- a/ui/src/app/jobs/new/page.tsx
+++ b/ui/src/app/jobs/new/page.tsx
@@ -64,7 +64,7 @@ export default function TrainingForm() {
   useEffect(() => {
     if (isGPUInfoLoaded) {
       if (gpuIDs === null && gpuList.length > 0) {
-        setGpuIDs(`${gpuList[0]}`);
+        setGpuIDs(`${gpuList[0].index}`);
       }
     }
   }, [gpuList, isGPUInfoLoaded]);
@@ -154,7 +154,7 @@ export default function TrainingForm() {
                 value={`${gpuIDs}`}
                 className="pt-2"
                 onChange={value => setGpuIDs(value)}
-                options={gpuList.map(gpu => ({ value: `${gpu}`, label: `GPU #${gpu}` }))}
+                options={gpuList.map(gpu => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
               />
             </Card>
 
diff --git a/ui/src/components/GPUWidget.tsx b/ui/src/components/GPUWidget.tsx
index e9f219b1..adb71ea7 100644
--- a/ui/src/components/GPUWidget.tsx
+++ b/ui/src/components/GPUWidget.tsx
@@ -28,7 +28,6 @@ export default function GPUWidget({ gpu }: GPUWidgetProps) {
             #{gpu.index}
           </span>
         </div>
-        <ChevronRight className="w-4 h-4 text-gray-400" />
       </div>
 
       <div className="p-4 space-y-4">
diff --git a/ui/src/components/JobOverview.tsx b/ui/src/components/JobOverview.tsx
index 3df5152c..61c48b84 100644
--- a/ui/src/components/JobOverview.tsx
+++ b/ui/src/components/JobOverview.tsx
@@ -1,23 +1,99 @@
 import { Job } from '@prisma/client';
+import useGPUInfo from '@/hooks/useGPUInfo';
+import GPUWidget from '@/components/GPUWidget';
+import { getJobConfig, getTotalSteps } from '@/utils/jobs';
+import { Cpu, HardDrive, Info } from 'lucide-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 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';
+    }
+  };
+
+  let status = job.status;
+  if (isStopping) {
+    status = 'stopping';
+  }
+
   return (
-    <>
-      <div className="grid grid-cols-1 gap-4 md:grid-cols-2">
-        <div className="">
-          <h2 className="text-lg font-semibold">Job Details</h2>
-          <p className="text-gray-400">ID: {job.id}</p>
-          <p className="text-gray-400">Name: {job.name}</p>
-          <p className="text-gray-400">GPUs: {job.gpu_ids}</p>
-          <p className="text-gray-400">Status: {job.status}</p>
-          <p className="text-gray-400">Info: {job.info}</p>
-          <p className="text-gray-400">Step: {job.step}</p>
+    <div className="grid grid-cols-1 gap-6 md:grid-cols-3">
+      {/* Job Information Panel */}
+      <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>
+        </div>
+
+        <div className="p-4 space-y-6">
+          {/* Progress Bar */}
+          <div className="space-y-2">
+            <div className="flex items-center justify-between text-sm">
+              <span className="text-gray-400">Progress</span>
+              <span className="text-gray-200">
+                Step {job.step} of {totalSteps}
+              </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>
+          </div>
+
+          {/* Job Info Grid */}
+          <div className="grid gap-4">
+            <div className="flex items-center space-x-4">
+              <HardDrive className="w-5 h-5 text-blue-400" />
+              <div>
+                <p className="text-sm font-medium text-gray-200">{job.name}</p>
+                <p className="text-xs text-gray-400">Job Name</p>
+              </div>
+            </div>
+
+            <div className="flex items-center space-x-4">
+              <Cpu className="w-5 h-5 text-purple-400" />
+              <div>
+                <p className="text-sm font-medium text-gray-200">GPUs: {job.gpu_ids}</p>
+                <p className="text-xs text-gray-400">Assigned GPUs</p>
+              </div>
+            </div>
+
+            <div className="flex items-center space-x-4">
+              <Info className="w-5 h-5 text-amber-400" />
+              <div>
+                <p className="text-sm font-medium text-gray-200">{job.info}</p>
+                <p className="text-xs text-gray-400">Additional Information</p>
+              </div>
+            </div>
+          </div>
         </div>
       </div>
-    </>
+
+      {/* GPU Widget Panel */}
+      <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 dd7e32fe..f8474701 100644
--- a/ui/src/hooks/useGPUInfo.tsx
+++ b/ui/src/hooks/useGPUInfo.tsx
@@ -1,10 +1,10 @@
 'use client';
 
-import { GPUApiResponse } from '@/types';
+import { GPUApiResponse, GpuInfo } from '@/types';
 import { useEffect, useState } from 'react';
 
-export default function useGPUInfo() {
-  const [gpuList, setGpuList] = useState<number[]>([]);
+export default function useGPUInfo(gpuIds: null | number[] = null) {
+  const [gpuList, setGpuList] = useState<GpuInfo[]>([]);
   const [isGPUInfoLoaded, setIsLoaded] = useState(false);
   useEffect(() => {
     const fetchGpuInfo = async () => {
@@ -16,7 +16,12 @@ export default function useGPUInfo() {
         }
 
         const data: GPUApiResponse = await response.json();
-        setGpuList(data.gpus.map(gpu => gpu.index).sort());
+        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 {