feat: Nitro-Tensorrt-LLM Extension

[louis-menlo]

Describe Your Changes

The extension mainly focuses on distributing its model and running the engine.
NOTES: I have not yet moved to the core, but I will align very soon.

graph LR

TensorRTLLMExtension[TensorRTLLMExtension] -- implements --> DefineModels[Define Models]
TensorRTLLMExtension -- extends --> LocalOAIEngine[LocalOAIEngine]
LocalOAIEngine[LocalOAIEngine] -- implements  --> StartModel[Start Model]
LocalOAIEngine[LocalOAIEngine] -- implements  --> StopModel[Stop Model]
LocalOAIEngine-- extends --> OAIEngine[OAIEngine] 
OAIEngine[OAIEngine]  -- implements --> Inference[Inference]
OAIEngine[OAIEngine] -- implements --> Stop[Stop Inference]
OAIEngine[OAIEngine] -- extends --> AIEngine[AIEngine]
AIEngine[AIEngine]  -- implements --> Populate[Pre-populate Models]

Loading

Main controller (index)

/**
 * A class that implements the InferenceExtension interface from the @janhq/core package.
 * The class provides methods for initializing and stopping a model, and for making inference requests.
 * It also subscribes to events emitted by the @janhq/core package and handles new message requests.
 */
export default class TensorRTLLMExtension extends LocalOAIEngine{
  override provider = 'nitro-tensorrt-llm'

  // Configure node module for updating by Rollup bundler.
  override inference_url = INFERENCE_URL
  override nodeModule = NODE

  /**
   * models implemented by the extension
   * define pre-populated models
   */
  models(): Model[] {
    return models as unknown as Model[]
  }
}

Node module (run engine)

graph LR

LoadModel[Load Model] -- kill port process --> TerminateCheck
TerminateCheck -->|Success| Nitro[Run Nitro]
TerminateCheck -->|Failed \n Port owned by another process| Terminated[Terminated] 

Nitro --> LoadRequest[Send Load Model Request]
LoadRequest --> Done[Done]

Loading

/**
 * Initializes a engine subprocess to load a machine learning model.
 * @param params - The model load settings.
 */
async function loadModel(params: any): Promise<{ error: Error | undefined }> {
  const settings: ModelLoadParams = {
    engine_path: params.modelFolder,
    ctx_len: params.model.settings.ctx_len ?? 2048,
  }
  return runEngine().then(() => loadModelRequest(settings))
}


/**
 * Loads a LLM model into the Engine subprocess by sending a HTTP POST request.
 */
function loadModelRequest(settings: ModelLoadParams): Promise<any> {
  return fetchRetry(LOAD_MODEL_URL, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
    },
    body: JSON.stringify(settings),
    retries: 3,
    retryDelay: 500,
  }).then((res) => { return Promise.resolve({ error: undefined }) })
}

/**
 * Spawns engine subprocess.
 */
function async runEngine() {
    // Current directory by default
    let binaryFolder = path.join(__dirname, '..', 'bin')

    // Binary path
    const binary = path.join(binaryFolder, process.platform === 'win32' ? 'nitro.exe' : 'nitro')

    const args: string[] = ['1', ENGINE_HOST, ENGINE_PORT]
    // Execute the binary
    subprocess = spawn(binary, args, { cwd: binaryFolder })

    // Wait for engine ready
    await tcpPortUsed.waitUntilUsed(parseInt(ENGINE_PORT), 300, 30000)
}

Fixes Issues

• Closes #

Self Checklist

• Base Inference Provider
• Distribute models via the extension
• Download according assets (GPU specific)