mirror of
https://github.com/NVIDIA/dgx-spark-playbooks.git
synced 2026-04-23 02:23:53 +00:00
433 lines
14 KiB
TypeScript
433 lines
14 KiB
TypeScript
//
|
|
// SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
"use client"
|
|
|
|
import React, { useState, useEffect, useRef } from 'react'
|
|
import { Button } from '@/components/ui/button'
|
|
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card'
|
|
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select'
|
|
import { Badge } from '@/components/ui/badge'
|
|
import { Switch } from '@/components/ui/switch'
|
|
import { Loader2, Zap, Activity, Cpu, MonitorSpeaker } from 'lucide-react'
|
|
import { useToast } from '@/hooks/use-toast'
|
|
|
|
interface GraphData {
|
|
nodes: Array<{
|
|
id: string
|
|
name: string
|
|
group?: string
|
|
[key: string]: any
|
|
}>
|
|
links: Array<{
|
|
source: string
|
|
target: string
|
|
name: string
|
|
[key: string]: any
|
|
}>
|
|
}
|
|
|
|
interface LocalGPUViewerProps {
|
|
graphData: GraphData
|
|
onError?: (error: Error) => void
|
|
}
|
|
|
|
interface GPUCapabilities {
|
|
has_rapids: boolean
|
|
has_torch_geometric: boolean
|
|
gpu_available: boolean
|
|
supported_layouts: string[]
|
|
supported_clustering: string[]
|
|
}
|
|
|
|
interface ProcessedData {
|
|
nodes: any[]
|
|
edges: any[]
|
|
gpu_processed: boolean
|
|
layout_positions: Record<string, [number, number]>
|
|
clusters: Record<string, number>
|
|
centrality: Record<string, Record<string, number>>
|
|
stats: {
|
|
node_count: number
|
|
edge_count: number
|
|
gpu_accelerated: boolean
|
|
layout_computed: boolean
|
|
clusters_computed: boolean
|
|
centrality_computed: boolean
|
|
}
|
|
timestamp: string
|
|
}
|
|
|
|
export function LocalGPUViewer({ graphData, onError }: LocalGPUViewerProps) {
|
|
const [isProcessing, setIsProcessing] = useState(false)
|
|
const [processedData, setProcessedData] = useState<ProcessedData | null>(null)
|
|
const [capabilities, setCapabilities] = useState<GPUCapabilities | null>(null)
|
|
const [serviceHealth, setServiceHealth] = useState<'unknown' | 'healthy' | 'error'>('unknown')
|
|
|
|
// Processing options
|
|
const [layoutAlgorithm, setLayoutAlgorithm] = useState('force_atlas2')
|
|
const [clusteringAlgorithm, setClusteringAlgorithm] = useState('leiden')
|
|
const [gpuAcceleration, setGpuAcceleration] = useState(true)
|
|
const [computeCentrality, setComputeCentrality] = useState(true)
|
|
|
|
const { toast } = useToast()
|
|
const wsRef = useRef<WebSocket | null>(null)
|
|
|
|
// Check service health and capabilities on mount
|
|
useEffect(() => {
|
|
checkServiceHealth()
|
|
getCapabilities()
|
|
setupWebSocket()
|
|
|
|
return () => {
|
|
if (wsRef.current) {
|
|
wsRef.current.close()
|
|
}
|
|
}
|
|
}, [])
|
|
|
|
const checkServiceHealth = async () => {
|
|
try {
|
|
const response = await fetch('/api/local-gpu/health')
|
|
if (response.ok) {
|
|
setServiceHealth('healthy')
|
|
} else {
|
|
setServiceHealth('error')
|
|
}
|
|
} catch (error) {
|
|
console.error('Local GPU service health check failed:', error)
|
|
setServiceHealth('error')
|
|
}
|
|
}
|
|
|
|
const getCapabilities = async () => {
|
|
try {
|
|
const response = await fetch('/api/local-gpu/capabilities')
|
|
if (response.ok) {
|
|
const caps = await response.json()
|
|
setCapabilities(caps)
|
|
|
|
// Set defaults based on capabilities
|
|
if (caps.supported_layouts?.length > 0) {
|
|
setLayoutAlgorithm(caps.supported_layouts[0])
|
|
}
|
|
if (caps.supported_clustering?.length > 0) {
|
|
setClusteringAlgorithm(caps.supported_clustering[0])
|
|
}
|
|
}
|
|
} catch (error) {
|
|
console.error('Failed to get GPU capabilities:', error)
|
|
}
|
|
}
|
|
|
|
const setupWebSocket = () => {
|
|
try {
|
|
const ws = new WebSocket('ws://localhost:8081/ws')
|
|
|
|
ws.onopen = () => {
|
|
console.log('WebSocket connected to local GPU service')
|
|
}
|
|
|
|
ws.onmessage = (event) => {
|
|
try {
|
|
const message = JSON.parse(event.data)
|
|
if (message.type === 'graph_processed') {
|
|
setProcessedData(message.data)
|
|
setIsProcessing(false)
|
|
|
|
toast({
|
|
title: "GPU Processing Complete",
|
|
description: `Processed ${message.data.stats.node_count} nodes with ${message.data.gpu_processed ? 'GPU' : 'CPU'} acceleration.`,
|
|
})
|
|
}
|
|
} catch (error) {
|
|
console.error('WebSocket message error:', error)
|
|
}
|
|
}
|
|
|
|
ws.onerror = (error) => {
|
|
console.error('WebSocket error:', error)
|
|
}
|
|
|
|
wsRef.current = ws
|
|
} catch (error) {
|
|
console.error('WebSocket setup failed:', error)
|
|
}
|
|
}
|
|
|
|
const processWithLocalGPU = async () => {
|
|
if (!graphData?.nodes?.length || !graphData?.links?.length) {
|
|
toast({
|
|
title: "No Graph Data",
|
|
description: "Please ensure graph data is loaded before processing.",
|
|
variant: "destructive"
|
|
})
|
|
return
|
|
}
|
|
|
|
setIsProcessing(true)
|
|
|
|
try {
|
|
const requestData = {
|
|
graph_data: {
|
|
nodes: graphData.nodes,
|
|
links: graphData.links
|
|
},
|
|
layout_algorithm: layoutAlgorithm,
|
|
clustering_algorithm: clusteringAlgorithm,
|
|
gpu_acceleration: gpuAcceleration,
|
|
compute_centrality: computeCentrality
|
|
}
|
|
|
|
const response = await fetch('/api/local-gpu/process', {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json'
|
|
},
|
|
body: JSON.stringify(requestData)
|
|
})
|
|
|
|
if (!response.ok) {
|
|
throw new Error(`Local GPU processing failed: ${response.statusText}`)
|
|
}
|
|
|
|
// The result will come via WebSocket, so we just wait
|
|
toast({
|
|
title: "Processing Started",
|
|
description: "Graph processing started on local GPU. Results will appear shortly.",
|
|
})
|
|
|
|
} catch (error) {
|
|
console.error('Local GPU processing error:', error)
|
|
const errorMessage = error instanceof Error ? error.message : 'Unknown error occurred'
|
|
|
|
toast({
|
|
title: "Processing Failed",
|
|
description: errorMessage,
|
|
variant: "destructive"
|
|
})
|
|
|
|
setIsProcessing(false)
|
|
|
|
if (onError) {
|
|
onError(error instanceof Error ? error : new Error(errorMessage))
|
|
}
|
|
}
|
|
}
|
|
|
|
const ServiceStatus = () => (
|
|
<div className="flex items-center gap-2 text-sm">
|
|
<div className={`w-2 h-2 rounded-full ${
|
|
serviceHealth === 'healthy' ? 'bg-green-500' :
|
|
serviceHealth === 'error' ? 'bg-red-500' : 'bg-yellow-500'
|
|
}`} />
|
|
<span>
|
|
Local GPU Service: {serviceHealth === 'healthy' ? 'Connected' :
|
|
serviceHealth === 'error' ? 'Disconnected' : 'Checking...'}
|
|
</span>
|
|
{capabilities?.gpu_available && (
|
|
<Badge variant="outline" className="ml-2">
|
|
<Zap className="w-3 h-3 mr-1" />
|
|
cuGraph Ready
|
|
</Badge>
|
|
)}
|
|
</div>
|
|
)
|
|
|
|
const StatsDisplay = ({ stats }: { stats: ProcessedData['stats'] }) => (
|
|
<div className="grid grid-cols-2 md:grid-cols-3 gap-4 text-sm">
|
|
<div className="flex items-center gap-2">
|
|
<Activity className="w-4 h-4 text-blue-500" />
|
|
<span>{stats.node_count.toLocaleString()} nodes</span>
|
|
</div>
|
|
<div className="flex items-center gap-2">
|
|
<Activity className="w-4 h-4 text-green-500" />
|
|
<span>{stats.edge_count.toLocaleString()} edges</span>
|
|
</div>
|
|
<div className="flex items-center gap-2">
|
|
{stats.gpu_accelerated ? (
|
|
<Zap className="w-4 h-4 text-yellow-500" />
|
|
) : (
|
|
<Cpu className="w-4 h-4 text-gray-500" />
|
|
)}
|
|
<span>{stats.gpu_accelerated ? 'GPU' : 'CPU'} processed</span>
|
|
</div>
|
|
{stats.layout_computed && (
|
|
<div className="flex items-center gap-2">
|
|
<MonitorSpeaker className="w-4 h-4 text-purple-500" />
|
|
<span>Layout computed</span>
|
|
</div>
|
|
)}
|
|
{stats.clusters_computed && (
|
|
<div className="flex items-center gap-2">
|
|
<Activity className="w-4 h-4 text-orange-500" />
|
|
<span>Clusters detected</span>
|
|
</div>
|
|
)}
|
|
{stats.centrality_computed && (
|
|
<div className="flex items-center gap-2">
|
|
<Activity className="w-4 h-4 text-pink-500" />
|
|
<span>Centrality computed</span>
|
|
</div>
|
|
)}
|
|
</div>
|
|
)
|
|
|
|
return (
|
|
<div className="space-y-4">
|
|
{/* Service Status */}
|
|
<Card>
|
|
<CardHeader className="pb-3">
|
|
<CardTitle className="flex items-center justify-between">
|
|
<span className="flex items-center gap-2">
|
|
<Zap className="w-5 h-5" />
|
|
Local GPU Visualization
|
|
</span>
|
|
<ServiceStatus />
|
|
</CardTitle>
|
|
</CardHeader>
|
|
<CardContent className="space-y-4">
|
|
{/* GPU Capabilities */}
|
|
{capabilities && (
|
|
<div className="p-3 bg-muted rounded-lg">
|
|
<h4 className="font-medium mb-2">GPU Capabilities</h4>
|
|
<div className="grid grid-cols-2 gap-2 text-sm">
|
|
<div className="flex items-center gap-2">
|
|
<div className={`w-2 h-2 rounded-full ${capabilities.has_rapids ? 'bg-green-500' : 'bg-red-500'}`} />
|
|
<span>RAPIDS cuGraph: {capabilities.has_rapids ? 'Available' : 'Not Available'}</span>
|
|
</div>
|
|
<div className="flex items-center gap-2">
|
|
<div className={`w-2 h-2 rounded-full ${capabilities.has_torch_geometric ? 'bg-green-500' : 'bg-red-500'}`} />
|
|
<span>PyTorch Geometric: {capabilities.has_torch_geometric ? 'Available' : 'Not Available'}</span>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
)}
|
|
|
|
{/* Processing Controls */}
|
|
<div className="grid grid-cols-1 md:grid-cols-2 gap-4">
|
|
<div className="space-y-4">
|
|
<div className="space-y-2">
|
|
<label className="text-sm font-medium">Layout Algorithm</label>
|
|
<Select value={layoutAlgorithm} onValueChange={setLayoutAlgorithm}>
|
|
<SelectTrigger>
|
|
<SelectValue />
|
|
</SelectTrigger>
|
|
<SelectContent>
|
|
{capabilities?.supported_layouts?.map(layout => (
|
|
<SelectItem key={layout} value={layout}>
|
|
{layout.replace('_', ' ').replace(/\b\w/g, l => l.toUpperCase())}
|
|
</SelectItem>
|
|
))}
|
|
</SelectContent>
|
|
</Select>
|
|
</div>
|
|
|
|
<div className="space-y-2">
|
|
<label className="text-sm font-medium">Clustering Algorithm</label>
|
|
<Select value={clusteringAlgorithm} onValueChange={setClusteringAlgorithm}>
|
|
<SelectTrigger>
|
|
<SelectValue />
|
|
</SelectTrigger>
|
|
<SelectContent>
|
|
{capabilities?.supported_clustering?.map(clustering => (
|
|
<SelectItem key={clustering} value={clustering}>
|
|
{clustering.replace('_', ' ').replace(/\b\w/g, l => l.toUpperCase())}
|
|
</SelectItem>
|
|
))}
|
|
</SelectContent>
|
|
</Select>
|
|
</div>
|
|
</div>
|
|
|
|
<div className="space-y-4">
|
|
<div className="flex items-center justify-between">
|
|
<label className="text-sm font-medium">GPU Acceleration</label>
|
|
<Switch
|
|
checked={gpuAcceleration}
|
|
onCheckedChange={setGpuAcceleration}
|
|
disabled={!capabilities?.gpu_available}
|
|
/>
|
|
</div>
|
|
|
|
<div className="flex items-center justify-between">
|
|
<label className="text-sm font-medium">Compute Centrality</label>
|
|
<Switch
|
|
checked={computeCentrality}
|
|
onCheckedChange={setComputeCentrality}
|
|
/>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
|
|
{/* Action Buttons */}
|
|
<div className="flex gap-2">
|
|
<Button
|
|
onClick={processWithLocalGPU}
|
|
disabled={isProcessing || serviceHealth !== 'healthy'}
|
|
className="flex-1"
|
|
>
|
|
{isProcessing ? (
|
|
<>
|
|
<Loader2 className="w-4 h-4 mr-2 animate-spin" />
|
|
Processing on {gpuAcceleration ? 'GPU' : 'CPU'}...
|
|
</>
|
|
) : (
|
|
<>
|
|
<Zap className="w-4 h-4 mr-2" />
|
|
Process with Local GPU
|
|
</>
|
|
)}
|
|
</Button>
|
|
</div>
|
|
</CardContent>
|
|
</Card>
|
|
|
|
{/* Results */}
|
|
{processedData && (
|
|
<Card className="flex-1">
|
|
<CardHeader className="pb-3">
|
|
<CardTitle>Local GPU Processing Results</CardTitle>
|
|
</CardHeader>
|
|
<CardContent className="space-y-4">
|
|
{/* Statistics */}
|
|
<StatsDisplay stats={processedData.stats} />
|
|
|
|
{/* Visualization */}
|
|
<div className="w-full h-96 border rounded-lg overflow-hidden">
|
|
<div className="p-4 bg-blue-50 border-b">
|
|
<div className="text-sm font-medium text-blue-700">
|
|
🚀 Local GPU Visualization (cuGraph Powered)
|
|
</div>
|
|
<div className="text-xs text-blue-600 mt-1">
|
|
Processing completed locally with {processedData.gpu_processed ? 'GPU' : 'CPU'} acceleration
|
|
</div>
|
|
</div>
|
|
<div className="p-4">
|
|
<iframe
|
|
src="http://localhost:8081"
|
|
className="w-full h-80"
|
|
title="Local GPU Visualization"
|
|
style={{ border: 'none' }}
|
|
/>
|
|
</div>
|
|
</div>
|
|
</CardContent>
|
|
</Card>
|
|
)}
|
|
</div>
|
|
)
|
|
}
|