Faisal Amir dde260e723
fix: gpu hardware state (#4650)
* fix: gpu hardware state

* chore: cleanup state

* chore: clear state
2025-02-14 11:21:18 +07:00

390 lines
16 KiB
TypeScript

/* eslint-disable @typescript-eslint/no-explicit-any */
import * as React from 'react'
import { useState } from 'react'
import { DragDropContext, Draggable, Droppable } from '@hello-pangea/dnd'
import { Progress, ScrollArea, Switch } from '@janhq/joi'
import { useAtom, useAtomValue } from 'jotai'
import { atomWithStorage } from 'jotai/utils'
import { ChevronDownIcon, GripVerticalIcon } from 'lucide-react'
import { twMerge } from 'tailwind-merge'
import {
useGetHardwareInfo,
setActiveGpus,
} from '@/hooks/useHardwareManagement'
import { toGibibytes } from '@/utils/converter'
import { utilizedMemory } from '@/utils/memory'
import {
cpuUsageAtom,
ramUtilitizedAtom,
totalRamAtom,
usedRamAtom,
gpusAtom,
} from '@/helpers/atoms/SystemBar.atom'
const orderGpusAtom = atomWithStorage<any>('orderGpus', [], undefined, {
getOnInit: true,
})
const Hardware = () => {
const { hardware, mutate } = useGetHardwareInfo()
const [isActivatingGpu, setIsActivatingGpu] = useState<Set<string>>(new Set())
const [openPanels, setOpenPanels] = useState<Record<number, boolean>>({})
const cpuUsage = useAtomValue(cpuUsageAtom)
const totalRam = useAtomValue(totalRamAtom)
const usedRam = useAtomValue(usedRamAtom)
const ramUtilitized = useAtomValue(ramUtilitizedAtom)
const [gpus, setGpus] = useAtom(gpusAtom)
const [orderGpus, setOrderGpus] = useAtom(orderGpusAtom)
const togglePanel = (index: number) => {
setOpenPanels((prev) => ({
...prev,
[index]: !prev[index], // Toggle the specific panel
}))
}
// Handle switch toggle for GPU activation
const handleSwitchChange = async (id: string, isActive: boolean) => {
setIsActivatingGpu((prev) => new Set(prev).add(id))
const updatedGpus = gpus.map((gpu) =>
gpu.id === id ? { ...gpu, activated: isActive } : gpu
)
// Call the API to update the active GPUs
try {
const activeGpuIds = updatedGpus
.filter((gpu: any) => gpu.activated)
.map((gpu: any) => Number(gpu.id))
await setActiveGpus({ gpus: activeGpuIds })
mutate()
} catch (error) {
console.error('Failed to update active GPUs:', error)
}
}
const handleDragEnd = (result: any) => {
if (!result.destination) return
const reorderedGpus = Array.from(gpus)
const [movedGpu] = reorderedGpus.splice(result.source.index, 1)
reorderedGpus.splice(result.destination.index, 0, movedGpu)
setGpus(reorderedGpus)
setOrderGpus(reorderedGpus.map((gpu) => gpu.id))
}
React.useEffect(() => {
if (hardware?.gpus) {
setGpus((prevGpus: any) => {
// Create a map of existing GPUs by UUID for quick lookup
const gpuMap = new Map(prevGpus.map((gpu: any) => [gpu.uuid, gpu]))
// Update existing GPUs or add new ones
const updatedGpus = hardware.gpus.map((newGpu) => {
const existingGpu: any = gpuMap.get(newGpu.uuid)
if (existingGpu) {
// Update the GPU properties while keeping the original order
if (existingGpu.activated !== newGpu.activated) {
setIsActivatingGpu((prev) => {
const updated = new Set(prev)
updated.delete(existingGpu.id)
updated.clear()
return updated
})
}
return {
...existingGpu,
activated: newGpu.activated,
free_vram: newGpu.free_vram,
total_vram: newGpu.total_vram,
}
}
// Return the new GPU if not already in the state
return newGpu
})
// Append GPUs from the previous state that are not in the hardware.gpus
// This preserves user-reordered GPUs that aren't present in the new data
const remainingGpus = prevGpus.filter(
(prevGpu: any) =>
!hardware.gpus?.some((gpu) => gpu.uuid === prevGpu.uuid)
)
return [...updatedGpus, ...remainingGpus]
})
}
}, [hardware?.gpus, setGpus])
return (
<ScrollArea className="h-full w-full px-4">
<div className="block w-full py-4">
{/* CPU */}
<div className="flex w-full flex-col items-start justify-between gap-4 border-b border-[hsla(var(--app-border))] py-4 first:pt-0 last:border-none sm:flex-row">
<div className="flex-shrink-0 space-y-1">
<div className="flex gap-x-2">
<h6 className="font-semibold capitalize">CPU</h6>
</div>
</div>
<div className="w-full md:w-2/3">
<div className="flex flex-col items-end gap-2">
<div className="flex w-full justify-end gap-2 text-xs text-[hsla(var(--text-secondary))]">
<span>{hardware?.cpu.model}</span>
<span>|</span>
<span>Cores: {hardware?.cpu.cores}</span>
<span>|</span>
<span>Architecture: {hardware?.cpu.arch}</span>
</div>
<div className="flex w-2/3 items-center gap-3">
<Progress value={cpuUsage} size="small" className="w-full" />
<span className="font-medium">{cpuUsage}%</span>
</div>
</div>
</div>
</div>
{/* RAM */}
<div className="flex w-full flex-col items-start justify-between gap-4 border-b border-[hsla(var(--app-border))] py-4 first:pt-0 last:border-none sm:flex-row">
<div className="flex-shrink-0 space-y-1">
<div className="flex gap-x-2">
<h6 className="font-semibold capitalize">RAM</h6>
</div>
</div>
<div className="w-full md:w-2/3">
<div className="flex flex-col items-end gap-2">
<div className="flex w-full justify-end gap-2 text-xs text-[hsla(var(--text-secondary))]">
<span>
{toGibibytes(usedRam, { hideUnit: true })}GB /{' '}
{toGibibytes(totalRam, { hideUnit: true })}GB
</span>
{hardware?.ram.type && (
<>
<span>|</span>
<span>Type: {hardware?.ram.type}</span>
</>
)}
</div>
<div className="flex w-2/3 items-center gap-3">
<Progress
value={Math.round((usedRam / totalRam) * 100)}
size="small"
className="w-full"
/>
<span className="font-medium">{ramUtilitized}%</span>
</div>
</div>
</div>
</div>
{/* OS */}
<div className="flex w-full flex-col items-start justify-between gap-4 border-b border-[hsla(var(--app-border))] py-4 first:pt-0 last:border-none sm:flex-row">
<div className="flex-shrink-0 space-y-1">
<div className="flex gap-x-2">
<h6 className="font-semibold capitalize">OS</h6>
</div>
</div>
<div className="w-full md:w-2/3">
<div className="flex flex-col items-end gap-2">
<div className="flex w-full justify-end gap-2 text-xs text-[hsla(var(--text-secondary))]">
<span>{hardware?.os.name}</span>
<span>|</span>
<span>{hardware?.os.version}</span>
</div>
</div>
</div>
</div>
{/* GPUs */}
{!isMac && gpus.length > 0 && (
<div className="flex w-full flex-col items-start justify-between gap-4 border-b border-[hsla(var(--app-border))] py-4 first:pt-0 last:border-none sm:flex-row">
<div className="w-full flex-shrink-0">
<div className="flex gap-x-2">
<h6 className="font-semibold capitalize">GPUs</h6>
</div>
<p className="mt-1 font-medium leading-relaxed text-[hsla(var(--text-secondary))]">
{`Enhance model performance by utilizing your device's GPU for
acceleration.`}
</p>
<DragDropContext onDragEnd={handleDragEnd}>
<Droppable droppableId="gpu-list">
{(provided) => (
<div
{...provided.droppableProps}
ref={provided.innerRef}
className="mt-4"
>
{gpus
.sort((a, b) => {
const orderA = orderGpus.indexOf(a.id)
const orderB = orderGpus.indexOf(b.id)
return orderA - orderB
})
.map((item: any, i) => {
const gpuUtilization = utilizedMemory(
item.free_vram,
item.total_vram
)
const isLoading = isActivatingGpu.has(item.id)
return (
<Draggable
key={i}
draggableId={String(i)}
index={i}
>
{(provided, snapshot) => (
<div
ref={provided.innerRef}
{...provided.draggableProps}
{...provided.dragHandleProps}
className={twMerge(
'cursor-pointer border border-[hsla(var(--app-border))] bg-[hsla(var(--tertiary-bg))] p-4 first:rounded-t-lg last:rounded-b-lg',
gpus.length > 1 && 'last:rounded-t-none',
snapshot.isDragging
? 'border-b'
: 'border-b-0 last:border-b'
)}
onClick={() => togglePanel(i)}
>
<div className="flex flex-col items-start justify-start gap-4 sm:flex-row sm:items-center sm:justify-between">
<div className="flex w-full items-center justify-between">
<div className="flex h-full flex-shrink-0 items-center gap-2">
<GripVerticalIcon
size={14}
className="text-[hsla(var(--text-tertiary))]"
/>
<div
className={twMerge(
'h-2 w-2 rounded-full',
item.activated
? 'bg-green-400'
: 'bg-neutral-300'
)}
/>
<h6 title={item.name}>{item.name}</h6>
</div>
<div className="flex flex-shrink-0 items-end gap-4">
{item.activated && (
<div className="flex w-40 items-center gap-3">
<Progress
value={gpuUtilization}
size="small"
className="w-full"
/>
<span className="font-medium">
{gpuUtilization}%
</span>
</div>
)}
<div className="flex justify-end gap-2 text-xs text-[hsla(var(--text-secondary))]">
{item.activated && (
<span>
{(
(Number(item.total_vram) -
Number(item.free_vram)) /
1024
).toFixed(2)}
GB /{' '}
</span>
)}
<span>
{(
Number(item.total_vram) / 1024
).toFixed(2)}
GB
</span>
</div>
<Switch
checked={item.activated}
className={twMerge(
isLoading && 'pointer-events-none'
)}
disabled={
Boolean(isActivatingGpu.size) &&
!isLoading
}
onChange={(e) =>
handleSwitchChange(
item.id,
e.target.checked
)
}
/>
{isLoading && (
<div className="ml-2 h-4 w-4 animate-spin rounded-full border-t-2 border-solid border-blue-500" />
)}
<ChevronDownIcon
size={14}
className={twMerge(
'relative z-10 transform cursor-pointer transition-transform',
openPanels[i]
? 'rotate-180'
: 'rotate-0'
)}
/>
</div>
</div>
</div>
{openPanels[i] && (
<div className="space-y-4 p-4 pb-0 text-[hsla(var(--text-secondary))]">
<div className="flex">
<div className="w-[200px]">
Driver Version
</div>
<span>
{
item.additional_information
?.driver_version
}
</span>
</div>
<div className="flex">
<div className="w-[200px]">
Compute Capability
</div>
<span>
{
item.additional_information
?.compute_cap
}
</span>
</div>
</div>
)}
</div>
)}
</Draggable>
)
})}
{provided.placeholder}
</div>
)}
</Droppable>
</DragDropContext>
</div>
</div>
)}
</div>
</ScrollArea>
)
}
export default Hardware