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
2025-02-27 09:14:01 +00:00 · 2025-02-27 09:14:01 +00:00 · 1678815540
parent 7d8541a44b
commit 1678815540
10 changed files with 606 additions and 1056 deletions
--- a/packages/components/credentials/NvdiaNIMApi.credential.ts
+++ b/packages/components/credentials/NvdiaNIMApi.credential.ts
@ -1,6 +1,6 @@
 import { INodeParams, INodeCredential } from '../src/Interface'

-class NvdiaNIMApi implements INodeCredential {
+class NvidiaNIMApi implements INodeCredential {
    label: string
    name: string
    version: number
@ -9,16 +9,16 @@ class NvdiaNIMApi implements INodeCredential {

    constructor() {
        this.label = 'Nvdia NIM API Key'
-        this.name = 'nvdiaNIMApi'
+        this.name = 'nvidiaNIMApi'
        this.version = 1.0
        this.inputs = [
            {
-                label: 'Nvdia NIM API Key',
-                name: 'nvdiaNIMApiKey',
+                label: 'Nvidia NIM API Key',
+                name: 'nvidiaNIMApiKey',
                type: 'password'
            }
        ]
    }
 }

-module.exports = { credClass: NvdiaNIMApi }
+module.exports = { credClass: NvidiaNIMApi }
--- a/packages/components/nodes/chatmodels/ChatNvdiaNIM/ChatNvdiaNIM.ts
+++ b/packages/components/nodes/chatmodels/ChatNvdiaNIM/ChatNvdiaNIM.ts
@ -16,10 +16,10 @@ class ChatNvdiaNIM_ChatModels implements INode {
    inputs: INodeParams[]

    constructor() {
-        this.label = 'Chat Nvdia NIM'
-        this.name = 'chatNvdiaNIM'
+        this.label = 'Chat Nvidia NIM'
+        this.name = 'chatNvidiaNIM'
        this.version = 1.0
-        this.type = 'ChatNvdiaNIM'
+        this.type = 'ChatNvidiaNIM'
        this.icon = 'nvdia.svg'
        this.category = 'Chat Models'
        this.description = 'Wrapper around Nvdia NIM Inference API'
@ -28,7 +28,7 @@ class ChatNvdiaNIM_ChatModels implements INode {
            label: 'Connect Credential',
            name: 'credential',
            type: 'credential',
-            credentialNames: ['nvdiaNIMApi'],
+            credentialNames: ['nvidiaNIMApi'],
            optional: true
        }
        this.inputs = [
@ -44,6 +44,13 @@ class ChatNvdiaNIM_ChatModels implements INode {
                type: 'string',
                placeholder: 'microsoft/phi-3-mini-4k-instruct'
            },
+            {
+                label: 'Base Path',
+                name: 'basePath',
+                type: 'string',
+                description: 'Specify the URL of the deployed NIM Inference API',
+                placeholder: 'https://integrate.api.nvidia.com/v1'
+            },
            {
                label: 'Temperature',
                name: 'temperature',
@ -52,13 +59,6 @@ class ChatNvdiaNIM_ChatModels implements INode {
                default: 0.9,
                optional: true
            },
-            {
-                label: 'Base Path',
-                name: 'basePath',
-                type: 'string',
-                description: 'Specify the URL of the deployed NIM Inference API',
-                placeholder: 'https://integrate.api.nvidia.com/v1'
-            },
            {
                label: 'Streaming',
                name: 'streaming',
@ -131,12 +131,12 @@ class ChatNvdiaNIM_ChatModels implements INode {
        const cache = nodeData.inputs?.cache as BaseCache

        const credentialData = await getCredentialData(nodeData.credential ?? '', options)
-        const nvdiaNIMApiKey = getCredentialParam('nvdiaNIMApiKey', credentialData, nodeData)
+        const nvidiaNIMApiKey = getCredentialParam('nvidiaNIMApiKey', credentialData, nodeData)

        const obj: ChatOpenAIFields & { nvdiaNIMApiKey?: string } = {
            temperature: parseFloat(temperature),
            modelName,
-            openAIApiKey: nvdiaNIMApiKey,
+            openAIApiKey: nvidiaNIMApiKey ?? 'sk-',
            streaming: streaming ?? true
        }

--- a/packages/server/package.json
+++ b/packages/server/package.json
@ -97,6 +97,7 @@
        "multer": "^1.4.5-lts.1",
        "multer-s3": "^3.0.1",
        "mysql2": "^3.11.3",
+        "nim-container-manager": "^1.0.4",
        "openai": "^4.82.0",
        "pg": "^8.11.1",
        "posthog-node": "^3.5.0",
--- a/packages/server/src/controllers/nvidia-nim/index.ts
+++ b/packages/server/src/controllers/nvidia-nim/index.ts
@ -0,0 +1,106 @@
+import axios from 'axios'
+import { Request, Response, NextFunction } from 'express'
+
+const { NimContainerManager } = require('nim-container-manager')
+
+const getToken = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const headers = {
+            'Content-Type': 'application/json',
+            Accept: 'application/json'
+        }
+        const data = {
+            client_id: 'Flowise',
+            pdi: '0x1234567890abcdeg',
+            access_policy_name: 'nim-dev'
+        }
+        const response = await axios.post('https://nts.ngc.nvidia.com/v1/token', data, { headers })
+        const responseJson = response.data
+        return res.json(responseJson)
+    } catch (error) {
+        next(error)
+    }
+}
+
+const preload = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        await NimContainerManager.preload()
+        return res.send('Preloaded NIM')
+    } catch (error) {
+        next(error)
+    }
+}
+
+const downloadInstaller = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        await NimContainerManager.downloadInstaller()
+        return res.send('NIM Installer completed successfully!')
+    } catch (error) {
+        next(error)
+    }
+}
+
+const pullImage = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const imageTag = req.body.imageTag
+        const apiKey = req.body.apiKey
+        await NimContainerManager.pullImage(imageTag, apiKey)
+        return res.send(`Pulling image ${imageTag}`)
+    } catch (error) {
+        next(error)
+    }
+}
+
+const startContainer = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const imageTag = req.body.imageTag
+        const apiKey = req.body.apiKey
+        await NimContainerManager.startContainer(imageTag, apiKey)
+        return res.send(`Starting container ${imageTag}`)
+    } catch (error) {
+        next(error)
+    }
+}
+
+const getImage = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const imageTag = req.body.imageTag
+        const images = await NimContainerManager.userImageLibrary()
+        const image = images.find((img: any) => img.tag === imageTag)
+        if (!image) {
+            return res.status(404).send(`Image ${imageTag} not found`)
+        }
+        return res.json(image)
+    } catch (error) {
+        next(error)
+    }
+}
+
+const getContainer = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const imageTag = req.body.imageTag
+        const images = await NimContainerManager.userImageLibrary()
+        const image = images.find((img: any) => img.tag === imageTag)
+        if (!image) {
+            return res.status(404).send(`Image ${imageTag} not found`)
+        }
+        if (!image.container) {
+            return res.status(404).send(`Container of ${imageTag} not found`)
+        }
+        const container = image.container
+        container.image = image.name
+        return res.json(container)
+    } catch (error) {
+        next(error)
+    }
+}
+
+export default {
+    preload,
+    getToken,
+    downloadInstaller,
+    pullImage,
+    startContainer,
+    getImage,
+    getContainer
+}
--- a/packages/server/src/routes/index.ts
+++ b/packages/server/src/routes/index.ts
@ -42,6 +42,7 @@ import variablesRouter from './variables'
 import vectorRouter from './vectors'
 import verifyRouter from './verify'
 import versionRouter from './versions'
+import nvidiaNimRouter from './nvidia-nim'

 const router = express.Router()

@ -88,5 +89,6 @@ router.use('/vector', vectorRouter)
 router.use('/verify', verifyRouter)
 router.use('/version', versionRouter)
 router.use('/upsert-history', upsertHistoryRouter)
+router.use('/nvidia-nim', nvidiaNimRouter)

 export default router
--- a/packages/server/src/routes/nvidia-nim/index.ts
+++ b/packages/server/src/routes/nvidia-nim/index.ts
@ -0,0 +1,14 @@
+import express from 'express'
+import nimController from '../../controllers/nvidia-nim'
+const router = express.Router()
+
+// READ
+router.get('/preload', nimController.preload)
+router.get('/get-token', nimController.getToken)
+router.get('/download-installer', nimController.downloadInstaller)
+router.post('/pull-image', nimController.pullImage)
+router.post('/start-container', nimController.startContainer)
+router.post('/get-image', nimController.getImage)
+router.post('/get-container', nimController.getContainer)
+
+export default router
--- a/packages/server/src/utils/constants.ts
+++ b/packages/server/src/utils/constants.ts
@ -17,7 +17,8 @@ export const WHITELIST_URLS = [
    '/api/v1/ping',
    '/api/v1/version',
    '/api/v1/attachments',
-    '/api/v1/metrics'
+    '/api/v1/metrics',
+    '/api/v1/nvidia-nim'
 ]

 export const OMIT_QUEUE_JOB_DATA = ['componentNodes', 'appDataSource', 'sseStreamer', 'telemetry', 'cachePool']
--- a/packages/ui/src/ui-component/dialog/NvidiaNIMDialog.jsx
+++ b/packages/ui/src/ui-component/dialog/NvidiaNIMDialog.jsx
@ -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
--- a/packages/ui/src/views/canvas/NodeInputHandler.jsx
+++ b/packages/ui/src/views/canvas/NodeInputHandler.jsx
@ -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>
    )
 }
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml