Feature/nim container (#3966)

* add nim container setup

* check if image or container exist before pulling

* update NIM dialog

* update chat nvidia api key

* update nim container version

* update nim container version
This commit is contained in:
Henry Heng
2025-02-27 09:14:01 +00:00
committed by GitHub
parent 7d8541a44b
commit 1678815540
10 changed files with 606 additions and 1056 deletions
@@ -0,0 +1,328 @@
import { useState, useEffect } from 'react'
import { createPortal } from 'react-dom'
import axios from 'axios'
import PropTypes from 'prop-types'
import {
Dialog,
DialogTitle,
DialogContent,
DialogActions,
Button,
CircularProgress,
Stepper,
Step,
StepLabel,
Select,
MenuItem,
FormControl,
InputLabel
} from '@mui/material'
const NvidiaNIMDialog = ({ open, onClose, onComplete }) => {
const portalElement = document.getElementById('portal')
const modelOptions = {
'nv-mistralai/mistral-nemo-12b-instruct:latest': {
label: 'Mistral Nemo 12B Instruct',
licenseUrl: 'https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nv-mistralai/containers/mistral-nemo-12b-instruct'
},
'meta/llama-3.1-8b-instruct-rtx:latest': {
label: 'Llama 3.1 8B Instruct',
licenseUrl: 'https://catalog.ngc.nvidia.com/orgs/nim/teams/meta/containers/llama-3.1-8b-instruct'
}
}
const [activeStep, setActiveStep] = useState(0)
const [loading, setLoading] = useState(false)
const [imageTag, setImageTag] = useState('')
const [pollInterval, setPollInterval] = useState(null)
const steps = ['Download Installer', 'Pull Image', 'Start Container']
const handleDownloadInstaller = async () => {
try {
setLoading(true)
await axios.get('/api/v1/nvidia-nim/download-installer')
setLoading(false)
} catch (err) {
let errorData = err.message
if (typeof err === 'string') {
errorData = err
} else if (err.response?.data) {
errorData = err.response.data.message
}
alert('Failed to download installer: ' + errorData)
setLoading(false)
}
}
const preload = async () => {
try {
setLoading(true)
await axios.get('/api/v1/nvidia-nim/preload')
setLoading(false)
setActiveStep(1)
} catch (err) {
let errorData = err.message
if (typeof err === 'string') {
errorData = err
} else if (err.response?.data) {
errorData = err.response.data.message
}
alert('Failed to preload: ' + errorData)
setLoading(false)
}
}
const handlePullImage = async () => {
try {
setLoading(true)
try {
const imageResponse = await axios.post('/api/v1/nvidia-nim/get-image', { imageTag })
if (imageResponse.data && imageResponse.data.tag === imageTag) {
setLoading(false)
setActiveStep(2)
return
}
} catch (err) {
// Continue if image not found
if (err.response?.status !== 404) {
throw err
}
}
// Get token first
const tokenResponse = await axios.get('/api/v1/nvidia-nim/get-token')
const apiKey = tokenResponse.data.access_token
// Pull image
await axios.post('/api/v1/nvidia-nim/pull-image', {
imageTag,
apiKey
})
// Start polling for image status
const interval = setInterval(async () => {
try {
const imageResponse = await axios.post('/api/v1/nvidia-nim/get-image', { imageTag })
if (imageResponse.data) {
clearInterval(interval)
setLoading(false)
setActiveStep(2)
}
} catch (err) {
// Continue polling if image not found
if (err.response?.status !== 404) {
clearInterval(interval)
alert('Failed to check image status: ' + err.message)
setLoading(false)
}
}
}, 5000)
setPollInterval(interval)
} catch (err) {
let errorData = err.message
if (typeof err === 'string') {
errorData = err
} else if (err.response?.data) {
errorData = err.response.data.message
}
alert('Failed to pull image: ' + errorData)
setLoading(false)
}
}
const handleStartContainer = async () => {
try {
setLoading(true)
try {
const containerResponse = await axios.post('/api/v1/nvidia-nim/get-container', { imageTag })
if (containerResponse.data && containerResponse.data && containerResponse.data.status === 'running') {
// wait additional 10 seconds for container to be ready
await new Promise((resolve) => setTimeout(resolve, 10000))
setLoading(false)
onComplete(containerResponse.data)
onClose()
return
}
} catch (err) {
// Continue if container not found
if (err.response?.status !== 404) {
throw err
}
}
const tokenResponse = await axios.get('/api/v1/nvidia-nim/get-token')
const apiKey = tokenResponse.data.access_token
await axios.post('/api/v1/nvidia-nim/start-container', {
imageTag,
apiKey
})
// Start polling for container status
const interval = setInterval(async () => {
try {
const containerResponse = await axios.post('/api/v1/nvidia-nim/get-container', { imageTag })
if (containerResponse.data) {
clearInterval(interval)
setLoading(false)
onComplete(containerResponse.data)
onClose()
}
} catch (err) {
// Continue polling if container not found
if (err.response?.status !== 404) {
clearInterval(interval)
alert('Failed to check container status: ' + err.message)
setLoading(false)
}
}
}, 5000)
setPollInterval(interval)
} catch (err) {
let errorData = err.message
if (typeof err === 'string') {
errorData = err
} else if (err.response?.data) {
errorData = err.response.data.message
}
alert('Failed to start container: ' + errorData)
setLoading(false)
}
}
const handleNext = () => {
if (activeStep === 1 && !imageTag) {
alert('Please enter an image tag')
return
}
switch (activeStep) {
case 0:
preload()
break
case 1:
handlePullImage()
break
case 2:
handleStartContainer()
break
default:
setActiveStep((prev) => prev + 1)
}
}
// Cleanup polling on unmount
useEffect(() => {
return () => {
if (pollInterval) {
clearInterval(pollInterval)
}
}
}, [pollInterval])
// clear state on close
useEffect(() => {
if (!open) {
setActiveStep(0)
setLoading(false)
setImageTag('')
}
}, [open])
const component = open ? (
<Dialog open={open}>
<DialogTitle>NIM Setup</DialogTitle>
<DialogContent>
<Stepper activeStep={activeStep}>
{steps.map((label) => (
<Step key={label}>
<StepLabel>{label}</StepLabel>
</Step>
))}
</Stepper>
{activeStep === 0 && (
<div style={{ marginTop: 20 }}>
<p style={{ marginBottom: 20 }}>
Would you like to download the NIM installer? Click Next if it has been installed
</p>
{loading && <CircularProgress />}
</div>
)}
{activeStep === 1 && (
<div>
<FormControl fullWidth sx={{ mt: 2 }}>
<InputLabel>Model</InputLabel>
<Select label='Model' value={imageTag} onChange={(e) => setImageTag(e.target.value)}>
{Object.entries(modelOptions).map(([value, { label }]) => (
<MenuItem key={value} value={value}>
{label}
</MenuItem>
))}
</Select>
</FormControl>
{imageTag && (
<Button
variant='text'
size='small'
sx={{ mt: 1 }}
onClick={() => window.open(modelOptions[imageTag].licenseUrl, '_blank')}
>
View License
</Button>
)}
{loading && (
<div>
<div style={{ marginBottom: 20 }} />
<CircularProgress />
<p>Pulling image...</p>
</div>
)}
</div>
)}
{activeStep === 2 && (
<div>
{loading ? (
<>
<div style={{ marginBottom: 20 }} />
<CircularProgress />
<p>Starting container...</p>
</>
) : (
<p>Image is ready! Click Next to start the container.</p>
)}
</div>
)}
</DialogContent>
<DialogActions>
<Button onClick={onClose} variant='outline'>
Cancel
</Button>
{activeStep === 0 && (
<Button onClick={handleNext} variant='outline' color='secondary'>
Next
</Button>
)}
<Button onClick={activeStep === 0 ? handleDownloadInstaller : handleNext} disabled={loading}>
{activeStep === 0 ? 'Download' : 'Next'}
</Button>
</DialogActions>
</Dialog>
) : null
return createPortal(component, portalElement)
}
NvidiaNIMDialog.propTypes = {
open: PropTypes.bool,
onClose: PropTypes.func,
onComplete: PropTypes.func
}
export default NvidiaNIMDialog
@@ -37,6 +37,7 @@ import PromptLangsmithHubDialog from '@/ui-component/dialog/PromptLangsmithHubDi
import ManageScrapedLinksDialog from '@/ui-component/dialog/ManageScrapedLinksDialog'
import CredentialInputHandler from './CredentialInputHandler'
import InputHintDialog from '@/ui-component/dialog/InputHintDialog'
import NvidiaNIMDialog from '@/ui-component/dialog/NvidiaNIMDialog'
// utils
import { getInputVariables, getCustomConditionOutputs, isValidConnection, getAvailableNodesForVariable } from '@/utils/genericHelper'
@@ -95,6 +96,7 @@ const NodeInputHandler = ({
const [inputHintDialogProps, setInputHintDialogProps] = useState({})
const [showConditionDialog, setShowConditionDialog] = useState(false)
const [conditionDialogProps, setConditionDialogProps] = useState({})
const [isNvidiaNIMDialogOpen, setIsNvidiaNIMDialogOpen] = useState(false)
const [tabValue, setTabValue] = useState(0)
const onInputHintDialogClicked = (hint) => {
@@ -443,6 +445,13 @@ const NodeInputHandler = ({
setAsyncOptionEditDialog('')
}
const handleNvidiaNIMDialogComplete = (containerData) => {
if (containerData) {
data.inputs['basePath'] = containerData.baseUrl
data.inputs['modelName'] = containerData.image
}
}
useEffect(() => {
if (ref.current && ref.current.offsetTop && ref.current.clientHeight) {
setPosition(ref.current.offsetTop + ref.current.clientHeight / 2)
@@ -528,6 +537,22 @@ const NodeInputHandler = ({
></PromptLangsmithHubDialog>
</>
)}
{data.name === 'chatNvidiaNIM' && inputParam.name === 'modelName' && (
<>
<Button
style={{
display: 'flex',
flexDirection: 'row',
width: '100%'
}}
sx={{ borderRadius: '12px', width: '100%', mb: 2, mt: -1 }}
variant='outlined'
onClick={() => setIsNvidiaNIMDialogOpen(true)}
>
Setup NIM Locally
</Button>
</>
)}
<div style={{ display: 'flex', flexDirection: 'row' }}>
<Typography>
{inputParam.label}
@@ -893,6 +918,11 @@ const NodeInputHandler = ({
dialogProps={inputHintDialogProps}
onCancel={() => setShowInputHintDialog(false)}
></InputHintDialog>
<NvidiaNIMDialog
open={isNvidiaNIMDialogOpen}
onClose={() => setIsNvidiaNIMDialogOpen(false)}
onComplete={handleNvidiaNIMDialogComplete}
></NvidiaNIMDialog>
</div>
)
}