feat: Nitro-Tensorrt-LLM Extension (janhq#2280)

* feat: tensorrt-llm-extension * fix: loading * feat: add download tensorrt llm runner Signed-off-by: James <[email protected]> * feat: update to rollupjs instead of webpack for monitoring extension Signed-off-by: James <[email protected]> * feat: move update nvidia info to monitor extension Signed-off-by: James <[email protected]> * allow download tensorrt Signed-off-by: James <[email protected]> * update Signed-off-by: James <[email protected]> * allow download tensor rt based on gpu setting Signed-off-by: James <[email protected]> * update downloaded models Signed-off-by: James <[email protected]> * feat: add extension compatibility * dynamic tensor rt engines Signed-off-by: James <[email protected]> * update models Signed-off-by: James <[email protected]> * chore: remove ts-ignore * feat: getting installation state from extension Signed-off-by: James <[email protected]> * chore: adding type for decompress Signed-off-by: James <[email protected]> * feat: update according Louis's comment Signed-off-by: James <[email protected]> * feat: add progress for installing extension Signed-off-by: James <[email protected]> * chore: remove args from extension installation * fix: model download does not work properly * fix: do not allow user to stop tensorrtllm inference * fix: extension installed style * fix: download tensorrt does not update state Signed-off-by: James <[email protected]> * chore: replace int4 by fl16 * feat: modal for installing extension Signed-off-by: James <[email protected]> * fix: start download immediately after press install Signed-off-by: James <[email protected]> * fix: error switching between engines * feat: rename inference provider to ai engine and refactor to core * fix: missing ulid * fix: core bundler * feat: add cancel extension installing Signed-off-by: James <[email protected]> * remove mocking for mac Signed-off-by: James <[email protected]> * fix: show models only when extension is ready * add tensorrt badge for model Signed-off-by: James <[email protected]> * fix: copy * fix: add compatible check (janhq#2342) * fix: add compatible check Signed-off-by: James <[email protected]> * fix: copy * fix: font * fix: copy * fix: broken monitoring extension * chore: bump engine * fix: copy * fix: model copy * fix: copy * fix: model json --------- Signed-off-by: James <[email protected]> Co-authored-by: James <[email protected]> Co-authored-by: Louis <[email protected]> * fix: vulkan support * fix: installation button padding * fix: empty script * fix: remove hard code string --------- Signed-off-by: James <[email protected]> Co-authored-by: James <[email protected]> Co-authored-by: NamH <[email protected]>
thunhuanh · Mar 14, 2024 · d85d026 · d85d026
1 parent 24c6dd0
commit d85d026
Show file tree

Hide file tree

Showing 71 changed files with 2,496 additions and 625 deletions.
diff --git a/.gitignore b/.gitignore
@@ -22,16 +22,16 @@ package-lock.json
 core/lib/**
 
 # Nitro binary files
-extensions/inference-nitro-extension/bin/*/nitro
-extensions/inference-nitro-extension/bin/*/*.metal
-extensions/inference-nitro-extension/bin/*/*.exe
-extensions/inference-nitro-extension/bin/*/*.dll
-extensions/inference-nitro-extension/bin/*/*.exp
-extensions/inference-nitro-extension/bin/*/*.lib
-extensions/inference-nitro-extension/bin/saved-*
-extensions/inference-nitro-extension/bin/*.tar.gz
-extensions/inference-nitro-extension/bin/vulkaninfoSDK.exe
-extensions/inference-nitro-extension/bin/vulkaninfo
+extensions/*-extension/bin/*/nitro
+extensions/*-extension/bin/*/*.metal
+extensions/*-extension/bin/*/*.exe
+extensions/*-extension/bin/*/*.dll
+extensions/*-extension/bin/*/*.exp
+extensions/*-extension/bin/*/*.lib
+extensions/*-extension/bin/saved-*
+extensions/*-extension/bin/*.tar.gz
+extensions/*-extension/bin/vulkaninfoSDK.exe
+extensions/*-extension/bin/vulkaninfo
 
 
 # Turborepo

diff --git a/core/package.json b/core/package.json
@@ -45,11 +45,12 @@
     "start": "rollup -c rollup.config.ts -w"
   },
   "devDependencies": {
-    "jest": "^29.7.0",
     "@types/jest": "^29.5.12",
     "@types/node": "^12.0.2",
-    "eslint-plugin-jest": "^27.9.0",
     "eslint": "8.57.0",
+    "eslint-plugin-jest": "^27.9.0",
+    "jest": "^29.7.0",
+    "rimraf": "^3.0.2",
     "rollup": "^2.38.5",
     "rollup-plugin-commonjs": "^9.1.8",
     "rollup-plugin-json": "^3.1.0",
@@ -58,7 +59,10 @@
     "rollup-plugin-typescript2": "^0.36.0",
     "ts-jest": "^29.1.2",
     "tslib": "^2.6.2",
-    "typescript": "^5.3.3",
-    "rimraf": "^3.0.2"
+    "typescript": "^5.3.3"
+  },
+  "dependencies": {
+    "rxjs": "^7.8.1",
+    "ulid": "^2.3.0"
   }
 }
diff --git a/core/rollup.config.ts b/core/rollup.config.ts
@@ -64,7 +64,7 @@ export default [
       // Allow json resolution
       json(),
       // Compile TypeScript files
-      typescript({ useTsconfigDeclarationDir: true }),
+      typescript({ useTsconfigDeclarationDir: true, exclude: ['src/*.ts', 'src/extensions/**'] }),
       // Allow bundling cjs modules (unlike webpack, rollup doesn't understand cjs)
       commonjs(),
       // Allow node_modules resolution, so you can use 'external' to control

diff --git a/core/src/api/index.ts b/core/src/api/index.ts
@@ -33,6 +33,8 @@ export enum AppRoute {
   stopServer = 'stopServer',
   log = 'log',
   logServer = 'logServer',
+  systemInformations = 'systemInformations',
+  showToast = 'showToast',
 }
 
 export enum AppEvent {
@@ -56,6 +58,7 @@ export enum DownloadEvent {
   onFileDownloadUpdate = 'onFileDownloadUpdate',
   onFileDownloadError = 'onFileDownloadError',
   onFileDownloadSuccess = 'onFileDownloadSuccess',
+  onFileUnzipSuccess = 'onFileUnzipSuccess',
 }
 
 export enum LocalImportModelEvent {

diff --git a/core/src/core.ts b/core/src/core.ts
@@ -1,4 +1,4 @@
-import { FileStat } from './types'
+import { DownloadRequest, FileStat, NetworkConfig } from './types'
 
 /**
  * Execute a extension module function in main process
@@ -17,18 +17,16 @@ const executeOnMain: (extension: string, method: string, ...args: any[]) => Prom
 
 /**
  * Downloads a file from a URL and saves it to the local file system.
- * @param {string} url - The URL of the file to download.
- * @param {string} fileName - The name to use for the downloaded file.
- * @param {object} network - Optional object to specify proxy/whether to ignore SSL certificates.
+ *
+ * @param {DownloadRequest} downloadRequest - The request to download the file.
+ * @param {NetworkConfig} network - Optional object to specify proxy/whether to ignore SSL certificates.
+ *
  * @returns {Promise<any>} A promise that resolves when the file is downloaded.
  */
-const downloadFile: (
-  url: string,
-  fileName: string,
-  network?: { proxy?: string; ignoreSSL?: boolean }
-) => Promise<any> = (url, fileName, network) => {
-  return global.core?.api?.downloadFile(url, fileName, network)
-}
+const downloadFile: (downloadRequest: DownloadRequest, network?: NetworkConfig) => Promise<any> = (
+  downloadRequest,
+  network
+) => global.core?.api?.downloadFile(downloadRequest, network)
 
 /**
  * Aborts the download of a specific file.
@@ -108,6 +106,20 @@ const log: (message: string, fileName?: string) => void = (message, fileName) =>
 const isSubdirectory: (from: string, to: string) => Promise<boolean> = (from: string, to: string) =>
   global.core.api?.isSubdirectory(from, to)
 
+/**
+ * Get system information
+ * @returns {Promise<any>} - A promise that resolves with the system information.
+ */
+const systemInformations: () => Promise<any> = () => global.core.api?.systemInformations()
+
+/**
+ * Show toast message from browser processes.
+ * @param title
+ * @param message
+ * @returns
+ */
+const showToast: (title: string, message: string) => void = (title, message) =>
+  global.core.api?.showToast(title, message)
 /**
  * Register extension point function type definition
  */
@@ -134,5 +146,7 @@ export {
   log,
   isSubdirectory,
   getUserHomePath,
+  systemInformations,
+  showToast,
   FileStat,
 }
diff --git a/core/src/extension.ts b/core/src/extension.ts
@@ -10,6 +10,22 @@ export enum ExtensionTypeEnum {
 export interface ExtensionType {
   type(): ExtensionTypeEnum | undefined
 }
+
+export interface Compatibility {
+  platform: string[]
+  version: string
+}
+
+const ALL_INSTALLATION_STATE = [
+  'NotRequired', // not required.
+  'Installed', // require and installed. Good to go.
+  'NotInstalled', // require to be installed.
+  'Corrupted', // require but corrupted. Need to redownload.
+] as const
+
+export type InstallationStateTuple = typeof ALL_INSTALLATION_STATE
+export type InstallationState = InstallationStateTuple[number]
+
 /**
  * Represents a base extension.
  * This class should be extended by any class that represents an extension.
@@ -33,4 +49,32 @@ export abstract class BaseExtension implements ExtensionType {
    * Any cleanup logic for the extension should be put here.
    */
   abstract onUnload(): void
+
+  /**
+   * The compatibility of the extension.
+   * This is used to check if the extension is compatible with the current environment.
+   * @property {Array} platform
+   */
+  compatibility(): Compatibility | undefined {
+    return undefined
+  }
+
+  /**
+   * Determine if the prerequisites for the extension are installed.
+   *
+   * @returns {boolean} true if the prerequisites are installed, false otherwise.
+   */
+  async installationState(): Promise<InstallationState> {
+    return 'NotRequired'
+  }
+
+  /**
+   * Install the prerequisites for the extension.
+   *
+   * @returns {Promise<void>}
+   */
+  // @ts-ignore
+  async install(...args): Promise<void> {
+    return
+  }
 }
diff --git a/core/src/extensions/ai-engines/AIEngine.ts b/core/src/extensions/ai-engines/AIEngine.ts
@@ -0,0 +1,60 @@
+import { getJanDataFolderPath, joinPath } from '../../core'
+import { events } from '../../events'
+import { BaseExtension } from '../../extension'
+import { fs } from '../../fs'
+import { Model, ModelEvent } from '../../types'
+
+/**
+ * Base AIEngine
+ * Applicable to all AI Engines
+ */
+export abstract class AIEngine extends BaseExtension {
+  // The inference engine
+  abstract provider: string
+  // The model folder
+  modelFolder: string = 'models'
+
+  abstract models(): Promise<Model[]>
+
+  /**
+   * On extension load, subscribe to events.
+   */
+  onLoad() {
+    this.prePopulateModels()
+  }
+
+  /**
+   * Pre-populate models to App Data Folder
+   */
+  prePopulateModels(): Promise<void> {
+    return this.models().then((models) => {
+      const prePoluateOperations = models.map((model) =>
+        getJanDataFolderPath()
+          .then((janDataFolder) =>
+            // Attempt to create the model folder
+            joinPath([janDataFolder, this.modelFolder, model.id]).then((path) =>
+              fs
+                .mkdirSync(path)
+                .catch()
+                .then(() => path)
+            )
+          )
+          .then((path) => joinPath([path, 'model.json']))
+          .then((path) => {
+            // Do not overwite existing model.json
+            return fs.existsSync(path).then((exist: any) => {
+              if (!exist) return fs.writeFileSync(path, JSON.stringify(model, null, 2))
+            })
+          })
+          .catch((e: Error) => {
+            console.error('Error', e)
+          })
+      )
+      Promise.all(prePoluateOperations).then(() =>
+        // Emit event to update models
+        // So the UI can update the models list
+        events.emit(ModelEvent.OnModelsUpdate, {})
+      )
+    })
+  }
+}
diff --git a/core/src/extensions/ai-engines/LocalOAIEngine.ts b/core/src/extensions/ai-engines/LocalOAIEngine.ts
@@ -0,0 +1,63 @@
+import { executeOnMain, getJanDataFolderPath, joinPath } from '../../core'
+import { events } from '../../events'
+import { Model, ModelEvent } from '../../types'
+import { OAIEngine } from './OAIEngine'
+
+/**
+ * Base OAI Local Inference Provider
+ * Added the implementation of loading and unloading model (applicable to local inference providers)
+ */
+export abstract class LocalOAIEngine extends OAIEngine {
+  // The inference engine
+  loadModelFunctionName: string = 'loadModel'
+  unloadModelFunctionName: string = 'unloadModel'
+  isRunning: boolean = false
+
+  /**
+   * On extension load, subscribe to events.
+   */
+  onLoad() {
+    super.onLoad()
+    // These events are applicable to local inference providers
+    events.on(ModelEvent.OnModelInit, (model: Model) => this.onModelInit(model))
+    events.on(ModelEvent.OnModelStop, (model: Model) => this.onModelStop(model))
+  }
+
+  /**
+   * Load the model.
+   */
+  async onModelInit(model: Model) {
+    if (model.engine.toString() !== this.provider) return
+
+    const modelFolder = await joinPath([await getJanDataFolderPath(), this.modelFolder, model.id])
+
+    const res = await executeOnMain(this.nodeModule, this.loadModelFunctionName, {
+      modelFolder,
+      model,
+    })
+
+    if (res?.error) {
+      events.emit(ModelEvent.OnModelFail, {
+        ...model,
+        error: res.error,
+      })
+      return
+    } else {
+      this.loadedModel = model
+      events.emit(ModelEvent.OnModelReady, model)
+      this.isRunning = true
+    }
+  }
+  /**
+   * Stops the model.
+   */
+  onModelStop(model: Model) {
+    if (model.engine?.toString() !== this.provider) return
+
+    this.isRunning = false
+
+    executeOnMain(this.nodeModule, this.unloadModelFunctionName).then(() => {
+      events.emit(ModelEvent.OnModelStopped, {})
+    })
+  }
+}