mirror of
https://github.com/NVIDIA/dgx-spark-playbooks.git
synced 2026-04-24 10:53:52 +00:00
525 lines
19 KiB
TypeScript
525 lines
19 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 { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'
|
||
import { Loader2, Zap, Activity, Cpu, Cloud, ExternalLink, Settings } 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 UnifiedGPUViewerProps {
|
||
graphData: GraphData
|
||
onError?: (error: Error) => void
|
||
}
|
||
|
||
interface ProcessingCapabilities {
|
||
processing_modes: {
|
||
pygraphistry_cloud: { available: boolean, description: string }
|
||
local_gpu: { available: boolean, description: string }
|
||
local_cpu: { available: boolean, description: string }
|
||
}
|
||
has_rapids: boolean
|
||
has_torch_geometric: boolean
|
||
gpu_available: boolean
|
||
}
|
||
|
||
interface ProcessedData {
|
||
processed_nodes: any[]
|
||
processed_edges: any[]
|
||
processing_mode: string
|
||
embed_url?: string
|
||
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
|
||
has_embed_url?: boolean
|
||
layout_computed?: boolean
|
||
clusters_computed?: boolean
|
||
centrality_computed?: boolean
|
||
}
|
||
timestamp: string
|
||
}
|
||
|
||
export function UnifiedGPUViewer({ graphData, onError }: UnifiedGPUViewerProps) {
|
||
const [isProcessing, setIsProcessing] = useState(false)
|
||
const [processedData, setProcessedData] = useState<ProcessedData | null>(null)
|
||
const [capabilities, setCapabilities] = useState<ProcessingCapabilities | null>(null)
|
||
const [serviceHealth, setServiceHealth] = useState<'unknown' | 'healthy' | 'error'>('unknown')
|
||
|
||
// Processing mode and options
|
||
const [processingMode, setProcessingMode] = useState<'pygraphistry_cloud' | 'local_gpu' | 'local_cpu'>('pygraphistry_cloud')
|
||
|
||
// PyGraphistry options
|
||
const [layoutType, setLayoutType] = useState('force')
|
||
const [gpuAcceleration, setGpuAcceleration] = useState(true)
|
||
const [clustering, setClustering] = useState(false)
|
||
|
||
// Local GPU options
|
||
const [layoutAlgorithm, setLayoutAlgorithm] = useState('force_atlas2')
|
||
const [clusteringAlgorithm, setClusteringAlgorithm] = useState('leiden')
|
||
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/unified-gpu/health')
|
||
if (response.ok) {
|
||
setServiceHealth('healthy')
|
||
} else {
|
||
setServiceHealth('error')
|
||
}
|
||
} catch (error) {
|
||
console.error('Unified service health check failed:', error)
|
||
setServiceHealth('error')
|
||
}
|
||
}
|
||
|
||
const getCapabilities = async () => {
|
||
try {
|
||
const response = await fetch('/api/unified-gpu/capabilities')
|
||
if (response.ok) {
|
||
const caps = await response.json()
|
||
setCapabilities(caps)
|
||
|
||
// Set default mode based on capabilities
|
||
if (caps.processing_modes.pygraphistry_cloud.available) {
|
||
setProcessingMode('pygraphistry_cloud')
|
||
} else if (caps.processing_modes.local_gpu.available) {
|
||
setProcessingMode('local_gpu')
|
||
} else {
|
||
setProcessingMode('local_cpu')
|
||
}
|
||
}
|
||
} catch (error) {
|
||
console.error('Failed to get capabilities:', error)
|
||
}
|
||
}
|
||
|
||
const setupWebSocket = () => {
|
||
try {
|
||
const ws = new WebSocket('ws://localhost:8080/ws')
|
||
|
||
ws.onopen = () => {
|
||
console.log('WebSocket connected to unified service')
|
||
}
|
||
|
||
ws.onmessage = (event) => {
|
||
try {
|
||
const message = JSON.parse(event.data)
|
||
if (message.type === 'graph_processed') {
|
||
setProcessedData(message.data)
|
||
setIsProcessing(false)
|
||
|
||
toast({
|
||
title: "Processing Complete",
|
||
description: `Processed ${message.data.stats.node_count} nodes with ${message.data.processing_mode}`,
|
||
})
|
||
}
|
||
} catch (error) {
|
||
console.error('WebSocket message error:', error)
|
||
}
|
||
}
|
||
|
||
wsRef.current = ws
|
||
} catch (error) {
|
||
console.error('WebSocket setup failed:', error)
|
||
}
|
||
}
|
||
|
||
const processGraph = 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
|
||
},
|
||
processing_mode: processingMode,
|
||
|
||
// PyGraphistry options
|
||
layout_type: layoutType,
|
||
gpu_acceleration: gpuAcceleration,
|
||
clustering: clustering,
|
||
|
||
// Local GPU options
|
||
layout_algorithm: layoutAlgorithm,
|
||
clustering_algorithm: clusteringAlgorithm,
|
||
compute_centrality: computeCentrality
|
||
}
|
||
|
||
const response = await fetch('/api/unified-gpu/visualize', {
|
||
method: 'POST',
|
||
headers: {
|
||
'Content-Type': 'application/json'
|
||
},
|
||
body: JSON.stringify(requestData)
|
||
})
|
||
|
||
if (!response.ok) {
|
||
throw new Error(`Processing failed: ${response.statusText}`)
|
||
}
|
||
|
||
// Result will come via WebSocket or direct response
|
||
const result = await response.json()
|
||
if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
|
||
// Handle direct response if WebSocket not available
|
||
setProcessedData(result)
|
||
setIsProcessing(false)
|
||
|
||
toast({
|
||
title: "Processing Complete",
|
||
description: `Processed ${result.stats.node_count} nodes with ${result.processing_mode}`,
|
||
})
|
||
}
|
||
|
||
} catch (error) {
|
||
console.error('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>
|
||
Unified Service: {serviceHealth === 'healthy' ? 'Connected' :
|
||
serviceHealth === 'error' ? 'Disconnected' : 'Checking...'}
|
||
</span>
|
||
</div>
|
||
)
|
||
|
||
const ProcessingModeSelector = () => (
|
||
<div className="space-y-2">
|
||
<label className="text-sm font-medium">Processing Mode</label>
|
||
<Select value={processingMode} onValueChange={(value: any) => setProcessingMode(value)}>
|
||
<SelectTrigger>
|
||
<SelectValue />
|
||
</SelectTrigger>
|
||
<SelectContent>
|
||
{capabilities?.processing_modes.pygraphistry_cloud.available && (
|
||
<SelectItem value="pygraphistry_cloud">
|
||
<div className="flex items-center gap-2">
|
||
<Cloud className="w-4 h-4" />
|
||
PyGraphistry Cloud
|
||
</div>
|
||
</SelectItem>
|
||
)}
|
||
{capabilities?.processing_modes.local_gpu.available && (
|
||
<SelectItem value="local_gpu">
|
||
<div className="flex items-center gap-2">
|
||
<Zap className="w-4 h-4" />
|
||
Local GPU (cuGraph)
|
||
</div>
|
||
</SelectItem>
|
||
)}
|
||
<SelectItem value="local_cpu">
|
||
<div className="flex items-center gap-2">
|
||
<Cpu className="w-4 h-4" />
|
||
Local CPU
|
||
</div>
|
||
</SelectItem>
|
||
</SelectContent>
|
||
</Select>
|
||
<p className="text-xs text-muted-foreground">
|
||
{capabilities?.processing_modes[processingMode]?.description}
|
||
</p>
|
||
</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'} accelerated</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" />
|
||
Unified GPU Visualization
|
||
</span>
|
||
<ServiceStatus />
|
||
</CardTitle>
|
||
</CardHeader>
|
||
<CardContent className="space-y-4">
|
||
{/* Processing Mode Selection */}
|
||
<ProcessingModeSelector />
|
||
|
||
{/* Processing Options */}
|
||
<Tabs value={processingMode} className="w-full">
|
||
<TabsList className="grid w-full grid-cols-3">
|
||
<TabsTrigger value="pygraphistry_cloud" disabled={!capabilities?.processing_modes.pygraphistry_cloud.available}>
|
||
Cloud
|
||
</TabsTrigger>
|
||
<TabsTrigger value="local_gpu" disabled={!capabilities?.processing_modes.local_gpu.available}>
|
||
Local GPU
|
||
</TabsTrigger>
|
||
<TabsTrigger value="local_cpu">
|
||
Local CPU
|
||
</TabsTrigger>
|
||
</TabsList>
|
||
|
||
<TabsContent value="pygraphistry_cloud" className="space-y-4">
|
||
<div className="grid grid-cols-1 md:grid-cols-2 gap-4">
|
||
<div className="space-y-2">
|
||
<label className="text-sm font-medium">Layout Type</label>
|
||
<Select value={layoutType} onValueChange={setLayoutType}>
|
||
<SelectTrigger>
|
||
<SelectValue />
|
||
</SelectTrigger>
|
||
<SelectContent>
|
||
<SelectItem value="force">Force Directed</SelectItem>
|
||
<SelectItem value="circular">Circular</SelectItem>
|
||
<SelectItem value="hierarchical">Hierarchical</SelectItem>
|
||
</SelectContent>
|
||
</Select>
|
||
</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} />
|
||
</div>
|
||
<div className="flex items-center justify-between">
|
||
<label className="text-sm font-medium">Auto-Clustering</label>
|
||
<Switch checked={clustering} onCheckedChange={setClustering} />
|
||
</div>
|
||
</div>
|
||
</div>
|
||
</TabsContent>
|
||
|
||
<TabsContent value="local_gpu" className="space-y-4">
|
||
<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>
|
||
<SelectItem value="force_atlas2">Force Atlas 2</SelectItem>
|
||
<SelectItem value="spectral">Spectral Layout</SelectItem>
|
||
<SelectItem value="fruchterman_reingold">Fruchterman-Reingold</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>
|
||
<SelectItem value="leiden">Leiden</SelectItem>
|
||
<SelectItem value="louvain">Louvain</SelectItem>
|
||
<SelectItem value="spectral">Spectral Clustering</SelectItem>
|
||
</SelectContent>
|
||
</Select>
|
||
</div>
|
||
</div>
|
||
<div className="space-y-4">
|
||
<div className="flex items-center justify-between">
|
||
<label className="text-sm font-medium">Compute Centrality</label>
|
||
<Switch checked={computeCentrality} onCheckedChange={setComputeCentrality} />
|
||
</div>
|
||
{capabilities && (
|
||
<div className="p-2 bg-muted rounded text-xs">
|
||
RAPIDS cuGraph: {capabilities.has_rapids ? '✓ Available' : '✗ Not Available'}
|
||
</div>
|
||
)}
|
||
</div>
|
||
</div>
|
||
</TabsContent>
|
||
|
||
<TabsContent value="local_cpu" className="space-y-4">
|
||
<div className="p-4 bg-muted rounded-lg text-center">
|
||
<p className="text-sm text-muted-foreground">
|
||
CPU fallback mode - basic processing without GPU acceleration
|
||
</p>
|
||
</div>
|
||
</TabsContent>
|
||
</Tabs>
|
||
|
||
{/* Action Button */}
|
||
<Button
|
||
onClick={processGraph}
|
||
disabled={isProcessing || serviceHealth !== 'healthy'}
|
||
className="w-full"
|
||
>
|
||
{isProcessing ? (
|
||
<>
|
||
<Loader2 className="w-4 h-4 mr-2 animate-spin" />
|
||
Processing with {processingMode.replace('_', ' ')}...
|
||
</>
|
||
) : (
|
||
<>
|
||
<Zap className="w-4 h-4 mr-2" />
|
||
Process Graph
|
||
</>
|
||
)}
|
||
</Button>
|
||
</CardContent>
|
||
</Card>
|
||
|
||
{/* Results */}
|
||
{processedData && (
|
||
<Card className="flex-1">
|
||
<CardHeader className="pb-3">
|
||
<div className="flex items-center justify-between">
|
||
<CardTitle>Processing Results</CardTitle>
|
||
<Badge variant="outline">
|
||
{processedData.processing_mode.replace('_', ' ')}
|
||
</Badge>
|
||
</div>
|
||
</CardHeader>
|
||
<CardContent className="space-y-4">
|
||
{/* Statistics */}
|
||
<StatsDisplay stats={processedData.stats} />
|
||
|
||
{/* Visualization */}
|
||
<div className="w-full h-96 border rounded-lg overflow-hidden">
|
||
{processedData.embed_url ? (
|
||
// PyGraphistry Cloud embed
|
||
<>
|
||
<div className="p-3 bg-blue-50 border-b flex items-center justify-between">
|
||
<div>
|
||
<div className="text-sm font-medium text-blue-700">
|
||
☁️ PyGraphistry Cloud Visualization
|
||
</div>
|
||
<div className="text-xs text-blue-600">
|
||
Interactive GPU-accelerated visualization
|
||
</div>
|
||
</div>
|
||
<Button
|
||
variant="outline"
|
||
size="sm"
|
||
onClick={() => window.open(processedData.embed_url, '_blank')}
|
||
>
|
||
<ExternalLink className="w-4 h-4 mr-1" />
|
||
Open
|
||
</Button>
|
||
</div>
|
||
<iframe
|
||
src={processedData.embed_url}
|
||
className="w-full h-80"
|
||
title="PyGraphistry Visualization"
|
||
style={{ border: 'none' }}
|
||
/>
|
||
</>
|
||
) : (
|
||
// Local processing result
|
||
<div className="p-4 bg-green-50 border-b">
|
||
<div className="text-sm font-medium text-green-700">
|
||
🚀 Local {processedData.gpu_processed ? 'GPU' : 'CPU'} Processing Complete
|
||
</div>
|
||
<div className="text-xs text-green-600 mt-1">
|
||
{processedData.stats.layout_computed && '✓ Layout computed '}
|
||
{processedData.stats.clusters_computed && '✓ Clusters detected '}
|
||
{processedData.stats.centrality_computed && '✓ Centrality computed'}
|
||
</div>
|
||
</div>
|
||
)}
|
||
</div>
|
||
</CardContent>
|
||
</Card>
|
||
)}
|
||
</div>
|
||
)
|
||
}
|