Refactor & improve stability

CHANGES.md

@@ -1,5 +1,23 @@
 # CHANGELOG
 
+## v1.1.0
+
+- Rename EmbeddingGenerator to EmbeddingEncoder
+- Fixed serverOptions weren't passed through properly in test cases
+- Upgrade to @huggingface/transformers v3.2.4
+- Upgrade onnxruntime-node v1.20.1
+- Avoid including unused models in docker images (smaller image size)
+- Increase probe timeout seconds
+- Use worker pool
+- Process sentence list with separate model runs
+- set default `workerTaskTimeout` to `60` seconds
+- use quantized version (q8) default model
+- set default `limits.memory` to `850M`
+- set default replicas number to `2`
+- Add max_length config to model config (configurable via helm config)
+- set max_length of default model to 1024 due to excessive memory usage when working on text longer than 2048 (the default model supports up to 8192)
+- only use padding for multiple inputs received when encoding the input
+
 ## v1.0.0
 
-- #1: Initial implementation
+- #1: Initial implementation

Dockerfile

@@ -8,12 +8,16 @@ RUN mkdir -p /usr/src/app /etc/config && \
     chmod -R g=u /usr/src/app /etc/config
 
 COPY . /usr/src/app
-# make local cache folder writable by 1000 user
-RUN chown -R 1000 /usr/src/app/component/node_modules/@xenova/transformers/.cache
 # Reinstall onnxruntime-node based on current building platform architecture
 RUN cd /usr/src/app/component/node_modules/onnxruntime-node && npm run postinstall
 # Reinstall sharp based on current building platform architecture
-RUN cd /usr/src/app/component/node_modules/sharp && npm run clean && node install/libvips && node install/dll-copy && node ../prebuild-install/bin.js
+RUN cd /usr/src/app/component/node_modules && rm -Rf @img && cd sharp && npm install
+# remove downloaded model cache
+RUN cd /usr/src/app/component/node_modules/@huggingface/transformers && rm -Rf .cache && mkdir .cache
+# download default model
+RUN cd /usr/src/app/component && npm run download-default-model
+# make local cache folder writable by 1000 user
+RUN chown -R 1000 /usr/src/app/component/node_modules/@huggingface/transformers/.cache
 
 USER 1000
 WORKDIR /usr/src/app/component

README.md

@@ -49,7 +49,7 @@ Kubernetes: `>= 1.21.0`
 | appConfig | object | `{}` | Application configuration of the service. You can supply a list of key-value pairs to be used as the application configuration. Currently, the only supported config field is `modelList`. Via the `modelList` field, you can specify a list of LLM models that the service supports. Although you can specify multiple models, only one model will be used at this moment. Each model item have the following fields:  <ul> <li> `name` (string): The huggingface registered model name. We only support ONNX model at this moment. This field is required. </li> <li> `default` (bool): Optional; Whether this model is the default model. If not specified, the first model in the list will be the default model. Only default model will be loaded. </li> <li> `quantized` (bool): Optional; Whether the quantized version of model will be used. If not specified, the quantized version model will be loaded. </li> <li> `config` (object): Optional; The configuration object that will be passed to the model.  </li> <li> `cache_dir` (string): Optional; The cache directory of the downloaded models. If not specified, the default cache directory will be used. </li> <li> `local_files_only` (bool): Optional; Whether to only load the model from local files. If not specified, the model will be downloaded from the huggingface model hub. </li> <li> `revision` (string) Optional, Default to 'main'; The specific model version to use. It can be a branch name, a tag name, or a commit id.    Since we use a git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any identifier allowed by git. NOTE: This setting is ignored for local requests. </li> <li> `model_file_name` (string) Optional; </li> <li> `extraction_config` (object) Optional; The configuration object that will be passed to the model extraction function for embedding generation. <br/>   <ul>   <li> `pooling`: ('none' or 'mean' or 'cls') Default to 'none'. The pooling method to use. </li>   <li> `normalize`: (bool) Default to true. Whether or not to normalize the embeddings in the last dimension. </li>   <li> `quantize`: (bool) Default to `false`. Whether or not to quantize the embeddings. </li>   <li> `precision`: ("binary" or "ubinary") default to "binary". The precision to use for quantization. Only used when `quantize` is true. </li>   </ul>   </li> </ul> Please note: The released docker image only contains "Alibaba-NLP/gte-base-en-v1.5" model.  If you specify other models, the server will download the model from the huggingface model hub at the startup. You might want to adjust the `startupProbe` settings to accommodate the model downloading time. Depends on the model size, you might also want to adjust the `resources.limits.memory` & `resources.requests.memory`value. |
 | autoscaling.hpa.enabled | bool | `false` |  |
 | autoscaling.hpa.maxReplicas | int | `3` |  |
-| autoscaling.hpa.minReplicas | int | `1` |  |
+| autoscaling.hpa.minReplicas | int | `2` |  |
 | autoscaling.hpa.targetCPU | int | `90` |  |
 | autoscaling.hpa.targetMemory | string | `""` |  |
 | bodyLimit | int | Default to 10485760 (10MB). | Defines the maximum payload, in bytes, that the server is allowed to accept |
@@ -79,9 +79,11 @@ Kubernetes: `>= 1.21.0`
 | livenessProbe.successThreshold | int | `1` |  |
 | livenessProbe.timeoutSeconds | int | `5` |  |
 | logLevel | string | `"warn"` | The log level of the application. one of 'fatal', 'error', 'warn', 'info', 'debug', 'trace';  also 'silent' is supported to disable logging.  Any other value defines a custom level and requires supplying a level value via levelVal. |
+| maxWorkers | int | Default to 1. | The maximum number of workers that run the model to serve the request. |
+| minWorkers | int | Default to 1. | The maximum number of workers that run the model to serve the request. |
 | nameOverride | string | `""` |  |
 | nodeSelector | object | `{}` |  |
-| pluginTimeout | int | Default to 10000 (10 seconds). | The maximum amount of time in milliseconds in which a fastify plugin can load.  If not, ready will complete with an Error with code 'ERR_AVVIO_PLUGIN_TIMEOUT'. |
+| pluginTimeout | int | Default to 180000 (180 seconds). | The maximum amount of time in milliseconds in which a fastify plugin can load.  If not, ready will complete with an Error with code 'ERR_AVVIO_PLUGIN_TIMEOUT'. |
 | podAnnotations | object | `{}` |  |
 | podSecurityContext.runAsNonRoot | bool | `true` |  |
 | podSecurityContext.runAsUser | int | `1000` |  |
@@ -97,10 +99,10 @@ Kubernetes: `>= 1.21.0`
 | readinessProbe.periodSeconds | int | `20` |  |
 | readinessProbe.successThreshold | int | `1` |  |
 | readinessProbe.timeoutSeconds | int | `5` |  |
-| replicas | int | `1` |  |
-| resources.limits.memory | string | `"1100M"` | the memory limit of the container Due to [this issue of ONNX runtime](https://github.com/microsoft/onnxruntime/issues/15080), the peak memory usage of the service is much higher than the model file size.  When change the default model, be sure to test the peak memory usage of the service before setting the memory limit. quantized model will be used by default, the memory limit is set to 1100M to accommodate the default model size. |
+| replicas | int | `2` |  |
+| resources.limits.memory | string | `"850M"` | the memory limit of the container Due to [this issue of ONNX runtime](https://github.com/microsoft/onnxruntime/issues/15080), the peak memory usage of the service is much higher than the model file size.  When change the default model, be sure to test the peak memory usage of the service before setting the memory limit. When test your model memory requirement, please note that the memory usage of the model often goes much higher with long context length. E.g. the default model supports up to 8192 tokens (default max_length set to 1024), but when the content go beyond 512 tokens, the memory usage will be much higher (requires around 2G). |
 | resources.requests.cpu | string | `"100m"` |  |
-| resources.requests.memory | string | `"650M"` | the memory request of the container Once the model is loaded, the memory usage of the service for serving request would be much lower. Set to 650M for default model. |
+| resources.requests.memory | string | `"650M"` | the memory request of the container Once the model is loaded, the memory usage of the service for serving request would be much lower. Set to 850M for default model. |
 | service.annotations | object | `{}` |  |
 | service.httpPortName | string | `"http"` |  |
 | service.labels | object | `{}` |  |
@@ -120,6 +122,7 @@ Kubernetes: `>= 1.21.0`
 | startupProbe.timeoutSeconds | int | `5` |  |
 | tolerations | list | `[]` |  |
 | topologySpreadConstraints | list | `[]` | This is the pod topology spread constraints https://kubernetes.io/docs/concepts/workloads/pods/pod-topology-spread-constraints/ |
+| workerTaskTimeout | int | Default to 60000 (60 seconds). | The maximum time in milliseconds that a worker can run before being killed. |
 
 ### Build & Run for Local Development
 

deploy/magda-embedding-api/Chart.yaml

@@ -1,7 +1,7 @@
 apiVersion: v2
 name: magda-embedding-api
 description: An OpenAI embeddings API compatible microservice for Magda.
-version: "1.0.0"
+version: "1.1.0-alpha.0"
 kubeVersion: ">= 1.21.0"
 home: "https://github.com/magda-io/magda-embedding-api"
 sources: [ "https://github.com/magda-io/magda-embedding-api" ]

deploy/magda-embedding-api/templates/deployment.yaml

@@ -105,6 +105,12 @@ spec:
         - "start"
         - "dist/app.js"
         - "--"
+        - "--minWorkers"
+        - {{ .Values.minWorkers | quote }}
+        - "--maxWorkers"
+        - {{ .Values.maxWorkers | quote }}
+        - "--workerTaskTimeout"
+        - {{ .Values.workerTaskTimeout | quote }}
         - "--appConfigFile"
         - "/etc/config/appConfig.json"
         {{- if .Values.readinessProbe }}

deploy/magda-embedding-api/values.yaml

@@ -14,8 +14,8 @@ logLevel: "warn"
 
 # -- (int) The maximum amount of time in milliseconds in which a fastify plugin can load. 
 # If not, ready will complete with an Error with code 'ERR_AVVIO_PLUGIN_TIMEOUT'.
-# @default -- Default to 10000 (10 seconds).
-pluginTimeout: 
+# @default -- Default to 180000 (180 seconds).
+pluginTimeout: 180000
 
 # -- (int) Defines the maximum payload, in bytes, that the server is allowed to accept
 # @default -- Default to 10485760 (10MB).
@@ -57,6 +57,18 @@ closeGraceDelay: 25000
 # Depends on the model size, you might also want to adjust the `resources.limits.memory` & `resources.requests.memory`value.
 appConfig: {}
 
+# -- (int) The maximum number of workers that run the model to serve the request.
+# @default -- Default to 1.
+maxWorkers: 1
+
+# -- (int) The maximum number of workers that run the model to serve the request.
+# @default -- Default to 1.
+minWorkers: 1
+
+# -- (int) The maximum time in milliseconds that a worker can run before being killed.
+# @default -- Default to 60000 (60 seconds).
+workerTaskTimeout: 60000
+
 # image setting loadding order: (from higher priority to lower priority)
 # - Values.image.x
 # - Values.defaultImage.x
@@ -76,12 +88,12 @@ defaultImage:
 nameOverride: ""
 fullnameOverride: ""
 
-replicas: 1
+replicas: 2
 
 autoscaling:
   hpa:
     enabled: false
-    minReplicas: 1
+    minReplicas: 2
     maxReplicas: 3
     targetCPU: 90
     targetMemory: ""  
@@ -191,11 +203,12 @@ resources:
   requests:
     cpu: "100m"
     # -- (string) the memory request of the container
-    # Once the model is loaded, the memory usage of the service for serving request would be much lower. Set to 650M for default model.
+    # Once the model is loaded, the memory usage of the service for serving request would be much lower. Set to 850M for default model.
     memory: "650M"
   limits:
     # -- (string) the memory limit of the container
     # Due to [this issue of ONNX runtime](https://github.com/microsoft/onnxruntime/issues/15080), the peak memory usage of the service is much higher than the model file size. 
     # When change the default model, be sure to test the peak memory usage of the service before setting the memory limit.
-    # quantized model will be used by default, the memory limit is set to 1100M to accommodate the default model size.
-    memory: "1100M"
+    # When test your model memory requirement, please note that the memory usage of the model often goes much higher with long context length.
+    # E.g. the default model supports up to 8192 tokens (default max_length set to 1024), but when the content go beyond 512 tokens, the memory usage will be much higher (requires around 2G).
+    memory: "850M"

deploy/test-deploy.yaml

@@ -19,7 +19,11 @@ appConfig:
   - name: Xenova/bge-small-en-v1.5
     # set quantized to false to use the non-quantized version of the model
     # by default, the quantized version of the model will be used
-    quantized: true
+    dtype: "q8"
+    # optional set max length of the input text
+    # if not set, the value in model config will be used
+    # if model config does not have max_length, the default value (512) will be used
+    max_length: 512,
     extraction_config:
       pooling: "mean"
       normalize: true

package.json

@@ -1,7 +1,7 @@
 {
   "name": "magda-embedding-api",
   "type": "module",
-  "version": "1.0.0",
+  "version": "1.1.0-alpha.0",
   "description": "An OpenAI embeddings API compatible microservice for Magda.",
   "main": "app.ts",
   "directories": {
@@ -22,7 +22,8 @@
     "prebuild": "rimraf dist",
     "build": "tsc",
     "watch": "tsc -w",
-    "start": "npm run build && fastify start -l info dist/app.js",
+    "start": "npm run build && fastify start -T 60000 -l info dist/app.js -- --workerTaskTimeout 10000",
+    "download-default-model": "node ./dist/loadDefaultModel.js",
     "dev": "npm run build && concurrently -k -p \"[{name}]\" -n \"TypeScript,App\" -c \"yellow.bold,cyan.bold\" \"npm:watch\" \"npm:dev:start\"",
     "dev:start": "fastify start --ignore-watch=.ts$ -w -l info -P dist/app.js",
     "docker-build-local": "create-docker-context-for-node-component --build --push --tag auto --repository example.com",
@@ -41,20 +42,22 @@
     "@fastify/autoload": "^5.10.0",
     "@fastify/sensible": "^5.6.0",
     "@fastify/type-provider-typebox": "^4.0.0",
+    "@huggingface/transformers": "^3.2.4",
     "@sinclair/typebox": "^0.32.34",
-    "@xenova/transformers": "^2.17.2",
     "fastify": "^4.28.1",
     "fastify-cli": "^6.2.1",
     "fastify-plugin": "^4.5.1",
     "fs-extra": "^11.2.0",
-    "onnxruntime-node": "^1.14.0"
+    "onnxruntime-node": "^1.20.1",
+    "workerpool": "^9.2.0"
   },
   "devDependencies": {
-    "@langchain/openai": "^0.2.1",
+    "@langchain/openai": "^0.2.8",
     "@magda/ci-utils": "^1.0.5",
     "@magda/docker-utils": "^4.2.1",
     "@types/fs-extra": "^11.0.4",
     "@types/node": "^18.19.3",
+    "@types/workerpool": "^6.4.7",
     "concurrently": "^8.2.2",
     "eslint": "^9.6.0",
     "fastify-tsconfig": "^2.0.0",
@@ -75,7 +78,7 @@
     "exclude": [
       "test/helper.ts"
     ],
-    "timeout": 120,
+    "timeout": 12000,
     "jobs": 1
   },
   "config": {

src/app.ts

@@ -9,12 +9,16 @@ const __dirname = path.dirname(__filename);
 
 export type AppOptions = {
     appConfigFile: string;
+    maxWorkers: number;
+    minWorkers: number;
     // Place your custom options for app below here.
 } & Partial<AutoloadPluginOptions>;
 
 // Pass --options via CLI arguments in command to enable these options.
 const options: AppOptions = {
-    appConfigFile: ""
+    appConfigFile: "",
+    maxWorkers: 1,
+    minWorkers: 1
 };
 
 const app: FastifyPluginAsync<AppOptions> = async (