Skip to content

Commit

Permalink
Checkpoint visualization tweaks (#906)
Browse files Browse the repository at this point in the history 
Co-authored-by: Aakaash Meduri <[email protected]>
  • Loading branch information
@jeffmerrick @acashmoney
jeffmerrick and acashmoney authored Mar 7, 2024
1 parent 9a1ec9d commit 56e2379
Show file tree
Hide file tree
Showing 12 changed files with 3,237 additions and 245 deletions.
5 changes: 1 addition & 4 deletions frontend/app/experiments/ExperimentDetail.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,8 @@

import dayjs from "dayjs";
import relativeTime from "dayjs/plugin/relativeTime";
import { ChevronsUpDownIcon } from "lucide-react";
import Link from "next/link";
import React, { useEffect, useState } from "react";
import React, { useEffect } from "react";
import { useDispatch, useSelector } from "react-redux";

import { CopyToClipboard } from "@/components/shared/CopyToClipboard";
Expand All @@ -15,9 +14,7 @@ import { Alert } from "@/components/ui/alert";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
import { AppDispatch, flowDetailThunk, selectFlowDetail, selectFlowDetailError, selectFlowDetailLoading } from "@/lib/redux";
import { cn } from "@/lib/utils";

import { aggregateJobStatus, ExperimentStatus } from "./ExperimentStatus";
import JobDetail from "./JobDetail";
Expand Down
152 changes: 11 additions & 141 deletions frontend/app/experiments/JobDetail.tsx
View file
Original file line number Diff line number Diff line change
@@ -1,45 +1,23 @@
"use client";

import { getAccessToken } from "@privy-io/react-auth";
import { ColumnDef } from "@tanstack/react-table";
import backendUrl from "lib/backendUrl";
import { DownloadIcon } from "lucide-react";
import React, { useEffect, useState } from "react";
import { useDispatch, useSelector } from "react-redux";
import { ScatterChart, Scatter, XAxis, YAxis, CartesianGrid, Tooltip } from 'recharts';

import MolstarComponent from "@/components/Molstar";
import { CopyToClipboard } from "@/components/shared/CopyToClipboard";
import { TruncatedString } from "@/components/shared/TruncatedString";
import { Alert } from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
import { Card, CardContent, CardTitle } from "@/components/ui/card";
import { DataTable } from "@/components/ui/data-table";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { DataFile } from "@/lib/redux";

import LogViewer from "./LogViewer";
import MetricsVisualizer from "./MetricsVisualizer";

interface JobDetailProps {
jobID: number;
}

interface CustomTooltipProps {
active?: boolean;
payload?: any[];
label?: string;
}

interface CheckpointData {
cycle: number;
proposal: number;
plddt: number;
i_pae: number;
dim1: number;
dim2: number;
pdbFilePath: string;
}

export interface JobDetail {
ID: number | null;
BacalhauJobID: string;
Expand All @@ -57,11 +35,7 @@ export interface JobDetail {
export default function JobDetail({ jobID }: JobDetailProps) {
const [job, setJob] = useState({} as JobDetail);
const [loading, setLoading] = useState(false);
const [checkpoints, setCheckpoints] = useState([]);
const [plotData, setPlotData] = useState([]);
const [moleculeUrl, setMoleculeUrl] = useState('');
const [activeTab, setActiveTab] = useState('dashboard');

const [activeTab, setActiveTab] = useState("metrics");

interface File {
CID: string;
Expand Down Expand Up @@ -99,107 +73,27 @@ export default function JobDetail({ jobID }: JobDetailProps) {
}
};

fetchData();

fetch(`${backendUrl()}/checkpoints/${jobID}`)
.then(response => response.json())
.then(data => {
setCheckpoints(data);
})
.catch(error => console.error('Error fetching checkpoints:', error));

}, [jobID]);

useEffect(() => {
const fetchData = async () => {
try {
const checkpointResponse = await fetch(`${backendUrl()}/checkpoints/${jobID}`);
const checkpointData = await checkpointResponse.json();
setCheckpoints(checkpointData);

const plotDataResponse = await fetch(`${backendUrl()}/checkpoints/${jobID}/get-data`);
const plotData = await plotDataResponse.json();
setPlotData(plotData);
} catch (error) {
console.error('Error fetching data:', error);
}
};

if (job.State === "running") {
fetchData();
const intervalId = setInterval(fetchData, 5000);

return () => clearInterval(intervalId);
} else { //(job.State === "completed") {
} else {
fetchData();
}
}, [jobID, job.State]);

const handlePointClick = (data: CheckpointData) => {
console.log('Clicked point data:', data);
setMoleculeUrl(data.pdbFilePath);
console.log("set molecule url:", data.pdbFilePath);
// Switch to the visualize tab
console.log(activeTab);
setActiveTab('visualize');
console.log(activeTab);
};

const CustomTooltip: React.FC<CustomTooltipProps> = ({ active, payload }) => {
if (active && payload && payload.length) {
const data = payload[0].payload;
const keysToShow = Object.keys(data).filter(key => key !== 'pdbFilePath');

return (
<div className="bg-white p-3 border rounded shadow-lg">
{keysToShow.map((key) => (
<p key={key}>
{key}: {data[key]}
</p>
))}
</div>
);
}

return null;
};

return (
<Tabs value={activeTab} onValueChange={setActiveTab} className="w-full @container ">
<Tabs value={activeTab} onValueChange={setActiveTab} className="w-full @container">
<TabsList className="justify-start w-full px-6 pt-0 rounded-t-none">
<TabsTrigger value="dashboard">Dashboard</TabsTrigger>
<TabsTrigger value="visualize">Visualize</TabsTrigger>
<TabsTrigger value="metrics">Metrics</TabsTrigger>
<TabsTrigger value="parameters">Parameters</TabsTrigger>
<TabsTrigger value="outputs">Outputs</TabsTrigger>
<TabsTrigger value="inputs">Inputs</TabsTrigger>
<TabsTrigger value="logs">Logs</TabsTrigger>
</TabsList>
<TabsContent value="dashboard">
<div style={{ display: 'flex', alignItems: 'center', marginTop: '15px', justifyContent: 'space-around' }}>
<div>
<div className="font-heading" style={{ marginLeft: '70px', textAlign: 'center' }}>Metrics Space</div>
<ScatterChart width={730} height={250} margin={{ top: 20, right: 20, bottom: 20, left: 20}}>
<CartesianGrid />
<XAxis type="number" dataKey="plddt" name="plddt" label={{ value: 'plddt', position: 'insideBottom', offset: -15, dx: 0 }}/>
<YAxis type="number" dataKey="i_pae" name="i_pae" label={{ value: 'i_pae', angle: -90, position: 'insideLeft', dx: 0, dy: 15 }}/>
{/* <Tooltip cursor={{ strokeDasharray: '3 3' }} /> */}
<Tooltip content={<CustomTooltip />} cursor={{ strokeDasharray: '3 3' }} />
<Scatter name="Checkpoints" data={plotData} fill="#8884d8" onClick={handlePointClick} />
</ScatterChart>
</div>
<div style={{ maxWidth: '300px', fontSize: '0.75rem', color: 'grey', textAlign: 'left' }}>
<p><b>plddt:</b> larger value indicates higher confidence in the predicted local structure</p>
<p><b>i_pae:</b> larger value indicates higher confidence in the predicted interface residue distance</p>
<p style = {{fontSize: '0.80rem', color: 'black'}}>Click on a datapoint to visualize</p>
</div>
</div>
<CheckpointsList checkpoints={checkpoints} />
</TabsContent>
<TabsContent value="visualize">
<MolstarComponent
moleculeUrl={moleculeUrl}
customDataFormat="pdb"
/>
<TabsContent value="metrics">
<MetricsVisualizer job={job} />
</TabsContent>
<TabsContent value="parameters" className="px-6 pt-0">
{Object.entries(job.Inputs || {}).map(([key, val]) => (
Expand All @@ -217,37 +111,13 @@ export default function JobDetail({ jobID }: JobDetailProps) {
</TabsContent>
<TabsContent value="logs">
<div className="w-full">
<LogViewer jobID={job.BacalhauJobID} />
<LogViewer bacalhauJobID={job.BacalhauJobID} />
</div>
</TabsContent>
</Tabs>
);
}

function CheckpointsList({ checkpoints }: { checkpoints: Array<{ fileName: string, url: string }> }) {
const safeCheckpoints = checkpoints || [];
return (
<div>
{safeCheckpoints.length > 0 ? (
safeCheckpoints.map((checkpoint, index) => (
<div key={index} className="flex items-center justify-between px-6 py-2 text-xs border-b border-border/50 last:border-none">
<div>
<span className="text-accent">{checkpoint.fileName}</span>
</div>
<Button size="icon" variant="outline" asChild>
<a href={checkpoint.url} download target="_blank" rel="noopener noreferrer">
<DownloadIcon />
</a>
</Button>
</div>
))
) : (
<p>Checkpoints will start appearing soon... Please check back later</p>
)}
</div>
);
}

function FileList({ files }: { files: DataFile[] }) {
const handleDownload = async (file: DataFile) => {
try {
Expand Down Expand Up @@ -297,11 +167,11 @@ function FileList({ files }: { files: DataFile[] }) {
</a>
</Button>
</div>
))}
))}
</>
) : (
<div className="px-6 py-5">No files found.</div>
)}
</div>
</div>
);
}
10 changes: 5 additions & 5 deletions frontend/app/experiments/LogViewer.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,18 +6,18 @@ import React, { useEffect, useState } from "react";
import { CopyToClipboard } from "@/components/shared/CopyToClipboard";
import { ScrollArea, ScrollBar } from "@/components/ui/scroll-area";

const LogViewer = ({ jobID }: { jobID: string }) => {
const LogViewer = ({ bacalhauJobID }: { bacalhauJobID: string }) => {
const [logs, setLogs] = useState("");

useEffect(() => {
if (jobID) {
setLogs(`Connecting to stream with Bacalhau Job Id ${jobID}`);
if (bacalhauJobID) {
setLogs(`Connecting to stream with Bacalhau Job Id ${bacalhauJobID}`);

let formattedBackendUrl = backendUrl().replace("http://", "").replace("https://", "");
let wsProtocol = backendUrl().startsWith("https://") ? "wss" : "ws";

console.log(formattedBackendUrl);
const ws = new WebSocket(`${wsProtocol}://${formattedBackendUrl}/jobs/${jobID}/logs`);
const ws = new WebSocket(`${wsProtocol}://${formattedBackendUrl}/jobs/${bacalhauJobID}/logs`);

ws.onopen = () => {
console.log("connected");
Expand All @@ -39,7 +39,7 @@ const LogViewer = ({ jobID }: { jobID: string }) => {
} else {
setLogs(`Logs will stream after job is sent to the Bacalhau network`);
}
}, [jobID]);
}, [bacalhauJobID]);

return (
<div className="relative">
Expand Down
Loading

0 comments on commit 56e2379

Please sign in to comment.