From 556f1a4182450e55ba4e89946477226dc3e7d815 Mon Sep 17 00:00:00 2001
From: nameless <nameless@noreply@foxo.me>
Date: Fri, 10 Nov 2023 02:50:02 +0000
Subject: [PATCH] feat: show worker GPU

---
 api/workersRoute.ts        |  1 +
 app/generationQueue.ts     | 21 +++++++++++++++++++++
 app/workerInstanceStore.ts |  1 +
 ui/WorkersPage.tsx         |  9 ++++++---
 4 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/api/workersRoute.ts b/api/workersRoute.ts
index bf31c5a..9b7d0fa 100644
--- a/api/workersRoute.ts
+++ b/api/workersRoute.ts
@@ -55,6 +55,7 @@ async function getWorkerData(workerInstance: Model<WorkerInstance>): Promise<Wor
     id: workerInstance.id,
     key: workerInstance.value.key,
     name: workerInstance.value.name,
+    gpu: workerInstance.value.gpu,
     lastError: workerInstance.value.lastError,
     lastOnlineTime: workerInstance.value.lastOnlineTime,
     isActive: activeGenerationWorkers.get(workerInstance.id)?.isProcessing ?? false,
diff --git a/app/generationQueue.ts b/app/generationQueue.ts
index 49700d0..6ff615d 100644
--- a/app/generationQueue.ts
+++ b/app/generationQueue.ts
@@ -79,6 +79,27 @@ export async function processGenerationQueue() {
         continue;
       }
 
+      await workerSdClient.GET("/internal/sysinfo", {
+        params: {
+          query: {
+            attachment: false,
+          },
+        },
+        signal: AbortSignal.timeout(10_000),
+      })
+        .then((response) => {
+          if (!response.data) {
+            throw new SdError("Failed to get worker sysinfo", response.response, response.error);
+          }
+
+          // @ts-ignore there is no schema for /internal/sysinfo endpoint response
+          const nvidiaGPUModels = response.data["Torch env info"].nvidia_gpu_models ?? null;
+
+          if (nvidiaGPUModels !== null) {
+            workerInstance.update({ gpu: nvidiaGPUModels });
+          }
+        });
+
       // create worker
       const newWorker = generationQueue.createWorker(async ({ state }, updateJob) => {
         await processGenerationJob(state, updateJob, workerInstance.id);
diff --git a/app/workerInstanceStore.ts b/app/workerInstanceStore.ts
index 06325de..ea0c701 100644
--- a/app/workerInstanceStore.ts
+++ b/app/workerInstanceStore.ts
@@ -9,6 +9,7 @@ export const workerInstanceSchema = {
     key: { type: "string" },
     // used for display
     name: { type: ["string", "null"] },
+    gpu: { type: ["string", "null"] },
     sdUrl: { type: "string" },
     sdAuth: {
       type: ["object", "null"],
diff --git a/ui/WorkersPage.tsx b/ui/WorkersPage.tsx
index 0310cfe..e6cf60b 100644
--- a/ui/WorkersPage.tsx
+++ b/ui/WorkersPage.tsx
@@ -204,9 +204,12 @@ function WorkerListItem(props: {
 
   return (
     <li className="flex flex-col gap-2 rounded-md bg-zinc-100 dark:bg-zinc-800 p-2">
-      <p className="font-bold">
-        {worker.name ?? worker.key}
-      </p>
+      <div class="clearfix">
+        <span className="font-bold">
+          {worker.name ?? worker.key}
+        </span>
+        <span className="text-zinc-400 float-right">{worker.gpu}</span>
+      </div>
       {worker.isActive ? <p>✅ Active</p> : (
         <>
           <p>