dtamind-nim-container-manager/index.js

NVIDIA NIM Container Manager for Dtamind

const Docker = require('dockerode')
const axios = require('axios')

class NimContainerManager {
    constructor() {
        this.docker = new Docker()
        this.baseUrl = 'https://nts.ngc.nvidia.com/v1'
    }

    /**
     * Get authentication token for NVIDIA NGC
     */
    async getToken() {
        try {
            const headers = {
                'Content-Type': 'application/json',
                Accept: 'application/json'
            }
            const data = {
                client_id: 'Dtamind', // Changed from 'Flowise' to 'Dtamind'
                pdi: '0x1234567890abcdeg',
                access_policy_name: 'nim-dev'
            }
            const response = await axios.post(`${this.baseUrl}/token`, data, { headers })
            return response.data
        } catch (error) {
            throw new Error(`Failed to get token: ${error.message}`)
        }
    }

    /**
     * Preload NIM containers
     */
    async preload() {
        try {
            // Implementation for preloading NIM containers
            console.log('Preloading NIM containers...')
            return 'Preloaded NIM'
        } catch (error) {
            throw new Error(`Failed to preload: ${error.message}`)
        }
    }

    /**
     * Download NIM installer
     */
    async downloadInstaller() {
        try {
            // Implementation for downloading installer
            console.log('Downloading NIM installer...')
            return 'NIM Installer completed successfully!'
        } catch (error) {
            throw new Error(`Failed to download installer: ${error.message}`)
        }
    }

    /**
     * Pull NIM image
     */
    async pullImage(imageTag, apiKey) {
        try {
            // Implementation for pulling image
            console.log(`Pulling image ${imageTag}...`)
            return `Pulling image ${imageTag}`
        } catch (error) {
            throw new Error(`Failed to pull image: ${error.message}`)
        }
    }

    /**
     * Start NIM container
     */
    async startContainer(imageTag, apiKey, hostPort, nimRelaxMemConstraints) {
        try {
            // Validate nimRelaxMemConstraints
            if (isNaN(nimRelaxMemConstraints) || (nimRelaxMemConstraints !== 0 && nimRelaxMemConstraints !== 1)) {
                throw new Error('nimRelaxMemConstraints must be 0 or 1')
            }

            // Implementation for starting container
            console.log(`Starting container ${imageTag} on port ${hostPort}...`)
            return `Starting container ${imageTag}`
        } catch (error) {
            throw new Error(`Failed to start container: ${error.message}`)
        }
    }

    /**
     * Get NIM image information
     */
    async getImage(imageTag) {
        try {
            const images = await this.userImageLibrary()
            const image = images.find((img) => img.tag === imageTag)
            if (!image) {
                throw new Error(`Image ${imageTag} not found`)
            }
            return image
        } catch (error) {
            throw new Error(`Failed to get image: ${error.message}`)
        }
    }

    /**
     * Get container information
     */
    async getContainer(imageTag, port) {
        try {
            // First check if the image exists
            const images = await this.userImageLibrary()
            const image = images.find((img) => img.tag === imageTag)
            if (!image) {
                throw new Error(`Image ${imageTag} not found`)
            }

            const containers = await this.listRunningContainers()
            const portInUse = containers.find((cont) => cont.port === port)
            if (portInUse) {
                const isModelContainer = portInUse.image === image.tag
                return {
                    container: portInUse,
                    isModelContainer
                }
            }

            return null
        } catch (error) {
            throw new Error(`Failed to get container: ${error.message}`)
        }
    }

    /**
     * List running containers
     */
    async listRunningContainers() {
        try {
            const containers = await this.docker.listContainers()
            return containers.map(container => ({
                id: container.Id,
                name: container.Names[0],
                image: container.Image,
                status: container.Status,
                port: container.Ports[0]?.PublicPort || null
            }))
        } catch (error) {
            throw new Error(`Failed to list containers: ${error.message}`)
        }
    }

    /**
     * Stop container
     */
    async stopContainer(imageTag) {
        try {
            const containers = await this.listRunningContainers()
            const container = containers.find(cont => cont.image === imageTag)
            
            if (container) {
                const dockerContainer = this.docker.getContainer(container.id)
                await dockerContainer.stop()
                return `Stopped container ${imageTag}`
            } else {
                throw new Error(`Container for image ${imageTag} not found`)
            }
        } catch (error) {
            throw new Error(`Failed to stop container: ${error.message}`)
        }
    }

    /**
     * Get user image library
     */
    async userImageLibrary() {
        try {
            // Implementation for getting user image library
            // This would typically call NVIDIA NGC API
            return []
        } catch (error) {
            throw new Error(`Failed to get image library: ${error.message}`)
        }
    }
}

module.exports = {
    NimContainerManager
}