feat: add tvm wasm export

.github/workflows/pipeline.yml

@@ -30,3 +30,10 @@ jobs:
         ./packages/cli/dist/bin/pipcook daemon start
     - name: pipeline
       run: npm run test:pipeline -- ${{ matrix.pipeline }}
+    - name: wasm export pipeline:
+      run: |
+        git clone https://github.com/emscripten-core/emsdk.git && cd emsdk
+        ./emsdk install latest && source ./emsdk_env.sh
+        ./packages/cli/dist/bin/pipcook daemon stop
+        export WASM=1
+        ./packages/cli/dist/bin/pipcook run ./example/pipelines/databinding-image-classification-mobilenet.json

docs/INSTALL.md

@@ -10,6 +10,8 @@ Before starting the installation, please make sure the following environments ar
 - macOS, Linux
 - Node.js 12
 
+**Note:** To use `wasm` output format, you need to manually install [`emsdk`](https://emscripten.org/docs/introducing_emscripten/index.html) and export the `emcc` and `emsdk` to the environmental variable.
+
 ## Install via NPM
 
 Installing [Pipcook][] via NPM is easy, just run:

packages/boa/lib/index.js

@@ -332,6 +332,10 @@ function _internalWrap(T, src={}) {
       writable: false,
       value: () => T.__hash__(),
     },
+    __getattr__: {
+      configurable: true,
+      value: (k) => wrap(T.__getattr__(k)),
+    }
   });
 
   // Create the proxy object for handlers

packages/daemon/binary/preload.js

@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/* eslint-disable no-unused-vars */
+/**
+ * JS config used by --pre-js in emcc.
+ * Wrap module as a LibraryProvider.
+ */
+
+var __wasmLib = {};
+
+function __wasmLibInstantiateWasm(imports, successCallback) {
+    __wasmLib.imports = imports;
+    __wasmLib.successCallback = successCallback;
+}
+
+function __wasmLibStart(wasmInstance) {
+    __wasmLib.successCallback(wasmInstance);
+}
+
+__wasmLib.start = __wasmLibStart;
+
+var Module = {
+    "instantiateWasm": __wasmLibInstantiateWasm,
+    "wasmLibraryProvider": __wasmLib
+};

packages/daemon/package.json

@@ -4,6 +4,7 @@
   "description": "pipcook daemon",
   "dependencies": {
     "@pipcook/app": "^1.2.0",
+    "@pipcook/boa": "^1.2.0",
     "@pipcook/costa": "^1.2.0",
     "@pipcook/pipcook-core": "^1.2.0",
     "axios": "^0.18.1",
@@ -61,7 +62,8 @@
     "test": "midway-bin test --ts --full-trace",
     "migration": "sequelize-cli db:migrate",
     "migration:undo": "sequelize-cli db:migrate:undo",
-    "benchmark": "node benchmark/bootstrap.js"
+    "benchmark": "node benchmark/bootstrap.js",
+    "postinstall": "python3 -m pip install https://files.pythonhosted.org/packages/67/ff/011f588d54153c5d8ee3841d137acf752933f78eb1e3d1c5ffc6c1cb5a32/pipcook_tvm-0.7.dev1-cp37-cp37m-macosx_10_15_x86_64.whl"
   },
   "ci": {
     "version": "10"

packages/daemon/src/service/pipeline.ts

@@ -23,6 +23,8 @@ import { Tracer, JobStatusChangeEvent } from './trace-manager';
 import { pluginQueue } from '../utils';
 import { PipelineDB } from '../runner/helper';
 
+const boa = require('@pipcook/boa');
+
 interface QueryOptions {
   limit: number;
   offset: number;
@@ -39,7 +41,7 @@ interface GenerateOptions {
   datasetProcess?: PluginPackage;
   pipeline: PipelineModel;
   workingDir: string;
-  template: string;
+  template: 'node' | 'wasm';
 }
 
 interface PluginInfo {
@@ -329,15 +331,15 @@ export class PipelineService {
       job.evaluatePass = result.pass;
       job.endTime = Date.now();
       job.status = PipelineStatus.SUCCESS;
-
+      
       await this.generateOutput(job, {
         modelPath,
         modelPlugin,
         dataProcess,
         datasetProcess,
         pipeline,
         workingDir: runnable.workingDir,
-        template: 'node' // set node by default
+        template: process.env.WASM ? 'wasm' : 'node' // set node by default
       });
 
       await job.save();
@@ -383,6 +385,66 @@ export class PipelineService {
     return path.join(CoreConstants.PIPCOOK_RUN, id, 'output.tar.gz');
   }
 
+  private _generateWASMOutput(dist: string, opts: GenerateOptions, fileQueue: Array<Promise<void | string>>): void {
+    const relay = boa.import("tvm.relay");
+    const emcc = boa.import("tvm.contrib.emcc")
+    const keras = boa.import("tensorflow.keras");
+    const {dict, open} = boa.builtins();
+
+    // download tvm runtime from oss
+    const tvmjsPromise = execAsync(`wget http://ai-sample.oss-cn-hangzhou.aliyuncs.com/tvmjs/dist/tvmjs.bundle.js`, {cwd: dist});
+    fileQueue.push(tvmjsPromise);
+
+    const model = keras.models.load_model(path.join(opts.modelPath, "model.h5"));
+
+    const inputName = 'input_1';
+    const inputShape = model.layers[0].input_shape[0];
+    const shape = [1];
+    shape.push(inputShape[3]);
+    shape.push(inputShape[1]);
+    shape.push(inputShape[2]);
+
+    const [ mod, params ] = relay.frontend.from_keras(model, dict(boa.kwargs({[inputName]: shape})));
+    const [ graph, lib, param ] = relay.build(mod, boa.kwargs({
+      params: params,
+      target: "llvm -mtriple=wasm32--unknown-emcc -system-lib"
+    }));
+
+    lib.save(path.join(dist, "model.bc"));
+
+    const jsonWriter = open(path.join(dist, "modelDesc.json"), "w");
+    jsonWriter.write(graph);
+    const paramWriter = open(path.join(dist,"modelParams.parmas"), "wb");
+    paramWriter.write(relay.save_param_dict(param));
+    emcc.create_tvmjs_wasm(path.join(dist, "model.wasi.js"), path.join(dist, "model.bc"), boa.kwargs({
+      options: ["-O3", "-std=c++14", "-Wno-ignored-attributes", "-s", "ALLOW_MEMORY_GROWTH=1", "-s", "STANDALONE_WASM=1", "-s", "ERROR_ON_UNDEFINED_SYMBOLS=0", "-s", "ASSERTIONS=1", "--no-entry", "--pre-js", "./packages/daemon/binary/preload.js"]
+    }));
+
+    const templateHead = `function EmccWASI() {`;
+    const templateTail = `
+      this.Module = Module;
+      this.start = Module.wasmLibraryProvider.start;
+      this.imports = Module.wasmLibraryProvider.imports;
+      this.wasiImport = this.imports["wasi_snapshot_preview1"];
+    }
+
+    if (typeof module !== "undefined" && module.exports) {
+      module.exports = EmccWASI;
+    }
+    `
+
+    const result = templateHead + open(path.join(dist, "model.wasi.js")).read() + templateTail;
+    const resultWriter = open(path.join(dist, "model.wasi.js"), "w");
+    resultWriter.write(result);
+
+    const jsonPromise = fs.writeJSON(path.join(dist, "modelSpec.json"), {
+      shape,
+      inputName  
+    });
+
+    fileQueue.push(jsonPromise);
+  }
+
   /**
    * Generate the output package for a given job.
    * @param job the job model for output.
@@ -393,18 +455,34 @@ export class PipelineService {
     const dist = path.join(opts.workingDir, 'output');
     await fs.remove(dist);
     await fs.ensureDir(dist);
+
+    const fileQueue: Array<Promise<void | string>> = new Array();
+
+    // Only support tensorflow at this moment.
+    if (opts.template == 'wasm' && opts.modelPlugin.name.includes("tensorflow")) {
+      this._generateWASMOutput(dist, opts, fileQueue);
+    }
+
     await execAsync('npm init -y', { cwd: dist });
 
     // post processing the package.json
     const projPackage = await fs.readJSON(dist + '/package.json');
-    projPackage.dependencies = {
-      [opts.modelPlugin.name]: opts.modelPlugin.version,
-    };
-    projPackage.scripts = {
-      postinstall: 'node boapkg.js'
-    };
-    if (opts.dataProcess) {
-      projPackage.dependencies[opts.dataProcess.name] = opts.dataProcess.version;
+
+    if (opts.template == 'node') {
+      projPackage.dependencies = {
+        [opts.modelPlugin.name]: opts.modelPlugin.version,
+      };
+      projPackage.scripts = {
+        postinstall: 'node boapkg.js'
+      };
+      if (opts.dataProcess) {
+        projPackage.dependencies[opts.dataProcess.name] = opts.dataProcess.version;
+      }
+    } else {
+      projPackage.main = 'index.js',
+      projPackage.dependencies = {
+        'ws': '^7.3.1'
+      }
     }
 
     const jsonWriteOpts = { spaces: 2 } as fs.WriteOptions;
@@ -413,18 +491,21 @@ export class PipelineService {
       output: job.toJSON(),
     };
 
-    await Promise.all([
+    if (opts.template == 'node') {
       // copy base components
-      fs.copy(opts.modelPath, dist + '/model'),
-      fs.copy(path.join(__dirname, `../../templates/${opts.template}/predict.js`), `${dist}/index.js`),
-      fs.copy(path.join(__dirname, '../../templates/boapkg.js'), `${dist}/boapkg.js`),
-      // copy logs
-      fs.copy(opts.workingDir + '/logs', `${dist}/logs`),
-      // write package.json
-      fs.outputJSON(dist + '/package.json', projPackage, jsonWriteOpts),
-      // write metadata.json
-      fs.outputJSON(dist + '/metadata.json', metadata, jsonWriteOpts),
-    ]);
+      fileQueue.push(fs.copy(opts.modelPath, dist + '/model'));
+      fileQueue.push(fs.copy(path.join(__dirname, '../../templates/boapkg.js'), `${dist}/boapkg.js`));
+    }
+
+    fileQueue.push(fs.copy(path.join(__dirname, `../../templates/${opts.template}/predict.js`), `${dist}/index.js`));
+    // copy logs
+    fileQueue.push(fs.copy(opts.workingDir + '/logs', `${dist}/logs`));
+    // write package.json
+    fileQueue.push(fs.outputJSON(dist + '/package.json', projPackage, jsonWriteOpts));
+    // write metadata.json
+    fileQueue.push(fs.outputJSON(dist + '/metadata.json', metadata, jsonWriteOpts));
+
+    await Promise.all(fileQueue);
     console.info(`trained the model to ${dist}`);
 
     // packing the output directory.

packages/daemon/templates/wasm/predict.js

@@ -0,0 +1,41 @@
+const tvmjs = require("./tvmjs.bundle");
+const EmccWASI = require("./model.wasi");
+const fs = require('fs');
+const modelSpec = require("./modelSpec.json");
+
+const loadModel = async () => {
+  const wasmSource = fs.readFileSync('./model.wasi.wasm');
+  const tvm = await tvmjs.instantiate(wasmSource, new EmccWASI());
+
+  const graph = JSON.parse(fs.readFileSync('./modelDesc.json'));
+  const param = new Uint8Array(fs.readFileSync('./modelParams.parmas'));
+
+  const ctx = tvm.cpu(0);
+  const sysLib = tvm.systemLib();
+  model = tvm.createGraphRuntime(JSON.stringify(graph), sysLib, ctx);
+  model.loadParams(param);
+
+  return {model, tvm, ctx};
+}
+
+let model, tvm, ctx;
+
+const predict = async (input) => {
+  if (!model) {
+    const rets = await loadModel();
+    model = rets.model;
+    tvm = rets.tvm;
+    ctx = rets.ctx;
+  }
+
+  const inputData = tvm.empty(modelSpec.shape, "float32", tvm.cpu());
+  const output = model.getOutput(0);
+  inputData.copyFrom(input);
+  model.setInput(modelSpec.inputName, inputData);
+  model.run();
+  await ctx.sync();
+  console.log(output.toArray())
+  return output.toArray();
+}
+
+module.exports = predict;