[janhq/jan#1635] feat: decouple nitro engine into a library

.github/workflows/build.yml

@@ -346,7 +346,7 @@ jobs:
           ./install_deps.sh
           mkdir build && cd build
           cmake -DWHISPER_COREML=1 -DNITRO_VERSION=${{ needs.set-nitro-version.outputs.version }} .. 
-          CC=gcc-8 make -j $(sysctl -n hw.ncp)
+          CC=gcc-8 make -j $(sysctl -n hw.ncpu)
           ls -la
 
       - name: Package

.gitignore

@@ -86,6 +86,7 @@ CMakeFiles
 CMakeScripts
 Testing
 Makefile
+!nitro-node/Makefile
 cmake_install.cmake
 install_manifest.txt
 compile_commands.json

nitro-node/.gitignore

@@ -0,0 +1,19 @@
+.vscode
+.env
+
+# Jan inference
+.yarn
+node_modules
+*.tgz
+yarn.lock
+dist
+build
+.DS_Store
+package-lock.json
+
+*.log
+
+# Nitro binary directory
+bin/
+# Locally downloaded model for testing
+test/test_assets/*.gguf

nitro-node/.npmrc

@@ -0,0 +1,2 @@
+scripts-prepend-node-path=true
+engine-strict=true

nitro-node/.yarnrc.yml

@@ -0,0 +1 @@
+nodeLinker: node-modules

nitro-node/Makefile

@@ -0,0 +1,51 @@
+# Makefile for Nitro-node - Build and Clean
+
+# Default target, build all
+.PHONY: all
+
+all: publish
+
+# Installs yarn dependencies
+#install: build-core
+install:
+ifeq ($(OS),Windows_NT)
+	yarn config set network-timeout 300000
+endif
+	yarn install
+
+# Build
+build: install
+	yarn run build
+
+# Download Nitro
+download-nitro: install
+	yarn run downloadnitro
+
+test-ci: install
+	yarn test
+
+# Note, this make target is just for testing on *NIX systems
+test: install
+	@test -e test/test_assets/*.gguf &&  echo "test/test_assets/*.gguf is already downloaded" || (mkdir -p test/test_assets && cd test/test_assets/ && curl -JLO "https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF/resolve/main/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf")
+	yarn test
+
+# Builds and pack
+pack: build
+	yarn run build:publish
+
+# Test that installation will also download nitro binaries
+test-e2e-installation: pack
+ifeq ($(OS),Windows_NT)
+	$$env:NITRO_NODE_VERSION=(npm version --json | jq '.["@janhq/nitro-node"]' | foreach {$$_.replace('"','')}) $$env:NITRO_NODE_PKG=(Resolve-Path -Path "janhq-nitro-node$$NITRO_NODE_VERSION.tgz") node ..\.github\scripts\e2e-test-install-nitro-node.js
+else
+	NITRO_NODE_VERSION=$$(npm version --json | jq '.["@janhq/nitro-node"]' | tr -d '"') NITRO_NODE_PKG=$$(realpath "janhq-nitro-node-$${NITRO_NODE_VERSION}.tgz") node ../.github/scripts/e2e-test-install-nitro-node.js
+endif
+
+clean:
+ifeq ($(OS),Windows_NT)
+	powershell -Command "Remove-Item -Recurse -Force -Path *.tgz, .yarn, yarn.lock, package-lock.json, bin, dist"
+	powershell -Command "Get-ChildItem -Path . -Include node_modules -Recurse -Directory | Remove-Item -Recurse -Force"
+else
+	rm -rf *.tgz .yarn yarn.lock package-lock.json bin dist
+	find . -name "node_modules" -type d -prune -exec rm -rf '{}' +
+endif

nitro-node/README.md

@@ -0,0 +1,3 @@
+# NodeJS wrapper for Nitro
+
+**TODO** Documenting on usage

nitro-node/jest.config.ts

@@ -0,0 +1,14 @@
+import type { JestConfigWithTsJest } from "ts-jest";
+
+const jestConfig: JestConfigWithTsJest = {
+  preset: "ts-jest",
+  testEnvironment: "node",
+  transformIgnorePatterns: ["/node_modules/"],
+  globals: {
+    RELEASE_URL_PREFIX: "https://api.github.com/repos/janhq/nitro/releases/",
+    TAGGED_RELEASE_URL_PREFIX:
+      "https://api.github.com/repos/janhq/nitro/releases/tags",
+  },
+};
+
+export default jestConfig;

nitro-node/package.json

@@ -0,0 +1,78 @@
+{
+  "name": "@janhq/nitro-node",
+  "version": "1.0.0",
+  "description": "This NodeJS library is a wrapper for Nitro, a lightweight (3mb) inference engine written in C++. See nitro.jan.ai",
+  "main": "./dist/index",
+  "module": "./dist/index.esm.js",
+  "types": "./dist/types/index",
+  "author": "Jan <service@jan.ai>",
+  "license": "AGPL-3.0",
+  "scripts": {
+    "test": "jest --verbose --detectOpenHandles",
+    "build": "tsc --module commonjs && rollup -c rollup.config.ts",
+    "predownloadnitro": "npm run build",
+    "downloadnitro": "node dist/scripts/index.cjs",
+    "build:publish": "npm pack",
+    "postinstall": "node -r @janhq/nitro-node/postinstall"
+  },
+  "exports": {
+    ".": {
+      "import": "./dist/index.esm.js",
+      "require": "./dist/index.cjs",
+      "types": "./dist/types/index.d.ts"
+    },
+    "./scripts": {
+      "import": "./dist/scripts/index.esm.js",
+      "require": "./dist/scripts/index.cjs",
+      "types": "./dist/types/scripts/index.d.ts"
+    },
+    "./postinstall": "./postinstall.js"
+  },
+  "typesVersions": {
+    "*": {
+      ".": [
+        "./dist/index.esm.js.map",
+        "./dist/types/index.d.ts"
+      ],
+      "scripts": [
+        "./dist/scripts/index.esm.js.map",
+        "./dist/types/scripts/index.d.ts"
+      ]
+    }
+  },
+  "devDependencies": {
+    "@rollup/plugin-commonjs": "^25.0.7",
+    "@rollup/plugin-json": "^6.1.0",
+    "@rollup/plugin-node-resolve": "^15.2.3",
+    "@rollup/plugin-replace": "^5.0.5",
+    "@types/download": "^8.0.5",
+    "@types/jest": "^29.5.11",
+    "@types/node": "^20.11.4",
+    "@types/os-utils": "^0.0.4",
+    "@types/tcp-port-used": "^1.0.4",
+    "jest": "^29.7.0",
+    "rollup": "^2.38.5",
+    "rollup-plugin-sourcemaps": "^0.6.3",
+    "rollup-plugin-typescript2": "^0.36.0",
+    "ts-jest": "^29.1.2",
+    "ts-node": "^10.9.2",
+    "typescript": "^5.3.3"
+  },
+  "dependencies": {
+    "cross-fetch": "^4.0.0",
+    "download": "^8.0.0",
+    "fetch-retry": "^5.0.6",
+    "os-utils": "^0.0.14",
+    "tcp-port-used": "^1.0.2"
+  },
+  "engines": {
+    "node": ">=18.0.0"
+  },
+  "files": [
+    "postinstall.js",
+    "dist",
+    "package.json",
+    "tsconfig.json",
+    "README.md"
+  ]
+}

nitro-node/postinstall.js

@@ -0,0 +1,6 @@
+// Only run if this package is installed as dependency
+if (process.env.INIT_CWD === process.cwd()) process.exit();
+
+const path = require("node:path");
+const { downloadNitro } = require("@janhq/nitro-node/scripts");
+downloadNitro(path.join(__dirname, "bin"));

nitro-node/rollup.config.ts

@@ -0,0 +1,109 @@
+import resolve from "@rollup/plugin-node-resolve";
+import commonjs from "@rollup/plugin-commonjs";
+import sourceMaps from "rollup-plugin-sourcemaps";
+import typescript from "rollup-plugin-typescript2";
+import json from "@rollup/plugin-json";
+import replace from "@rollup/plugin-replace";
+
+export default [
+  {
+    input: `src/index.ts`,
+    output: [
+      { file: "dist/index.cjs", format: "cjs", sourcemap: true },
+      { file: "dist/index.esm.js", format: "es", sourcemap: true },
+    ],
+    // Indicate here external modules you don't wanna include in your bundle (i.e.: 'lodash')
+    external: [
+      // `download@8.0.0` requires `got@^8.3.1` which then optionally requires `electron`, result in wrong dependency
+      // Ref: https://github.com/kubernetes-client/javascript/issues/350#issue-500860208
+      // Ref: https://github.com/kubernetes-client/javascript/issues/350#issuecomment-553644659
+      "got",
+    ],
+    watch: {
+      include: "src/**",
+    },
+    plugins: [
+      replace({
+        RELEASE_URL_PREFIX: JSON.stringify(
+          "https://api.github.com/repos/janhq/nitro/releases/",
+        ),
+        TAGGED_RELEASE_URL_PREFIX: JSON.stringify(
+          "https://api.github.com/repos/janhq/nitro/releases/tags",
+        ),
+      }),
+      // Allow json resolution
+      json(),
+      // Allow node_modules resolution, so you can use 'external' to control
+      // which external modules to include in the bundle
+      // https://github.com/rollup/rollup-plugin-node-resolve#usage
+      resolve({
+        extensions: [".ts", ".js", ".json"],
+      }),
+      // Allow bundling cjs modules (unlike webpack, rollup doesn't understand cjs)
+      // This should be after resolve() plugin
+      commonjs(),
+      // Compile TypeScript files
+      typescript({
+        useTsconfigDeclarationDir: true,
+      }),
+
+      // Resolve source maps to the original source
+      sourceMaps(),
+    ],
+  },
+  {
+    input: `src/scripts/index.ts`,
+    output: [
+      {
+        file: "dist/scripts/index.cjs",
+        format: "cjs",
+        sourcemap: true,
+      },
+      {
+        file: "dist/scripts/index.esm.js",
+        format: "es",
+        sourcemap: true,
+      },
+    ],
+    // Indicate here external modules you don't wanna include in your bundle (i.e.: 'lodash')
+    external: [
+      // `download@8.0.0` requires `got@^8.3.1` which then optionally requires `electron`, result in wrong dependency
+      // Ref: https://github.com/kubernetes-client/javascript/issues/350#issue-500860208
+      // Ref: https://github.com/kubernetes-client/javascript/issues/350#issuecomment-553644659
+      "got",
+    ],
+    watch: {
+      include: "src/scripts/**",
+    },
+    plugins: [
+      replace({
+        RELEASE_URL_PREFIX: JSON.stringify(
+          "https://api.github.com/repos/janhq/nitro/releases/",
+        ),
+        TAGGED_RELEASE_URL_PREFIX: JSON.stringify(
+          "https://api.github.com/repos/janhq/nitro/releases/tags",
+        ),
+      }),
+      // Allow json resolution
+      json(),
+
+      // Allow node_modules resolution, so you can use 'external' to control
+      // which external modules to include in the bundle
+      // https://github.com/rollup/rollup-plugin-node-resolve#usage
+      resolve({
+        extensions: [".ts", ".js", ".json"],
+      }),
+
+      // Allow bundling cjs modules (unlike webpack, rollup doesn't understand cjs)
+      // This should be after resolve() plugin
+      commonjs(),
+      // Compile TypeScript files
+      typescript({
+        useTsconfigDeclarationDir: true,
+      }),
+
+      // Resolve source maps to the original source
+      sourceMaps(),
+    ],
+  },
+];

nitro-node/src/execute.ts

@@ -0,0 +1,64 @@
+import path from "node:path";
+import { getNvidiaConfig } from "./nvidia";
+
+export interface NitroExecutableOptions {
+  executablePath: string;
+  cudaVisibleDevices: string;
+}
+
+/**
+ * Find which executable file to run based on the current platform.
+ * @returns The name of the executable file to run.
+ */
+export const executableNitroFile = (
+  binaryFolder: string,
+  // Default to GPU if CUDA is available when calling
+  runMode: "cpu" | "gpu" = getNvidiaConfig().cuda.exist ? "gpu" : "cpu",
+): NitroExecutableOptions => {
+  const nvidiaSettings = getNvidiaConfig();
+  let cudaVisibleDevices = "";
+  let binaryName = "nitro";
+  /**
+   * The binary folder is different for each platform.
+   */
+  if (process.platform === "win32") {
+    /**
+     *  For Windows: win-cpu, win-cuda-11-7, win-cuda-12-0
+     */
+    if (runMode === "cpu") {
+      binaryFolder = path.join(binaryFolder, "win-cpu");
+    } else {
+      if (nvidiaSettings["cuda"].version === "12") {
+        binaryFolder = path.join(binaryFolder, "win-cuda-12-0");
+      } else {
+        binaryFolder = path.join(binaryFolder, "win-cuda-11-7");
+      }
+      cudaVisibleDevices = nvidiaSettings["gpu_highest_vram"];
+    }
+    binaryName = "nitro.exe";
+  } else if (process.platform === "darwin") {
+    /**
+     *  For MacOS: mac-arm64 (Silicon), mac-x64 (InteL)
+     */
+    if (process.arch === "arm64") {
+      binaryFolder = path.join(binaryFolder, "mac-arm64");
+    } else {
+      binaryFolder = path.join(binaryFolder, "mac-x64");
+    }
+  } else {
+    if (runMode === "cpu") {
+      binaryFolder = path.join(binaryFolder, "linux-cpu");
+    } else {
+      if (nvidiaSettings["cuda"].version === "12") {
+        binaryFolder = path.join(binaryFolder, "linux-cuda-12-0");
+      } else {
+        binaryFolder = path.join(binaryFolder, "linux-cuda-11-7");
+      }
+      cudaVisibleDevices = nvidiaSettings["gpu_highest_vram"];
+    }
+  }
+  return {
+    executablePath: path.join(binaryFolder, binaryName),
+    cudaVisibleDevices,
+  };
+};

nitro-node/src/index.ts

@@ -0,0 +1,4 @@
+export * from "./types";
+export * from "./nitro";
+export { getNvidiaConfig, setNvidiaConfig } from "./nvidia";
+export { setLogger } from "./logger";

nitro-node/src/logger.ts

@@ -0,0 +1,14 @@
+import os from "node:os";
+import { NitroLogger } from "./types";
+
+// The logger to use, default to stdout
+export let log: NitroLogger = (message, ..._) =>
+  process.stdout.write(message + os.EOL);
+
+/**
+ * Set logger before running nitro
+ * @param {NitroLogger} logger The logger to use
+ */
+export async function setLogger(logger: NitroLogger): Promise<void> {
+  log = logger;
+}