diff --git a/frontend/3dmol.d.ts b/frontend/3dmol.d.ts new file mode 100644 index 00000000..8fdf2988 --- /dev/null +++ b/frontend/3dmol.d.ts @@ -0,0 +1,4 @@ +declare module '3dmol/build/3Dmol.js' { + const $3Dmol: any; + export default $3Dmol; +} \ No newline at end of file diff --git a/frontend/app/experiments/(experiment)/(forms)/NewExperimentForm.tsx b/frontend/app/experiments/(experiment)/(forms)/NewExperimentForm.tsx index 3959b999..793d4393 100644 --- a/frontend/app/experiments/(experiment)/(forms)/NewExperimentForm.tsx +++ b/frontend/app/experiments/(experiment)/(forms)/NewExperimentForm.tsx @@ -18,9 +18,17 @@ import { addExperimentThunk, addExperimentWithCheckoutThunk, AppDispatch, experi import { createExperiment } from "@/lib/redux/slices/experimentAddSlice/asyncActions"; import { DynamicArrayField } from "./DynamicArrayField"; +import ThreeDMolViewer from "./ThreeDMolViewer"; import { generateDefaultValues, generateSchema } from "./formGenerator"; import { groupInputs, transformJson } from "./formUtils"; +interface FormValues { + name: string; + model: string; + pdbData?: string; + selectedResidues?: string[]; +} + export default function NewExperimentForm({ task }: { task: any }) { const dispatch = useDispatch(); const router = useRouter(); @@ -30,6 +38,9 @@ export default function NewExperimentForm({ task }: { task: any }) { const walletAddress = user?.wallet?.address; const userTier = useSelector(selectUserTier); const isUserSubscribed = useSelector(selectIsUserSubscribed); + const handleViewerSubmit = (data: { pdb: File | null; binderLength: number; hotspots: string }) => { + console.log('Data from viewer:', data); +}; useEffect(() => { dispatch(refreshUserDataThunk()); @@ -49,12 +60,12 @@ export default function NewExperimentForm({ task }: { task: any }) { const groupedInputs = groupInputs(model.ModelJson?.inputs); const formSchema = generateSchema(model.ModelJson?.inputs); - const defaultValues = generateDefaultValues(model.ModelJson?.inputs, task, model); + // const defaultValues = generateDefaultValues(model.ModelJson?.inputs, task, model); - const form = useForm>({ - resolver: zodResolver(formSchema), - defaultValues: defaultValues, - }); + const form = useForm({ + resolver: zodResolver(formSchema), // Make sure formSchema reflects these types + defaultValues: generateDefaultValues(model.ModelJson?.inputs, task, model), +}); // Watch all form values const watchedValues = form.watch(); @@ -84,7 +95,9 @@ export default function NewExperimentForm({ task }: { task: any }) { if (userTier === 'Free') { // For free tier users, directly create the experiment regardless of subscription status const response = await dispatch(addExperimentThunk(transformedPayload)).unwrap(); - if (response && response.ID) { + if (response.redirectUrl) { + router.push(response.redirectUrl); + } else if (response && response.ID) { console.log("Experiment created", response); router.push(`/experiments/${response.ID}`, { scroll: false }); dispatch(experimentListThunk(walletAddress)); @@ -145,6 +158,12 @@ export default function NewExperimentForm({ task }: { task: any }) { + + + + + + {!!groupedInputs?.standard && ( <> diff --git a/frontend/app/experiments/(experiment)/(forms)/RerunExperimentForm.tsx b/frontend/app/experiments/(experiment)/(forms)/RerunExperimentForm.tsx index b32c8032..53dd456c 100644 --- a/frontend/app/experiments/(experiment)/(forms)/RerunExperimentForm.tsx +++ b/frontend/app/experiments/(experiment)/(forms)/RerunExperimentForm.tsx @@ -66,7 +66,9 @@ export default function RerunExperimentForm() { try { if (userTier === 'Free' || (userTier === 'Paid' && isUserSubscribed)) { const response = await addJobToExperiment(experimentID, transformedPayload); - if (response && response.ID) { + if (response.redirectUrl) { + router.push(response.redirectUrl); + } else if (response && response.ID) { console.log("Job added to experiment", response); dispatch(experimentDetailThunk(response.ID)); dispatch(experimentListThunk(walletAddress)); diff --git a/frontend/app/experiments/(experiment)/(forms)/ThreeDMolViewer.tsx b/frontend/app/experiments/(experiment)/(forms)/ThreeDMolViewer.tsx new file mode 100644 index 00000000..e616cb28 --- /dev/null +++ b/frontend/app/experiments/(experiment)/(forms)/ThreeDMolViewer.tsx @@ -0,0 +1,195 @@ +import React, { useState, useEffect, useRef } from 'react'; +import axios from 'axios'; + +interface AtomSpec { + chain: string; + resi: number; +} + +interface ThreeDMolViewerProps { + onSubmit: (data: { pdb: File | null; binderLength: number; hotspots: string }) => void; +} + +const ThreeDMolViewer: React.FC = ({ onSubmit }) => { + const [viewer, setViewer] = useState(null); + const [selectedResidues, setSelectedResidues] = useState>({}); + const [binderLength, setBinderLength] = useState(90); + const fileInput = useRef(null); + const [data, setData] = useState(null); + + useEffect(() => { + (async () => { + const module = await import('3dmol/build/3Dmol.js'); + const $3Dmol = module.default ? module.default : module; + const element = document.getElementById('container-01') as HTMLElement; + const config = { backgroundColor: 'white' }; + const viewer = $3Dmol.createViewer(element, config); + setViewer(viewer); + })(); + }, []); + + useEffect(() => { + if (viewer) { + updateStyles(); + } + }, [selectedResidues]); + + const handleFileUpload = (event: React.ChangeEvent) => { + const file = event.target.files?.[0]; + if (file && viewer) { + const reader = new FileReader(); + reader.onload = function(e) { + const result = e.target?.result as string; + setData(result); + viewer.addModel(result, "pdb"); + updateStyles(); + viewer.getModel().setClickable({}, true, (atom: AtomSpec) => { + const residueKey = `${atom.chain}${atom.resi}`; + setSelectedResidues(prev => { + const newResidues = { ...prev }; + if (newResidues[residueKey]) { + delete newResidues[residueKey]; // Remove the residue from selected if it's already selected + } else { + newResidues[residueKey] = true; // Add the residue as selected if it's not already + } + return newResidues; + }); + }); + viewer.render(); + }; + reader.readAsText(file); + } + }; + + const updateStyles = () => { + if (!viewer) return; + + const allChains = new Set(); + const chainsWithSelectedResidues = new Set(); + + // First, reset styles for all residues + viewer.selectedAtoms({}).forEach((atom: { chain: unknown; resi: any; }) => { + allChains.add(atom.chain); + const residueKey = `${atom.chain}${atom.resi}`; + if (selectedResidues[residueKey]) { + chainsWithSelectedResidues.add(atom.chain); + } + }); + + if (Object.keys(selectedResidues).length === 0) { + // No residues selected, set all chains to fully visible + viewer.setStyle({}, { cartoon: { color: 'grey', opacity: 1.0 } }); + } else { + // Set default style for all chains with higher transparency + viewer.setStyle({}, { cartoon: { color: 'grey', opacity: 0.2 } }); + + // Set style for chains with selected residues to fully visible + chainsWithSelectedResidues.forEach(chain => { + viewer.setStyle({ chain }, { cartoon: { color: 'grey', opacity: 1.0 } }); + }); + + // Set style for selected residues + Object.keys(selectedResidues).forEach(residueKey => { + const chain = residueKey[0]; + const resi = residueKey.slice(1); + viewer.setStyle({ chain, resi }, { cartoon: { color: 'red', opacity: 1.0 } }); + }); + } + + viewer.render(); + }; + + useEffect(() => { + if (viewer) { + const loadAndMakeClickable = async () => { + const data = await axios.get('https://files.rcsb.org/download/1UBQ.pdb').then(res => res.data); + viewer.addModel(data, "pdb"); + viewer.getModel().setClickable({}, (atom: { chain: any; resi: any; }) => { + const key = `${atom.chain}:${atom.resi}`; + setSelectedResidues(prev => ({ + ...prev, + [key]: !prev[key] + })); + }); + updateStyles(); + }; + loadAndMakeClickable(); + } + }, [viewer]); + + + + const formatSelectedResidues = () => { + return Object.keys(selectedResidues).map((residueKey) => { + const chain = residueKey[0]; + const index = residueKey.slice(1); + return `${chain}${index}`; + }).join(', '); + }; + + const handleBinderLengthChange = (event: React.ChangeEvent) => { + setBinderLength(Number(event.target.value)); + }; + + const handleSubmit = async () => { + const pdbFile = fileInput.current?.files?.[0] ?? null; + const hotspots = formatSelectedResidues(); + onSubmit({ pdb: pdbFile, binderLength, hotspots }); + }; + + const handleResidueInputChange = (event: React.ChangeEvent) => { + const input = event.target.value; + const newSelectedResidues: Record = {}; + + const residues = input.split(',').map(res => res.trim()).filter(res => res); + + residues.forEach(residue => { + const match = residue.match(/^([a-zA-Z]+)(\d+)$/); + + if (match) { + const chain = match[1].toUpperCase(); + const resi = parseInt(match[2], 10); + if (!isNaN(resi)) { + const key = `${chain}${resi}`; + newSelectedResidues[key] = true; + } + } + }); + + setSelectedResidues(newSelectedResidues); + updateStyles(); + }; + + return ( +
+
+
+
+ + +
+ + + {binderLength} + +
+
+ ); +}; + +export default ThreeDMolViewer; \ No newline at end of file diff --git a/frontend/app/experiments/(experiment)/(forms)/ThreeDMolViewer_backup.tsx b/frontend/app/experiments/(experiment)/(forms)/ThreeDMolViewer_backup.tsx new file mode 100644 index 00000000..7c2de186 --- /dev/null +++ b/frontend/app/experiments/(experiment)/(forms)/ThreeDMolViewer_backup.tsx @@ -0,0 +1,252 @@ +import React, { useState, useEffect, useRef } from 'react'; +import axios from 'axios'; +import {FileSelect} from '@/components/shared/FileSelect'; + +interface AtomSpec { + chain: string; + resi: number; +} + +interface ThreeDMolViewerProps { + onSubmit: (data: { pdb: File | null; binderLength: number; hotspots: string }) => void; +} + +const ThreeDMolViewer: React.FC = ({ onSubmit }) => { + const [viewer, setViewer] = useState(null); + const [selectedResidues, setSelectedResidues] = useState>({}); + const [binderLength, setBinderLength] = useState(90); + const fileInput = useRef(null); + const [pdbData, setPdbData] = useState(null); + const [data, setData] = useState(null); + const [fileName, setFileName] = useState(''); + + useEffect(() => { + (async () => { + const module = await import('3dmol/build/3Dmol.js'); + const $3Dmol = module.default ? module.default : module; + const element = document.getElementById('container-01') as HTMLElement; + const config = { backgroundColor: 'white' }; + const viewer = $3Dmol.createViewer(element, config); + setViewer(viewer); + })(); + }, []); + + useEffect(() => { + const loadPDBFile = async () => { + if (viewer && fileName) { // Ensure there's a filename and viewer instance + try { + const response = await axios.get(fileName, { responseType: 'text' }); + viewer.addModel(response.data, "pdb"); + viewer.zoomTo(); // Adjust the camera to the new model + viewer.render(); + updateStyles(); // Update visual styles if needed + } catch (error) { + console.error("Error loading PDB file:", error); + } + } + }; + loadPDBFile(); + }, [viewer, fileName]); // Depend on fileName to refetch when it changes + + useEffect(() => { + if (viewer) { + updateStyles(); + } + }, [selectedResidues]); + + const handleFileUpload = (fileName: string) => { + // Access the file object directly from the input reference + setFileName(fileName); + const file = fileInput.current?.files?.[0]; + if (file && viewer && file.name === fileName) { + const reader = new FileReader(); + reader.onload = function(e) { + const pdbData = e.target?.result as string; + setData(pdbData); + viewer.addModel(pdbData, "pdb"); + viewer.zoomTo(); // Ensure the viewer adjusts to the loaded model + viewer.render(); + // Enable clicking to select hotspots only after model is loaded + viewer.getModel().setClickable({}, true, (atom: AtomSpec) => { + const residueKey = `${atom.chain}${atom.resi}`; + setSelectedResidues(prev => ({ + ...prev, + [residueKey]: !prev[residueKey] + })); + }); + }; + reader.readAsText(file); + } else { + console.error("File selected does not match or viewer is not initialized"); + } + }; + + // Ensure fileInput is accessible and used correctly +// const handleFileSelection = (fileData: any, fileName: string) => { +// // Trigger native file input click +// setData(fileData); +// setFileName(fileName); // Assuming you keep a fileName state +// if (viewer) { +// viewer.addModel(fileData, "pdb"); +// viewer.render(); +// } +// }; + +// const handleFileSelection = (selectedFileName: string) => { +// setFileName(selectedFileName); +// }; + +// Handling file change from the native file input +const handleFileChange = () => { + const file = fileInput.current?.files?.[0]; + if (file && viewer) { + const reader = new FileReader(); + reader.onload = (e) => { + const result = e.target?.result as string; + setData(result); + viewer.addModel(result, "pdb"); + viewer.render(); + updateStyles(); + }; + reader.readAsText(file); + } +}; + + const updateStyles = () => { + if (!viewer) return; + + const allChains = new Set(); + const chainsWithSelectedResidues = new Set(); + + // First, reset styles for all residues + viewer.selectedAtoms({}).forEach((atom: { chain: unknown; resi: any; }) => { + allChains.add(atom.chain); + const residueKey = `${atom.chain}${atom.resi}`; + if (selectedResidues[residueKey]) { + chainsWithSelectedResidues.add(atom.chain); + } + }); + + if (Object.keys(selectedResidues).length === 0) { + // No residues selected, set all chains to fully visible + viewer.setStyle({}, { cartoon: { color: 'grey', opacity: 1.0 } }); + } else { + // Set default style for all chains with higher transparency + viewer.setStyle({}, { cartoon: { color: 'grey', opacity: 0.2 } }); + + // Set style for chains with selected residues to fully visible + chainsWithSelectedResidues.forEach(chain => { + viewer.setStyle({ chain }, { cartoon: { color: 'grey', opacity: 1.0 } }); + }); + + // Set style for selected residues + Object.keys(selectedResidues).forEach(residueKey => { + const chain = residueKey[0]; + const resi = residueKey.slice(1); + viewer.setStyle({ chain, resi }, { cartoon: { color: 'red', opacity: 1.0 } }); + }); + } + + viewer.render(); + }; + + useEffect(() => { + if (viewer) { + const loadAndMakeClickable = async () => { + const data = await axios.get('https://files.rcsb.org/download/1UBQ.pdb').then(res => res.data); + viewer.addModel(data, "pdb"); + viewer.getModel().setClickable({}, (atom: { chain: any; resi: any; }) => { + const key = `${atom.chain}:${atom.resi}`; + setSelectedResidues(prev => ({ + ...prev, + [key]: !prev[key] + })); + }); + updateStyles(); + }; + loadAndMakeClickable(); + } + }, [viewer]); + + + + const formatSelectedResidues = () => { + return Object.keys(selectedResidues).map((residueKey) => { + const chain = residueKey[0]; + const index = residueKey.slice(1); + return `${chain}${index}`; + }).join(', '); + }; + + const handleBinderLengthChange = (event: React.ChangeEvent) => { + setBinderLength(Number(event.target.value)); + }; + + const handleSubmit = async () => { + const pdbFile = fileInput.current?.files?.[0] ?? null; + const hotspots = formatSelectedResidues(); + onSubmit({ pdb: pdbFile, binderLength, hotspots }); + }; + + const handleResidueInputChange = (event: React.ChangeEvent) => { + const input = event.target.value; + const newSelectedResidues: Record = {}; + + const residues = input.split(',').map(res => res.trim()).filter(res => res); + + residues.forEach(residue => { + const match = residue.match(/^([a-zA-Z]+)(\d+)$/); + + if (match) { + const chain = match[1].toUpperCase(); + const resi = parseInt(match[2], 10); + if (!isNaN(resi)) { + const key = `${chain}${resi}`; + newSelectedResidues[key] = true; + } + } + }); + + setSelectedResidues(newSelectedResidues); + updateStyles(); + }; + + return ( +
+
+
+
+ + + +
+ + + {binderLength} + +
+
+ ); +}; + +export default ThreeDMolViewer; \ No newline at end of file diff --git a/frontend/app/experiments/(experiment)/(forms)/formGenerator.ts b/frontend/app/experiments/(experiment)/(forms)/formGenerator.ts index 165d5232..3cb770ad 100644 --- a/frontend/app/experiments/(experiment)/(forms)/formGenerator.ts +++ b/frontend/app/experiments/(experiment)/(forms)/formGenerator.ts @@ -111,7 +111,7 @@ export function generateRerunSchema(inputs: InputType = {}, jobInputs: any = {}) schema[key] = z.array(z.object({ value: valueSchema })); } - console.log("Generated schema:", schema); // Add this for debugging + // console.log("Generated schema:", schema); // Add this for debugging return z.object(schema).optional().default({}); } @@ -153,7 +153,7 @@ export function generateValues(inputs: InputType = {}, jobInputs: any = {}) { values[key] = [{ value }]; } - console.log("Generated values:", values); // Add this for debugging + // console.log("Generated values:", values); // Add this for debugging return values; } diff --git a/frontend/app/experiments/(experiment)/(results)/JobsAccordion.tsx b/frontend/app/experiments/(experiment)/(results)/JobsAccordion.tsx index c3070188..b0f33dc5 100644 --- a/frontend/app/experiments/(experiment)/(results)/JobsAccordion.tsx +++ b/frontend/app/experiments/(experiment)/(results)/JobsAccordion.tsx @@ -1,4 +1,4 @@ -import React, { useContext, useEffect } from "react"; +import React, { use, useContext, useEffect, useState } from "react"; import { TruncatedString } from "@/components/shared/TruncatedString"; import { Accordion, AccordionContent, AccordionItem, AccordionTrigger } from "@/components/ui/accordion"; @@ -8,6 +8,8 @@ import { ExperimentDetail } from "@/lib/redux"; import { ExperimentUIContext } from "../ExperimentUIContext"; import JobDetail from "./JobDetail"; +import { getAccessToken } from "@privy-io/react-auth"; +import backendUrl from "@/lib/backendUrl"; interface JobsAccordionProps { experiment: ExperimentDetail; @@ -15,17 +17,66 @@ interface JobsAccordionProps { export default function JobsAccordion({ experiment }: JobsAccordionProps) { const { activeJobUUID, setActiveJobUUID } = useContext(ExperimentUIContext); + const [jobs, setJobs] = useState(experiment.Jobs); // Initialize with initial job data + useEffect(() => { if (!activeJobUUID) { - setActiveJobUUID(experiment.Jobs?.[0]?.RayJobID); + setActiveJobUUID(jobs?.[0]?.RayJobID); } // eslint-disable-next-line react-hooks/exhaustive-deps + }, [jobs]); + + useEffect(() => { + setJobs(experiment.Jobs); }, [experiment.Jobs]); + // Polling to update job status and ID + useEffect(() => { + const fetchJobUpdates = async () => { + try { + const authToken = await getAccessToken(); + const nonFinalJobs = jobs.filter( + (job) => !["failed", "succeeded", "stopped"].includes(job.JobStatus) + ); + + if (nonFinalJobs.length > 0) { + const updatedJobs = await Promise.all( + nonFinalJobs.map(async (job) => { + const response = await fetch(`${backendUrl()}/jobs/${job.ID}`, { + headers: { + Authorization: `Bearer ${authToken}`, + }, + }); + const data = await response.json(); + return data; + }) + ); + + // Merge updated jobs back into the full jobs list + const mergedJobs = jobs.map((job) => + updatedJobs.find((updatedJob) => updatedJob.ID === job.ID) || job + ); + + setJobs(mergedJobs); + } + } catch (error) { + console.error("Error fetching job updates:", error); + } + }; + + const intervalId = setInterval(() => { + console.log("Polling for job updates"); + fetchJobUpdates(); + }, 10000); // Poll every 10 seconds + + return () => clearInterval(intervalId); + }, [jobs]); + + return ( - {[...experiment.Jobs]?.sort((a, b) => (a.ID || 0) - (b.ID || 0)).map((job, index) => { + {[...jobs]?.sort((a, b) => (a.ID || 0) - (b.ID || 0)).map((job, index) => { const validStates = ["queued", "processing", "pending", "running", "failed", "succeeded", "stopped"]; const status = (validStates.includes(job.JobStatus) ? job.JobStatus : "unknown") as "queued" | "processing" | "pending" | "running" | "failed" | "succeeded" | "stopped" | "unknown"; diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 68365f50..82e7301b 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -31,11 +31,14 @@ "@tailwindcss/container-queries": "^0.1.1", "@tanstack/react-table": "^8.10.7", "@web3auth/modal": "^7.0.3", + "3dmol": "^2.3.0", + "axios": "^1.7.3", "class-variance-authority": "^0.7.0", "clsx": "^2.0.0", "cmdk": "^0.2.0", "dayjs": "^1.11.10", "ethereumjs-util": "^7.1.5", + "jquery": "^3.7.1", "jwt-decode": "^3.1.2", "lucide-react": "^0.290.0", "molstar": "^4.0.1", @@ -5943,6 +5946,21 @@ "resolved": "https://registry.npmjs.org/@xobotyi/scrollbar-width/-/scrollbar-width-1.9.5.tgz", "integrity": "sha512-N8tkAACJx2ww8vFMneJmaAgmjAG1tnVBZJRLRcx061tmsLRZHSEZSLuGWnwPtunsSLvSqXQ2wfp7Mgqg1I+2dQ==" }, + "node_modules/3dmol": { + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/3dmol/-/3dmol-2.3.0.tgz", + "integrity": "sha512-aRRaAtgQiel7Vl+9C+MJ4pUKlMyLpbQW9wYxG0TXUSX4NSb0kDdGNkMbZuR/Wpqtlp46D7EWy5gj8TeIJ2sXQA==", + "dependencies": { + "iobuffer": "^5.3.1", + "netcdfjs": "^3.0.0", + "pako": "^2.1.0", + "upng-js": "^2.1.0" + }, + "engines": { + "node": ">=16.16.0", + "npm": ">=8.11" + } + }, "node_modules/abitype": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/abitype/-/abitype-1.0.0.tgz", @@ -6432,6 +6450,16 @@ "node": ">=4" } }, + "node_modules/axios": { + "version": "1.7.3", + "resolved": "https://registry.npmjs.org/axios/-/axios-1.7.3.tgz", + "integrity": "sha512-Ar7ND9pU99eJ9GpoGQKhKf58GpUOgnzuaB7ueNQ5BMi0p+LZ5oaEnfF999fAArcTIBwXTCHAmGcHOZJaWPq9Nw==", + "dependencies": { + "follow-redirects": "^1.15.6", + "form-data": "^4.0.0", + "proxy-from-env": "^1.1.0" + } + }, "node_modules/axobject-query": { "version": "3.2.1", "resolved": "https://registry.npmjs.org/axobject-query/-/axobject-query-3.2.1.tgz", @@ -9876,6 +9904,25 @@ "integrity": "sha512-36yxDn5H7OFZQla0/jFJmbIKTdZAQHngCedGxiMmpNfEZM0sdEeT+WczLQrjK6D7o2aiyLYDnkw0R3JK0Qv1RQ==", "dev": true }, + "node_modules/follow-redirects": { + "version": "1.15.6", + "resolved": "https://registry.npmjs.org/follow-redirects/-/follow-redirects-1.15.6.tgz", + "integrity": "sha512-wWN62YITEaOpSK584EZXJafH1AGpO8RVgElfkuXbTOrPX4fIfOyEpW/CsiNd8JdYrAoOvafRTOEnvsO++qCqFA==", + "funding": [ + { + "type": "individual", + "url": "https://github.com/sponsors/RubenVerborgh" + } + ], + "engines": { + "node": ">=4.0" + }, + "peerDependenciesMeta": { + "debug": { + "optional": true + } + } + }, "node_modules/for-each": { "version": "0.3.3", "resolved": "https://registry.npmjs.org/for-each/-/for-each-0.3.3.tgz", @@ -10685,6 +10732,11 @@ "fp-ts": "^2.5.0" } }, + "node_modules/iobuffer": { + "version": "5.3.2", + "resolved": "https://registry.npmjs.org/iobuffer/-/iobuffer-5.3.2.tgz", + "integrity": "sha512-kO3CjNfLZ9t+tHxAMd+Xk4v3D/31E91rMs1dHrm7ikEQrlZ8mLDbQ4z3tZfDM48zOkReas2jx8MWSAmN9+c8Fw==" + }, "node_modules/ipaddr.js": { "version": "1.9.1", "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-1.9.1.tgz", @@ -11333,6 +11385,11 @@ "node": ">=10" } }, + "node_modules/jquery": { + "version": "3.7.1", + "resolved": "https://registry.npmjs.org/jquery/-/jquery-3.7.1.tgz", + "integrity": "sha512-m4avr8yL8kmFN8psrbFFFmB/If14iN5o9nw/NgnnM+kybDJpRsAynV2BsfpTYrTRysYUdADVD7CkUUizgkpLfg==" + }, "node_modules/js-cookie": { "version": "3.0.5", "resolved": "https://registry.npmjs.org/js-cookie/-/js-cookie-3.0.5.tgz", @@ -13261,6 +13318,14 @@ "node": ">= 0.6" } }, + "node_modules/netcdfjs": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/netcdfjs/-/netcdfjs-3.0.0.tgz", + "integrity": "sha512-LOvT8KkC308qtpUkcBPiCMBtii7ZQCN6LxcVheWgyUeZ6DQWcpSRFV9dcVXLj/2eHZ/bre9tV5HTH4Sf93vrFw==", + "dependencies": { + "iobuffer": "^5.3.2" + } + }, "node_modules/next": { "version": "14.1.0", "resolved": "https://registry.npmjs.org/next/-/next-14.1.0.tgz", @@ -13733,6 +13798,11 @@ "node": ">=6" } }, + "node_modules/pako": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/pako/-/pako-2.1.0.tgz", + "integrity": "sha512-w+eufiZ1WuJYgPXbV/PO3NCMEc3xqylkKHzp8bxp1uW4qaSNQUkwmLLEc3kKsfz8lpV1F8Ht3U1Cm+9Srog2ug==" + }, "node_modules/parent-module": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/parent-module/-/parent-module-1.0.1.tgz", @@ -14241,6 +14311,11 @@ "resolved": "https://registry.npmjs.org/proxy-compare/-/proxy-compare-2.5.1.tgz", "integrity": "sha512-oyfc0Tx87Cpwva5ZXezSp5V9vht1c7dZBhvuV/y3ctkgMVUmiAGDVeeB0dKhGSyT0v1ZTEQYpe/RXlBVBNuCLA==" }, + "node_modules/proxy-from-env": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/proxy-from-env/-/proxy-from-env-1.1.0.tgz", + "integrity": "sha512-D+zkORCbA9f1tdWRK0RaCR3GPv50cMxcrz4X8k5LTSUD1Dkw47mKJEZQNunItRTkWwgtaUSo1RVFRIG9ZXiFYg==" + }, "node_modules/pump": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/pump/-/pump-3.0.0.tgz", @@ -16810,6 +16885,19 @@ "browserslist": ">= 4.21.0" } }, + "node_modules/upng-js": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/upng-js/-/upng-js-2.1.0.tgz", + "integrity": "sha512-d3xzZzpMP64YkjP5pr8gNyvBt7dLk/uGI67EctzDuVp4lCZyVMo0aJO6l/VDlgbInJYDY6cnClLoBp29eKWI6g==", + "dependencies": { + "pako": "^1.0.5" + } + }, + "node_modules/upng-js/node_modules/pako": { + "version": "1.0.11", + "resolved": "https://registry.npmjs.org/pako/-/pako-1.0.11.tgz", + "integrity": "sha512-4hLB8Py4zZce5s4yd9XzopqwVv/yGNhV1Bl8NTmCq1763HeK2+EwVTv+leGeL13Dnh2wfbqowVPXCIO0z4taYw==" + }, "node_modules/uqr": { "version": "0.1.2", "resolved": "https://registry.npmjs.org/uqr/-/uqr-0.1.2.tgz", @@ -21575,6 +21663,17 @@ "resolved": "https://registry.npmjs.org/@xobotyi/scrollbar-width/-/scrollbar-width-1.9.5.tgz", "integrity": "sha512-N8tkAACJx2ww8vFMneJmaAgmjAG1tnVBZJRLRcx061tmsLRZHSEZSLuGWnwPtunsSLvSqXQ2wfp7Mgqg1I+2dQ==" }, + "3dmol": { + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/3dmol/-/3dmol-2.3.0.tgz", + "integrity": "sha512-aRRaAtgQiel7Vl+9C+MJ4pUKlMyLpbQW9wYxG0TXUSX4NSb0kDdGNkMbZuR/Wpqtlp46D7EWy5gj8TeIJ2sXQA==", + "requires": { + "iobuffer": "^5.3.1", + "netcdfjs": "^3.0.0", + "pako": "^2.1.0", + "upng-js": "^2.1.0" + } + }, "abitype": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/abitype/-/abitype-1.0.0.tgz", @@ -21946,6 +22045,16 @@ "integrity": "sha512-/dlp0fxyM3R8YW7MFzaHWXrf4zzbr0vaYb23VBFCl83R7nWNPg/yaQw2Dc8jzCMmDVLhSdzH8MjrsuIUuvX+6g==", "dev": true }, + "axios": { + "version": "1.7.3", + "resolved": "https://registry.npmjs.org/axios/-/axios-1.7.3.tgz", + "integrity": "sha512-Ar7ND9pU99eJ9GpoGQKhKf58GpUOgnzuaB7ueNQ5BMi0p+LZ5oaEnfF999fAArcTIBwXTCHAmGcHOZJaWPq9Nw==", + "requires": { + "follow-redirects": "^1.15.6", + "form-data": "^4.0.0", + "proxy-from-env": "^1.1.0" + } + }, "axobject-query": { "version": "3.2.1", "resolved": "https://registry.npmjs.org/axobject-query/-/axobject-query-3.2.1.tgz", @@ -24670,6 +24779,11 @@ "integrity": "sha512-36yxDn5H7OFZQla0/jFJmbIKTdZAQHngCedGxiMmpNfEZM0sdEeT+WczLQrjK6D7o2aiyLYDnkw0R3JK0Qv1RQ==", "dev": true }, + "follow-redirects": { + "version": "1.15.6", + "resolved": "https://registry.npmjs.org/follow-redirects/-/follow-redirects-1.15.6.tgz", + "integrity": "sha512-wWN62YITEaOpSK584EZXJafH1AGpO8RVgElfkuXbTOrPX4fIfOyEpW/CsiNd8JdYrAoOvafRTOEnvsO++qCqFA==" + }, "for-each": { "version": "0.3.3", "resolved": "https://registry.npmjs.org/for-each/-/for-each-0.3.3.tgz", @@ -25274,6 +25388,11 @@ "integrity": "sha512-zz2Z69v9ZIC3mMLYWIeoUcwWD6f+O7yP92FMVVaXEOSZH1jnVBmET/urd/uoarD1WGBY4rCj8TAyMPzsGNzMFQ==", "requires": {} }, + "iobuffer": { + "version": "5.3.2", + "resolved": "https://registry.npmjs.org/iobuffer/-/iobuffer-5.3.2.tgz", + "integrity": "sha512-kO3CjNfLZ9t+tHxAMd+Xk4v3D/31E91rMs1dHrm7ikEQrlZ8mLDbQ4z3tZfDM48zOkReas2jx8MWSAmN9+c8Fw==" + }, "ipaddr.js": { "version": "1.9.1", "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-1.9.1.tgz", @@ -25691,6 +25810,11 @@ "resolved": "https://registry.npmjs.org/joycon/-/joycon-3.1.1.tgz", "integrity": "sha512-34wB/Y7MW7bzjKRjUKTa46I2Z7eV62Rkhva+KkopW7Qvv/OSWBqvkSY7vusOPrNuZcUG3tApvdVgNB8POj3SPw==" }, + "jquery": { + "version": "3.7.1", + "resolved": "https://registry.npmjs.org/jquery/-/jquery-3.7.1.tgz", + "integrity": "sha512-m4avr8yL8kmFN8psrbFFFmB/If14iN5o9nw/NgnnM+kybDJpRsAynV2BsfpTYrTRysYUdADVD7CkUUizgkpLfg==" + }, "js-cookie": { "version": "3.0.5", "resolved": "https://registry.npmjs.org/js-cookie/-/js-cookie-3.0.5.tgz", @@ -27071,6 +27195,14 @@ "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-0.6.3.tgz", "integrity": "sha512-+EUsqGPLsM+j/zdChZjsnX51g4XrHFOIXwfnCVPGlQk/k5giakcKsuxCObBRu6DSm9opw/O6slWbJdghQM4bBg==" }, + "netcdfjs": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/netcdfjs/-/netcdfjs-3.0.0.tgz", + "integrity": "sha512-LOvT8KkC308qtpUkcBPiCMBtii7ZQCN6LxcVheWgyUeZ6DQWcpSRFV9dcVXLj/2eHZ/bre9tV5HTH4Sf93vrFw==", + "requires": { + "iobuffer": "^5.3.2" + } + }, "next": { "version": "14.1.0", "resolved": "https://registry.npmjs.org/next/-/next-14.1.0.tgz", @@ -27402,6 +27534,11 @@ "resolved": "https://registry.npmjs.org/p-try/-/p-try-2.2.0.tgz", "integrity": "sha512-R4nPAVTAU0B9D35/Gk3uJf/7XYbQcyohSKdvAxIRSNghFl4e71hVoGnBNQz9cWaXxO2I10KTC+3jMdvvoKw6dQ==" }, + "pako": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/pako/-/pako-2.1.0.tgz", + "integrity": "sha512-w+eufiZ1WuJYgPXbV/PO3NCMEc3xqylkKHzp8bxp1uW4qaSNQUkwmLLEc3kKsfz8lpV1F8Ht3U1Cm+9Srog2ug==" + }, "parent-module": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/parent-module/-/parent-module-1.0.1.tgz", @@ -27768,6 +27905,11 @@ "resolved": "https://registry.npmjs.org/proxy-compare/-/proxy-compare-2.5.1.tgz", "integrity": "sha512-oyfc0Tx87Cpwva5ZXezSp5V9vht1c7dZBhvuV/y3ctkgMVUmiAGDVeeB0dKhGSyT0v1ZTEQYpe/RXlBVBNuCLA==" }, + "proxy-from-env": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/proxy-from-env/-/proxy-from-env-1.1.0.tgz", + "integrity": "sha512-D+zkORCbA9f1tdWRK0RaCR3GPv50cMxcrz4X8k5LTSUD1Dkw47mKJEZQNunItRTkWwgtaUSo1RVFRIG9ZXiFYg==" + }, "pump": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/pump/-/pump-3.0.0.tgz", @@ -29602,6 +29744,21 @@ "picocolors": "^1.0.0" } }, + "upng-js": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/upng-js/-/upng-js-2.1.0.tgz", + "integrity": "sha512-d3xzZzpMP64YkjP5pr8gNyvBt7dLk/uGI67EctzDuVp4lCZyVMo0aJO6l/VDlgbInJYDY6cnClLoBp29eKWI6g==", + "requires": { + "pako": "^1.0.5" + }, + "dependencies": { + "pako": { + "version": "1.0.11", + "resolved": "https://registry.npmjs.org/pako/-/pako-1.0.11.tgz", + "integrity": "sha512-4hLB8Py4zZce5s4yd9XzopqwVv/yGNhV1Bl8NTmCq1763HeK2+EwVTv+leGeL13Dnh2wfbqowVPXCIO0z4taYw==" + } + } + }, "uqr": { "version": "0.1.2", "resolved": "https://registry.npmjs.org/uqr/-/uqr-0.1.2.tgz", diff --git a/frontend/package.json b/frontend/package.json index f46d6702..3484d876 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -33,11 +33,14 @@ "@tailwindcss/container-queries": "^0.1.1", "@tanstack/react-table": "^8.10.7", "@web3auth/modal": "^7.0.3", + "3dmol": "^2.3.0", + "axios": "^1.7.3", "class-variance-authority": "^0.7.0", "clsx": "^2.0.0", "cmdk": "^0.2.0", "dayjs": "^1.11.10", "ethereumjs-util": "^7.1.5", + "jquery": "^3.7.1", "jwt-decode": "^3.1.2", "lucide-react": "^0.290.0", "molstar": "^4.0.1", diff --git a/frontend/tsconfig.json b/frontend/tsconfig.json index 31f3bc8b..67223628 100644 --- a/frontend/tsconfig.json +++ b/frontend/tsconfig.json @@ -19,6 +19,6 @@ "incremental": true, "plugins": [{ "name": "next" }] }, - "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"], + "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts", "3dmol.d.ts"], "exclude": ["node_modules"] } diff --git a/gateway/handlers/api.go b/gateway/handlers/api.go index c732deda..4abc99c0 100644 --- a/gateway/handlers/api.go +++ b/gateway/handlers/api.go @@ -92,7 +92,7 @@ func ListAPIKeysHandler(db *gorm.DB) http.HandlerFunc { } var apiKeys []models.APIKey - if err := db.Where("wallet_address = ?", user.WalletAddress).Find(&apiKeys).Error; err != nil { + if err := db.Where("user_id = ?", user.WalletAddress).Find(&apiKeys).Error; err != nil { utils.SendJSONError(w, "Failed to get API keys: "+err.Error(), http.StatusInternalServerError) return } diff --git a/gateway/handlers/experiments.go b/gateway/handlers/experiments.go index 1ef3c5df..ff84b1e6 100644 --- a/gateway/handlers/experiments.go +++ b/gateway/handlers/experiments.go @@ -105,6 +105,26 @@ func AddExperimentHandler(db *gorm.DB) http.HandlerFunc { } log.Println("Initialized IO List") + totalComputeCost := len(ioList) * model.ComputeCost + + thresholdStr := os.Getenv("TIER_THRESHOLD") + if thresholdStr == "" { + utils.SendJSONError(w, "TIER_THRESHOLD environment variable is not set", http.StatusInternalServerError) + return + } + + threshold, err := strconv.Atoi(thresholdStr) + if err != nil { + utils.SendJSONError(w, fmt.Sprintf("Error converting TIER_THRESHOLD to integer: %v", err), http.StatusInternalServerError) + return + } + + if user.ComputeTally+totalComputeCost > threshold && user.Tier == 0 { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]string{"redirectUrl": "/subscribe"}) + return + } + experiment := models.Experiment{ WalletAddress: user.WalletAddress, Name: name, @@ -212,18 +232,6 @@ func AddExperimentHandler(db *gorm.DB) http.HandlerFunc { return } - thresholdStr := os.Getenv("TIER_THRESHOLD") - if thresholdStr == "" { - utils.SendJSONError(w, "TIER_THRESHOLD environment variable is not set", http.StatusInternalServerError) - return - } - - threshold, err := strconv.Atoi(thresholdStr) - if err != nil { - utils.SendJSONError(w, fmt.Sprintf("Error converting TIER_THRESHOLD to integer: %v", err), http.StatusInternalServerError) - return - } - err = UpdateUserTier(db, user.WalletAddress, threshold) if err != nil { utils.SendJSONError(w, fmt.Sprintf("Error updating user tier: %v", err), http.StatusInternalServerError) @@ -562,6 +570,26 @@ func AddJobToExperimentHandler(db *gorm.DB) http.HandlerFunc { } log.Println("Initialized IO List") + totalComputeCost := len(ioList) * model.ComputeCost + + thresholdStr := os.Getenv("TIER_THRESHOLD") + if thresholdStr == "" { + utils.SendJSONError(w, "TIER_THRESHOLD environment variable is not set", http.StatusInternalServerError) + return + } + + threshold, err := strconv.Atoi(thresholdStr) + if err != nil { + utils.SendJSONError(w, fmt.Sprintf("Error converting TIER_THRESHOLD to integer: %v", err), http.StatusInternalServerError) + return + } + + if user.ComputeTally+totalComputeCost > threshold && user.Tier == 0 { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]string{"redirectUrl": "/subscribe"}) + return + } + for _, ioItem := range ioList { log.Println("Creating job entry") inputsJSON, err := json.Marshal(ioItem.Inputs) @@ -655,18 +683,6 @@ func AddJobToExperimentHandler(db *gorm.DB) http.HandlerFunc { return } - thresholdStr := os.Getenv("TIER_THRESHOLD") - if thresholdStr == "" { - http.Error(w, "TIER_THRESHOLD environment variable is not set", http.StatusInternalServerError) - return - } - - threshold, err := strconv.Atoi(thresholdStr) - if err != nil { - http.Error(w, fmt.Sprintf("Error converting TIER_THRESHOLD to integer: %v", err), http.StatusInternalServerError) - return - } - err = UpdateUserTier(db, user.WalletAddress, threshold) if err != nil { http.Error(w, fmt.Sprintf("Error updating user tier: %v", err), http.StatusInternalServerError) diff --git a/gateway/utils/queue.go b/gateway/utils/queue.go index d60dc6d8..26e2ca69 100644 --- a/gateway/utils/queue.go +++ b/gateway/utils/queue.go @@ -87,7 +87,7 @@ func (rq *RayQueue) worker(workerID int) { fmt.Printf("Error marking dangling jobs to stopped: %v\n", err) } var job models.Job - err = fetchAndMarkOldestQueuedJobAsProcessing(&job, models.QueueTypeRay, rq.db) + err = fetchAndMarkOldestQueuedJobAsProcessing(&job, rq.db) if err != nil { state.Busy = false state.CurrentJob = nil @@ -140,9 +140,6 @@ func fetchAndMarkOldestRunningJobAsStopped(db *gorm.DB) error { func checkRunningJob(jobID uint, db *gorm.DB) error { var job models.Job err := fetchJobWithModelAndExperimentData(&job, jobID, db) - if err != nil { - return err - } if err != nil && strings.Contains(err.Error(), "Job not found") { fmt.Printf("Job %v , %v has missing Ray Job, failing Job\n", job.ID, job.RayJobID) return setJobStatus(&job, models.JobStateFailed, fmt.Sprintf("Ray job %v not found", job.RayJobID), db) @@ -154,6 +151,14 @@ func checkRunningJob(jobID uint, db *gorm.DB) error { fmt.Printf("Job %v , %v is still in pending state nothing to do\n", job.ID, job.RayJobID) return nil } else if ray.JobIsRunning(job.RayJobID) { + if job.JobStatus != models.JobStateRunning { + fmt.Printf("Job %v , %v is running, updating status\n", job.ID, job.RayJobID) + err := setJobStatus(&job, models.JobStateRunning, "", db) + if err != nil { + return err + } + createInferenceEvent(job.ID, models.JobStateRunning, job.RayJobID, job.RetryCount, db) + } fmt.Printf("Job %v , %v is still running nothing to do\n", job.ID, job.RayJobID) return nil } else if ray.JobFailed(job.RayJobID) { @@ -245,10 +250,10 @@ func MonitorRunningJobs(db *gorm.DB) error { } func fetchRunningJobsWithModelData(jobs *[]models.Job, db *gorm.DB) error { - return db.Preload("Model").Where("job_status = ?", models.JobStateRunning).Where("job_type", models.JobTypeJob).Find(jobs).Error + return db.Preload("Model").Where("job_status in (?,?)", models.JobStateRunning, models.JobStatePending).Where("job_type", models.JobTypeJob).Find(jobs).Error } -func fetchAndMarkOldestQueuedJobAsProcessing(job *models.Job, queueType models.QueueType, db *gorm.DB) error { +func fetchAndMarkOldestQueuedJobAsProcessing(job *models.Job, db *gorm.DB) error { return db.Transaction(func(tx *gorm.DB) error { if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("job_status = ?", models.JobStateQueued).Order("created_at ASC").First(job).Error; err != nil { return err @@ -304,14 +309,17 @@ func processRayJob(jobID uint, db *gorm.DB) error { rayJobID := uuid.New().String() log.Printf("Here is the UUID for the job: %v\n", rayJobID) job.RayJobID = rayJobID - job.JobStatus = models.JobStatePending if err := db.Save(&job).Error; err != nil { return err } - createInferenceEvent(job.ID, models.JobStatePending, job.RayJobID, 0, db) if err := submitRayJobAndUpdateID(&job, db); err != nil { return err } + job.JobStatus = models.JobStatePending + if err := db.Save(&job).Error; err != nil { + return err + } + createInferenceEvent(job.ID, models.JobStatePending, job.RayJobID, 0, db) } return nil @@ -417,9 +425,8 @@ func submitRayJobAndUpdateID(job *models.Job, db *gorm.DB) error { modelPath := job.Model.S3URI log.Printf("Submitting to Ray with inputs: %+v\n", inputs) - createInferenceEvent(job.ID, models.JobStateRunning, job.RayJobID, 0, db) - setJobStatusAndID(job, models.JobStateRunning, job.RayJobID, "", db) - log.Printf("setting job %v to running\n", job.ID) + setRayJobID(job, job.RayJobID, db) + // log.Printf("setting job %v to running\n", job.ID) resp, err := ray.SubmitRayJob(*job, modelPath, job.RayJobID, inputs, db) if err != nil { return err @@ -500,7 +507,11 @@ func submitRayJobAndUpdateID(job *models.Job, db *gorm.DB) error { } fmt.Printf("Job had id %v\n", job.ID) - fmt.Printf("Finished Job with Ray id %v and status %v\n", job.RayJobID, job.JobStatus) + if job.JobType == models.JobTypeJob { + fmt.Printf("Submitted Job with Ray id %v and status %v\n", job.RayJobID, job.JobStatus) + } else { + fmt.Printf("Finished Job with Ray id %v and status %v\n", job.RayJobID, job.JobStatus) + } err = db.Save(&job).Error if err != nil { return err @@ -508,15 +519,14 @@ func submitRayJobAndUpdateID(job *models.Job, db *gorm.DB) error { return nil } -func setJobStatusAndID(job *models.Job, state models.JobState, rayJobID string, errorMessage string, db *gorm.DB) error { - job.JobStatus = state - job.StartedAt = time.Now().UTC() - job.Error = errorMessage +func setRayJobID(job *models.Job, rayJobID string, db *gorm.DB) error { job.RayJobID = rayJobID + job.StartedAt = time.Now().UTC() err := db.Save(&job).Error if err != nil { return err } + createInferenceEvent(job.ID, models.JobStatePending, job.RayJobID, 0, db) return nil } @@ -587,19 +597,8 @@ func completeRayJobAndAddFiles(job *models.Job, body []byte, resultJSON models.R } } - fmt.Printf("Looping through files in RayJobResponse\n %v\n", resultJSON.Files) - // Iterate over all files in the RayJobResponse - for key, fileDetail := range resultJSON.Files { - fmt.Printf("AddFileToDB for file: %s, Key: %s\n", fileDetail.URI, key) - if err := addFileToDB(job, fileDetail, key, db); err != nil { - return fmt.Errorf("failed to add file (%s) to database: %v", key, err) - } - } - - fmt.Printf("Adding PDB file to DB\n %v\n", resultJSON.PDB) - // Special handling for PDB as it's a common file across many jobs - if err := addFileToDB(job, resultJSON.PDB, "pdb", db); err != nil { - return fmt.Errorf("failed to add PDB file to database: %v", err) + if err := addFiles(job, resultJSON, db); err != nil { + return fmt.Errorf("failed to add files: %v", err) } return nil @@ -612,7 +611,11 @@ func addFileToDB(job *models.Job, fileDetail models.FileDetail, fileType string, var file models.File result := db.Where("s3_uri = ?", fileDetail.URI).First(&file) if result.Error == nil { - fmt.Println("File already exists in DB:", fileDetail.URI) + fmt.Println("File already exists in DB, adding an entry in the job_output_files table:", fileDetail.URI) + job.OutputFiles = append(job.OutputFiles, file) + if err := db.Save(&job).Error; err != nil { + return fmt.Errorf("error updating job with new output file: %v", err) + } return nil // File already processed } @@ -651,7 +654,7 @@ func addFileToDB(job *models.Job, fileDetail models.FileDetail, fileType string, func processNewFiles(job *models.Job, db *gorm.DB) error { bucketName := os.Getenv("BUCKET_NAME") - prefix := fmt.Sprintf("%s-", job.RayJobID) // Adjusted to the new naming pattern + prefix := job.RayJobID // Adjusted to the new naming pattern s3client, err := s3client.NewS3Client() if err != nil { @@ -688,7 +691,7 @@ func processFile(fileName string, job *models.Job, db *gorm.DB) error { } // Add files and update related job data in the database without marking the job as completed - if err := addFilesAndUpdateJob(job, data, rayJobResponse, db); err != nil { + if err := addFiles(job, rayJobResponse, db); err != nil { log.Printf("Failed to add files and update job for job %d from file %s: %v", job.ID, fileName, err) return err } @@ -700,7 +703,7 @@ func processFile(fileName string, job *models.Job, db *gorm.DB) error { EventType: models.EventTypeFileProcessed, JobStatus: models.JobStateRunning, FileName: fileName, - EventTime: time.Now(), + EventTime: time.Now().UTC(), OutputJson: datatypes.JSON(data), // Storing the JSON output directly in the event } if err := db.Create(&event).Error; err != nil { @@ -730,15 +733,23 @@ func fileProcessed(fileName string, jobID uint, db *gorm.DB) bool { return count > 0 } -func addFilesAndUpdateJob(job *models.Job, data []byte, response models.RayJobResponse, db *gorm.DB) error { - fmt.Printf("Adding output files and updating job data for job %d\n", job.ID) +func addFiles(job *models.Job, response models.RayJobResponse, db *gorm.DB) error { + fmt.Printf("Adding output files for job %d\n", job.ID) - // Loop through the files detailed in the response + fmt.Printf("Looping through files in RayJobResponse\n %v\n", response.Files) + // Iterate over all files in the RayJobResponse for key, fileDetail := range response.Files { + fmt.Printf("AddFileToDB for file: %s, Key: %s\n", fileDetail.URI, key) if err := addFileToDB(job, fileDetail, key, db); err != nil { return fmt.Errorf("failed to add file (%s) to database: %v", key, err) } } + fmt.Printf("Adding PDB file to DB\n %v\n", response.PDB) + // Special handling for PDB as it's a common file across many jobs + if err := addFileToDB(job, response.PDB, "pdb", db); err != nil { + return fmt.Errorf("failed to add PDB file to database: %v", err) + } + return nil } diff --git a/internal/ray/ray.go b/internal/ray/ray.go index 3c7682b5..19673b27 100644 --- a/internal/ray/ray.go +++ b/internal/ray/ray.go @@ -35,7 +35,7 @@ func GetRayJobApiHost() string { if exists { return rayApiHost } else { - return "http://localhost:8265" // Default Ray API host + return "http://base-job:8265" // Default Ray API host } } @@ -123,7 +123,6 @@ func CreateRayJob(job *models.Job, modelPath string, rayJobID string, inputs map if err != nil { return nil, err } - log.Printf("Submitting Ray job with testeks inputs: %s\n", inputsJSON) rayServiceURL = GetRayJobApiHost() + model.RayEndpoint runtimeEnv := map[string]interface{}{ diff --git a/models/ray/colabfold_demo_ray_job.json b/models/ray/demo_ray_job.json similarity index 100% rename from models/ray/colabfold_demo_ray_job.json rename to models/ray/demo_ray_job.json diff --git a/models/ray/demo_ray_job_3dmol.json b/models/ray/demo_ray_job_3dmol.json new file mode 100644 index 00000000..dc02cfd7 --- /dev/null +++ b/models/ray/demo_ray_job_3dmol.json @@ -0,0 +1,96 @@ +{ + "name": "demo_ray_job_3dmol", + "description": "This model generates protein binders based on an input .pdb file and optional additional prompting. \nBinders are generated by combining a conditional protein structure diffusion model (RFDiffusion) with an inverse folding model (ProteinMPNN). After generation, candidates are evaluated using a protein folding model (AF Multimer via Colabfold).\nThe model's output is a set of candidate binders with .pdb files for each candidate and corresponding scoring metrics.", + "guide": "Note: steps marked (optional) and (advanced) give additional constraints for the designs but it is not necessary to add/edit their input in order to submit an experimental run.\n\n1. Select a .pdb file containing the *target protein* you would like to design a binder against.\n2. Specify a *target chain* from the pdb file.\n3. (optional) Define *hotspot* residues on the target chain as the binding sites, stating the target chain ID followed by the residue index (e.g. A30).\n4. (optional) Specify the *binder length* (total number of residues) you wish to design.\n5. (optional) Define an *expert prompt* to crop the target and refine the binder design specs. The prompt string has to respect the following syntax. First, information on the target is specified. This is done by specifying the target chain ID followed by the indices of the first and last residue of a cropped region on the target chain (e.g. A20:110). For the binder the following information can be provided. First, residues of the binder chain to hold fixed (aka scaffold) during design. These are specified by the binder chain ID followed by an index range (e.g. B6:9, or B6:6 for a single fixed residue). Second, segments of residues for free design. Use small *x* followed by the number of residues to be designed (e.g. x15). Different parts of the prompt are separated by */*. Providing a prompt, overwrites other inputs.\n6. (optional) Crop the target structure by specifying the *start residue* and *end residue* on the target chain.", + "author": "sokrypton", + "github": "https://github.com/sokrypton/ColabDesign", + "paper": "https://www.nature.com/articles/s41586-023-06415-8", + "task": "protein design", + "taskCategory": "protein-binder-design", + "modelType": "ray", + "rayEndpoint": "/api/jobs/", + "rayJobEntryPoint": "python demo.py", + "checkpointCompatible": true, + "xAxis": "plddt", + "yAxis": "i_pae", + "jobType": "job", + "metricsDescription": "*Stability Score:* larger value indicates higher confidence in the predicted local structure\n*Affinity Score:* larger value indicates higher confidence in the predicted interface structure\n*Note:* designs which lie within the green square are recommended for laboratory testing", + "computeCost": 50, + "inputs": { + "target_chain": { + "type": "string", + "description": "Chain ID of the relevant target protein in the PDB file. Assuming A as target chain ID unless provided otherwise.", + "array": false, + "glob": [ + "" + ], + "default": "A", + "min": "", + "max": "", + "example": "A", + "grouping": "target", + "position": "102", + "required": true + }, + "binder_length": { + "type": "number", + "description": "The length of the generated protein binder candidates.", + "array": false, + "glob": [ + "" + ], + "default": "", + "min": "0", + "max": "1000", + "example": "80", + "grouping": "target", + "position": "104", + "required": false + }, + "target_start_residue": { + "type": "number", + "description": "Use this setting to crop the target chain for faster binder generation. Integer index defining the first residue of a cropped section on the target chain.", + "array": false, + "glob": [ + "" + ], + "default": "", + "min": "0", + "max": "", + "example": "20", + "grouping": "_advanced", + "position": "201", + "required": false + }, + "target_end_residue": { + "type": "number", + "description": "Use this setting to crop the target chain for faster binder generation. Integer index defining the last residue of a cropped section on the target chain.", + "array": false, + "glob": [ + "" + ], + "default": "", + "min": "0", + "max": "", + "example": "120", + "grouping": "_advanced", + "position": "202", + "required": false + }, + "expert_prompt_for_refined_binder_design": { + "type": "string", + "default": "", + "description": "Use this setting for advanced prompts to crop and scaffold protein chains. For motif scaffolding, specify the chain ID of the motif followed by a range of residue indices. For free design, write x followed by the number of residues to generate. The prompt overrides all other GUI inputs.", + "example": "A10:110/B4:5/x6/B9:50/x51", + "grouping": "_advanced", + "position": "203", + "required": false + } + }, + "outputs": { + "string_message": { + "type": "File", + "glob": ["*.json"] + } + } +}