Migrate tfjs-tflite package to bazel (tensorflow#5358)

* save * save * save * save * save * save * save * save * fix * fix
mattsoulanille · Aug 6, 2021 · 19227b2 · 19227b2
1 parent b2e256e
commit 19227b2
Show file tree

Hide file tree

Showing 23 changed files with 1,252 additions and 3,068 deletions.
diff --git a/e2e/scripts/publish-tfjs-ci.sh b/e2e/scripts/publish-tfjs-ci.sh
@@ -58,7 +58,8 @@ PACKAGES=("tfjs-core" "tfjs-backend-cpu" "tfjs-backend-webgl" \
 "tfjs-node" "tfjs-node-gpu")
 
 # Packages that build with Bazel
-BAZEL_PACKAGES=("tfjs-core" "tfjs-backend-cpu" "tfjs-backend-webgl")
+BAZEL_PACKAGES=("tfjs-core" "tfjs-backend-cpu" "tfjs-backend-webgl" \
+"tfjs-tflite")
 
 for package in "${PACKAGES[@]}"
 do

diff --git a/link-package/package.json b/link-package/package.json
@@ -6,16 +6,18 @@
   "license": "Apache 2.0",
   "private": true,
   "scripts": {
-    "build": "yarn build-backend-cpu && yarn build-backend-webgl && yarn build-core && yarn reinstall",
+    "build": "yarn build-backend-cpu && yarn build-backend-webgo && yarn build-core && yarn build-tflite && yarn reinstall",
     "build-backend-cpu": "cd ../tfjs-backend-cpu && yarn && yarn build",
     "build-backend-webgl": "cd ../tfjs-backend-webgl && yarn && yarn build",
     "build-core": "cd ../tfjs-core && yarn && yarn build",
-    "reinstall": "yarn && yarn cache clean @tensorflow/tfjs-core && yarn cache clean @tensorflow/tfjs-backend-cpu && rimraf node_modules && yarn"
+    "build-tflite": "cd ../tfjs-tflite && yarn && yarn build",
+    "reinstall": "yarn && yarn cache clean @tensorflow/tfjs-core && yarn cache clean @tensorflow/tfjs-backend-cpu && yarn cache clean @tensorflow/tfjs-tflite && rimraf node_modules && yarn"
   },
   "devDependencies": {
     "@tensorflow/tfjs-backend-cpu": "file:../dist/bin/tfjs-backend-cpu/tfjs-backend-cpu_pkg",
     "@tensorflow/tfjs-backend-webgl": "file:../dist/bin/tfjs-backend-webgl/tfjs-backend-webgl_pkg",
     "@tensorflow/tfjs-core": "file:../dist/bin/tfjs-core/tfjs-core_pkg",
+    "@tensorflow/tfjs-tflite": "file:../dist/bin/tfjs-tflite/tfjs-tflite_pkg",
     "rimraf": "^3.0.2"
   }
 }
diff --git a/link-package/yarn.lock b/link-package/yarn.lock
@@ -33,6 +33,9 @@
     node-fetch "~2.6.1"
     seedrandom "2.4.3"
 
+"@tensorflow/tfjs-tflite@file:../dist/bin/tfjs-tflite/tfjs-tflite_pkg":
+  version "0.0.0"
+
 "@types/long@^4.0.1":
   version "4.0.1"
   resolved "https://registry.yarnpkg.com/@types/long/-/long-4.0.1.tgz#459c65fa1867dafe6a8f322c4c51695663cc55e9"

diff --git a/scripts/cloudbuild_tfjs_core_expected.yml b/scripts/cloudbuild_tfjs_core_expected.yml
@@ -238,64 +238,6 @@ steps:
       - build-tfjs-backend-wasm
       - yarn-common
       - yarn-link-package
-  - name: 'node:10'
-    dir: tfjs-tflite
-    entrypoint: yarn
-    id: yarn-tfjs-tflite
-    args:
-      - install
-    waitFor:
-      - yarn-common
-      - yarn-link-package
-  - name: gcr.io/google.com/cloudsdktool/cloud-sdk
-    dir: tfjs-tflite
-    id: download-tflite-web-api-tfjs-tflite
-    entrypoint: bash
-    args:
-      - scripts/download-tflite-web-api.sh
-    waitFor:
-      - yarn-tfjs-tflite
-      - yarn-common
-      - yarn-link-package
-  - name: 'node:10'
-    dir: tfjs-tflite
-    id: build-tfjs-tflite
-    entrypoint: yarn
-    args:
-      - build-ci
-    waitFor:
-      - yarn-tfjs-tflite
-      - download-tflite-web-api-tfjs-tflite
-      - yarn-common
-      - yarn-link-package
-  - name: 'node:10'
-    dir: tfjs-tflite
-    entrypoint: yarn
-    id: lint-tfjs-tflite
-    args:
-      - lint
-    waitFor:
-      - yarn-tfjs-tflite
-      - download-tflite-web-api-tfjs-tflite
-      - yarn-common
-      - yarn-link-package
-  - name: 'node:10'
-    dir: tfjs-tflite
-    entrypoint: yarn
-    id: test-tfjs-tflite-tfjs-tflite
-    args:
-      - test-ci
-    waitFor:
-      - yarn-tfjs-tflite
-      - lint-tfjs-tflite
-      - download-tflite-web-api-tfjs-tflite
-      - yarn-common
-      - yarn-link-package
-    env:
-      - BROWSERSTACK_USERNAME=deeplearnjs1
-      - NIGHTLY=$_NIGHTLY
-    secretEnv:
-      - BROWSERSTACK_KEY
   - name: 'node:10'
     dir: tfjs-data
     entrypoint: yarn

diff --git a/scripts/package_dependencies.json b/scripts/package_dependencies.json
@@ -11,7 +11,6 @@
   "tfjs-backend-webgpu": ["tfjs-core", "tfjs-converter"],
   "tfjs-inference": [],
   "tfjs-react-native": [],
-  "tfjs-tflite": ["tfjs-core"],
   "tfjs-vis": [],
   "e2e": ["tfjs", "tfjs-converter", "tfjs-core", "tfjs-data", "tfjs-layers", "tfjs-node", "tfjs-backend-wasm"]
 }
diff --git a/scripts/publish-npm.ts b/scripts/publish-npm.ts
@@ -28,8 +28,8 @@ import * as shell from 'shelljs';
 import {RELEASE_UNITS, question, $, printReleaseUnit, printPhase, getReleaseBranch, checkoutReleaseBranch} from './release-util';
 
 const TMP_DIR = '/tmp/tfjs-publish';
-const BAZEL_PACKAGES =
-    new Set(['tfjs-core', 'tfjs-backend-cpu', 'tfjs-backend-webgl']);
+const BAZEL_PACKAGES = new Set(
+    ['tfjs-core', 'tfjs-backend-cpu', 'tfjs-backend-webgl', 'tfjs-tflite']);
 
 const parser = new argparse.ArgumentParser();
 parser.addArgument('--git-protocol', {
@@ -89,10 +89,6 @@ async function main() {
     if (pkg === 'tfjs-node-gpu') {
       $('yarn prep-gpu');
     }
-    // tfjs-tflite needs to download the tflite web api from google storage.
-    if (pkg === 'tfjs-tflite') {
-      $('yarn prep');
-    }
 
     // Yarn above the other checks to make sure yarn doesn't change the lock
     // file.

diff --git a/tfjs-backend-cpu/package.json b/tfjs-backend-cpu/package.json
@@ -49,7 +49,7 @@
   "scripts": {
     "build-ci": "yarn build",
     "build": "bazel build :tfjs-backend-cpu_pkg",
-    "bundle": "bazel build :tfjs-core_pkg",
+    "bundle": "bazel build :tfjs-backend-cpu_pkg",
     "bundle-ci": "yarn bundle",
     "build-core": "cd ../tfjs-core && yarn && yarn build",
     "build-core-ci": "cd ../tfjs-core && yarn && yarn build-ci",

diff --git a/tfjs-tflite/BUILD.bazel b/tfjs-tflite/BUILD.bazel
@@ -0,0 +1,104 @@
+# Copyright 2021 Google LLC. All Rights Reserved.
+#
+# Licensed 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.
+# =============================================================================
+
+load("@build_bazel_rules_nodejs//:index.bzl", "pkg_npm")
+load("//tools:copy_to_dist.bzl", "copy_to_dist")
+load("//tools:tfjs_bundle.bzl", "tfjs_bundle")
+load("//tools:tfjs_web_test.bzl", "tfjs_web_test")
+load("//tools:ts_library_outputs.bzl", "ts_library_outputs")
+
+package(default_visibility = ["//visibility:public"])
+
+# Run rollup to bundle the compiled files for different execution environments.
+tfjs_bundle(
+    name = "tf-tflite",
+    entry_point = "//tfjs-tflite/src:index.ts",
+    external = [
+        "crypto",
+        "@tensorflow/tfjs-core",
+        "seedrandom",
+    ],
+    globals = {
+        "@tensorflow/tfjs-core": "tf",
+        "seedrandom": "seedrandom",
+    },
+    umd_name = "tflite",
+    deps = [
+        "//tfjs-tflite/src:tfjs-tflite_lib",
+        "//tfjs-tflite/src:tfjs-tflite_src_lib",
+        "//tfjs-tflite/src:tflite_web_api_client_js",
+        "//tfjs-tflite/wasm:wasm_files",
+    ],
+)
+
+# Gather .js, .mjs, and .d.ts files as output files.
+ts_library_outputs(
+    name = "tfjs-tflite_outputs",
+    srcs = [
+        "//tfjs-tflite/src:tfjs-tflite_lib",
+        "//tfjs-tflite/src:tfjs-tflite_src_lib",
+    ],
+)
+
+# Copy ouput files to dist/.
+copy_to_dist(
+    name = "copy_src_to_dist",
+    srcs = [
+        ":tfjs-tflite_outputs",
+    ],
+    root = "src",
+)
+
+# Copy bundled files to dist/.
+copy_to_dist(
+    name = "copy_bundles",
+    srcs = [
+        ":tf-tflite",
+        ":tf-tflite.es2017",
+        ":tf-tflite.es2017.min",
+        ":tf-tflite.fesm",
+        ":tf-tflite.fesm.min",
+        ":tf-tflite.min",
+    ],
+)
+
+# Copy wasm related files to dist/.
+copy_to_dist(
+    name = "copy_wasm_files",
+    srcs = [
+        "//tfjs-tflite/wasm:wasm_files",
+    ],
+    root = "wasm",
+)
+
+# Package for npm.
+pkg_npm(
+    name = "tfjs-tflite_pkg",
+    srcs = ["package.json"],
+    tags = ["ci"],
+    deps = [
+        ":copy_bundles",
+        ":copy_src_to_dist",
+        ":copy_wasm_files",
+    ],
+)
+
+# Tests.
+tfjs_web_test(
+    name = "tfjs-tflite_test",
+    srcs = [
+        "//tfjs-tflite/src:tfjs-tflite_test_bundle",
+    ],
+)
diff --git a/tfjs-tflite/cloudbuild.yml b/tfjs-tflite/cloudbuild.yml
diff --git a/tfjs-tflite/demo/README.md b/tfjs-tflite/demo/README.md
@@ -12,35 +12,15 @@ the "WebAssembly SIMD support" and "WebAssembly threads support" in
 
 ## Run the demo locally
 
-cd into the `tfjs-tflite` folder:
+Build the dependencies.
 
 ```sh
-$ cd tfjs/tfjs-tflite
-```
-
-Make sure `gsutil` is installed. It will be used by the script below.
-Follow the installation instructions [here][gsutil]. You can skip the
-"Setting Up Credentials to Access Protected Data" section.
-
-Build the `tfjs-tflite` package locally.
-```sh
+$ yarn build-deps
 $ yarn
-# This script will download the TFLite web API WASM module files and JS client
-# to deps/. This script will use `gsutil`.
-$ ./scripts/download-tflite-web-api.sh
-# This will bundle and copy everything needed to tfjs-tflite/dist/ which will
-# be used by the local demo.
-$ yarn build-npm
 ```
 
-cd into the `demo` folder
+Run the demo locally
 ```sh
-$ cd demo
-```
-
-Build and run the demo locally
-```sh
-$ yarn
 $ yarn watch
 ```
 

diff --git a/tfjs-tflite/demo/package.json b/tfjs-tflite/demo/package.json
@@ -4,9 +4,9 @@
   "main": "script.js",
   "license": "Apache-2.0",
   "dependencies": {
-    "@tensorflow/tfjs-backend-cpu": "link:../../tfjs-backend-cpu",
-    "@tensorflow/tfjs-core": "link:../../tfjs-core",
-    "@tensorflow/tfjs-tflite": "link:../dist"
+    "@tensorflow/tfjs-backend-cpu": "link:../../link-package/node_modules/@tensorflow/tfjs-backend-cpu",
+    "@tensorflow/tfjs-core": "link:../../link-package/node_modules/@tensorflow/tfjs-core",
+    "@tensorflow/tfjs-tflite": "link:../../link-package/node_modules/@tensorflow/tfjs-tflite"
   },
   "devDependencies": {
     "@babel/core": "7.13.14",
@@ -20,10 +20,11 @@
   },
   "staticFiles": {
     "staticPath": [
-      "../dist"
+      "../../link-package/node_modules/@tensorflow/tfjs-tflite/dist"
     ]
   },
   "scripts": {
+    "build-deps": "cd ../../link-package && yarn build",
     "watch": "cross-env NODE_ENV=development parcel src/index.html --no-hmr --open",
     "build": "NODE_ENV=production parcel build src/index.html --no-minify --public-url ./"
   }