Add aarch64 support for tfjs-node

tfjs-node-gpu/yarn.lock

@@ -1283,6 +1283,14 @@ minipass@^2.2.1, minipass@^2.3.5:
     safe-buffer "^5.1.2"
     yallist "^3.0.0"
 
+minipass@^2.8.6:
+  version "2.9.0"
+  resolved "https://registry.yarnpkg.com/minipass/-/minipass-2.9.0.tgz#e713762e7d3e32fed803115cf93e04bca9fcc9a6"
+  integrity sha512-wxfUjg9WebH+CUDX/CdbRlh5SmfZiy/hpkxaRI16Y9W56Pa75sWgd/rvFilSgrauD9NyFymP/+JFV3KwzIsJeg==
+  dependencies:
+    safe-buffer "^5.1.2"
+    yallist "^3.0.0"
+
 minizlib@^1.2.1:
   version "1.2.1"
   resolved "https://registry.yarnpkg.com/minizlib/-/minizlib-1.2.1.tgz#dd27ea6136243c7c880684e8672bb3a45fd9b614"
@@ -1348,10 +1356,10 @@ node-gyp@~5.0.3:
     tar "^4.4.8"
     which "1"
 
-node-pre-gyp@0.13.0:
-  version "0.13.0"
-  resolved "https://registry.yarnpkg.com/node-pre-gyp/-/node-pre-gyp-0.13.0.tgz#df9ab7b68dd6498137717838e4f92a33fc9daa42"
-  integrity sha512-Md1D3xnEne8b/HGVQkZZwV27WUi1ZRuZBij24TNaZwUPU3ZAFtvT6xxJGaUVillfmMKnn5oD1HoGsp2Ftik7SQ==
+node-pre-gyp@0.14.0:
+  version "0.14.0"
+  resolved "https://registry.yarnpkg.com/node-pre-gyp/-/node-pre-gyp-0.14.0.tgz#9a0596533b877289bcad4e143982ca3d904ddc83"
+  integrity sha512-+CvDC7ZttU/sSt9rFjix/P05iS43qHCOOGzcr3Ry99bXG7VX953+vFyEuph/tfqoYu8dttBkE86JSKBO2OzcxA==
   dependencies:
     detect-libc "^1.0.2"
     mkdirp "^0.5.1"
@@ -1362,7 +1370,7 @@ node-pre-gyp@0.13.0:
     rc "^1.2.7"
     rimraf "^2.6.1"
     semver "^5.3.0"
-    tar "^4"
+    tar "^4.4.2"
 
 "nopt@2 || 3":
   version "3.0.6"
@@ -2054,7 +2062,20 @@ supports-color@^6.1.0:
   dependencies:
     has-flag "^3.0.0"
 
-tar@^4, tar@^4.4.6, tar@^4.4.8:
+tar@^4.4.2:
+  version "4.4.13"
+  resolved "https://registry.yarnpkg.com/tar/-/tar-4.4.13.tgz#43b364bc52888d555298637b10d60790254ab525"
+  integrity sha512-w2VwSrBoHa5BsSyH+KxEqeQBAllHhccyMFVHtGtdMpF4W7IRWfZjFiQceJPChOeTsSDVUpER2T8FA93pr0L+QA==
+  dependencies:
+    chownr "^1.1.1"
+    fs-minipass "^1.2.5"
+    minipass "^2.8.6"
+    minizlib "^1.2.1"
+    mkdirp "^0.5.0"
+    safe-buffer "^5.1.2"
+    yallist "^3.0.3"
+
+tar@^4.4.6, tar@^4.4.8:
   version "4.4.10"
   resolved "https://registry.yarnpkg.com/tar/-/tar-4.4.10.tgz#946b2810b9a5e0b26140cf78bea6b0b0d689eba1"
   integrity sha512-g2SVs5QIxvo6OLp0GudTqEf05maawKUxXru104iaayWA09551tFCTI8f1Asb4lPfkBr91k07iL4c11XO3/b0tA==

tfjs-node/README.md

@@ -116,3 +116,18 @@ cp bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz ~/myproject/node_
 cd path-to-my-project/node_modules/@tensorflow/tfjs-node/deps
 tar -xf libtensorflow.tar.gz
 ```
+
+If you want to publish an addon library with your own libtensorflow binary, you can host the custom libtensorflow binary and optional pre-compiled node addon module on the cloud service you choose, and add a `custom-binary.json` file in `scripts` folder with the following information:
+
+```js
+{
+  "tf-lib": "url-to-download-customized-binary",
+  "addon": {
+    "host": "host-of-pre-compiled-addon",
+    "remote_path": "remote-path-of-pre-compiled-addon",
+    "package_name": "file-name-of-pre-compile-addon"
+  }
+}
+```
+
+The installation scripts will automatically catch this file and use the custom libtensorflow binary and addon. If `addon` is not provided, the installation script will compile addon from source.

tfjs-node/scripts/custom-binary.sample.json

@@ -0,0 +1,8 @@
+{
+  "tf-lib": "url-to-download-customized-binary",
+  "addon": {
+    "host": "host-of-pre-compiled-addon",
+    "remote_path": "remote-path-of-pre-compiled-addon",
+    "package_name": "file-name-of-pre-compile-addon"
+  }
+}

tfjs-node/scripts/deps-constants.js

@@ -15,14 +15,37 @@
  * =============================================================================
  */
 const os = require('os');
-const path = require('path');
+const fs = require('fs');
+const join = require('path').join;
 const module_path_napi = require('../package.json').binary.module_path;
 const modulePath =
     module_path_napi.replace('{napi_build_version}', process.versions.napi);
 
 /** Version of the libtensorflow shared library to depend on. */
 const LIBTENSORFLOW_VERSION = '1.15.0';
 
+/** Map the os.arch() to arch string in a file name */
+const ARCH_MAPPING = {
+  'x64': 'x86_64'
+};
+/** Map the os.platform() to the platform value in a file name */
+const PLATFORM_MAPPING = {
+  'darwin': 'darwin',
+  'linux': 'linux',
+  'win32': 'windows'
+};
+/** The extension of a compressed file */
+const PLATFORM_EXTENSION = os.platform() === 'win32' ? 'zip' : 'tar.gz';
+/**
+ * Current supported type, platform and architecture combinations
+ * `tf-lib` represents tensorflow shared libraries and `binding` represents
+ * node binding.
+ */
+const ALL_SUPPORTED_COMBINATION = [
+  'cpu-darwin-x86_64', 'gpu-linux-x86_64', 'cpu-linux-x86_64',
+  'cpu-windows-x86_64', 'gpu-windows-x86_64'
+];
+
 /** Get the MAJOR.MINOR-only version of libtensorflow. */
 function getLibTensorFlowMajorDotMinorVersion() {
   const items = LIBTENSORFLOW_VERSION.split('.');
@@ -62,12 +85,20 @@ if (os.platform() === 'win32') {
   throw Exception('Unsupported platform: ' + os.platform());
 }
 
-const depsPath = path.join(__dirname, '..', 'deps');
-const depsLibPath = path.join(depsPath, 'lib');
+const depsPath = join(__dirname, '..', 'deps');
+const depsLibPath = join(depsPath, 'lib');
 
-const depsLibTensorFlowPath = path.join(depsLibPath, depsLibTensorFlowName);
+const depsLibTensorFlowPath = join(depsLibPath, depsLibTensorFlowName);
 const depsLibTensorFlowFrameworkPath =
-    path.join(depsLibPath, depsLibTensorFlowFrameworkName);
+    join(depsLibPath, depsLibTensorFlowFrameworkName);
+
+// Get information for custom binary
+const CUSTOM_BINARY_FILENAME = 'custom-binary.json';
+function loadCustomBinary() {
+  const cfg = join(__dirname, CUSTOM_BINARY_FILENAME);
+  return fs.existsSync(cfg) ? require(cfg) : {};
+}
+const customBinaries = loadCustomBinary();
 
 module.exports = {
   depsPath,
@@ -80,5 +111,11 @@ module.exports = {
   destLibTensorFlowName,
   getLibTensorFlowMajorDotMinorVersion,
   modulePath,
-  LIBTENSORFLOW_VERSION
+  LIBTENSORFLOW_VERSION,
+  ARCH_MAPPING,
+  PLATFORM_MAPPING,
+  PLATFORM_EXTENSION,
+  ALL_SUPPORTED_COMBINATION,
+  customTFLibUri: customBinaries['tf-lib'],
+  customAddon: customBinaries['addon']
 };

tfjs-node/scripts/get-addon-name.js

@@ -20,29 +20,11 @@ const version = require('../package.json').version;
 
 const platform = os.platform();
 
-const CPU_DARWIN = `CPU-darwin-${version}.tar.gz`;
-const CPU_LINUX = `CPU-linux-${version}.tar.gz`;
-const GPU_LINUX = `GPU-linux-${version}.tar.gz`;
-const CPU_WINDOWS = `CPU-windows-${version}.zip`;
-const GPU_WINDOWS = `GPU-windows-${version}.zip`;
+const {PLATFORM_MAPPING, PLATFORM_EXTENSION} = require('./deps-constants.js');
 
-let addonName;
-
-if (name.includes('gpu')) {
-  if (platform === 'linux') {
-    addonName = GPU_LINUX;
-  } else if (platform === 'win32') {
-    addonName = GPU_WINDOWS;
-  }
-} else {
-  if (platform === 'linux') {
-    addonName = CPU_LINUX;
-  } else if (platform === 'darwin') {
-    addonName = CPU_DARWIN;
-  } else if (platform === 'win32') {
-    addonName = CPU_WINDOWS;
-  }
-}
+const type = name.includes('gpu') ? 'GPU' : 'CPU';
+const addonName = `${type}-${PLATFORM_MAPPING[platform]}-` +
+    `${version}.${PLATFORM_EXTENSION}`;
 
 // Print out the addon tarball name so that it can be used in bash script when
 // uploading the tarball to GCP bucket.

tfjs-node/scripts/install.js

@@ -25,9 +25,14 @@ const {
   depsPath,
   depsLibPath,
   depsLibTensorFlowPath,
-  getLibTensorFlowMajorDotMinorVersion,
   LIBTENSORFLOW_VERSION,
-  modulePath
+  PLATFORM_MAPPING,
+  ARCH_MAPPING,
+  PLATFORM_EXTENSION,
+  ALL_SUPPORTED_COMBINATION,
+  modulePath,
+  customTFLibUri,
+  customAddon
 } = require('./deps-constants.js');
 const resources = require('./resources');
 const {addonName} = require('./get-addon-name.js');
@@ -39,31 +44,36 @@ const rimrafPromise = util.promisify(rimraf);
 
 const BASE_URI =
     'https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-';
-const CPU_DARWIN = `cpu-darwin-x86_64-${LIBTENSORFLOW_VERSION}.tar.gz`;
-const CPU_LINUX = `cpu-linux-x86_64-${LIBTENSORFLOW_VERSION}.tar.gz`;
-const GPU_LINUX = `gpu-linux-x86_64-${LIBTENSORFLOW_VERSION}.tar.gz`;
-const CPU_WINDOWS = `cpu-windows-x86_64-${LIBTENSORFLOW_VERSION}.zip`;
-const GPU_WINDOWS = `gpu-windows-x86_64-${LIBTENSORFLOW_VERSION}.zip`;
 
 const platform = os.platform();
+// Use windows path
+if (platform === 'win32') {
+  path = path.win32;
+}
 let libType = process.argv[2] === undefined ? 'cpu' : process.argv[2];
+let system = `${libType}-${PLATFORM_MAPPING[platform]}-` +
+    `${ARCH_MAPPING[os.arch()]}`;
 let forceDownload = process.argv[3] === undefined ? undefined : process.argv[3];
 
 let packageJsonFile;
 
-async function setPackageJsonFile() {
+function setPackageJsonFile() {
   packageJsonFile =
       JSON.parse(fs.readFileSync(`${__dirname}/../package.json`).toString());
 }
 
-async function updateAddonName() {
-  packageJsonFile['binary']['package_name'] = addonName;
+function updateAddonName() {
+  if (customAddon !== undefined) {
+    Object.assign(packageJsonFile['binary'], customAddon);
+  } else {
+    packageJsonFile['binary']['package_name'] = addonName;
+  }
   const stringFile = JSON.stringify(packageJsonFile, null, 2);
   fs.writeFileSync((`${__dirname}/../package.json`), stringFile);
 }
 
-async function revertAddonName() {
-  delete packageJsonFile['binary']['package_name'];
+function revertAddonName(orig) {
+  packageJsonFile['binary'] = orig;
   const stringFile = JSON.stringify(packageJsonFile, null, 2).concat('\n');
   fs.writeFileSync((`${__dirname}/../package.json`), stringFile);
 }
@@ -72,33 +82,20 @@ async function revertAddonName() {
  * Returns the libtensorflow hosted path of the current platform.
  */
 function getPlatformLibtensorflowUri() {
-  let targetUri = BASE_URI;
-  if (platform === 'linux') {
-    if (os.arch() === 'arm') {
-      // TODO(kreeger): Handle arm64 as well:
-      targetUri =
-          'https://storage.googleapis.com/tf-builds/libtensorflow_r1_14_linux_arm.tar.gz';
-    } else {
-      if (libType === 'gpu') {
-        targetUri += GPU_LINUX;
-      } else {
-        targetUri += CPU_LINUX;
-      }
-    }
-  } else if (platform === 'darwin') {
-    targetUri += CPU_DARWIN;
-  } else if (platform === 'win32') {
-    // Use windows path
-    path = path.win32;
-    if (libType === 'gpu') {
-      targetUri += GPU_WINDOWS;
-    } else {
-      targetUri += CPU_WINDOWS;
-    }
-  } else {
-    throw new Error(`Unsupported platform: ${platform}`);
+  // Exception for mac+gpu user
+  if (platform === 'darwin') {
+    system = `cpu-${PLATFORM_MAPPING[platform]}-${ARCH_MAPPING[os.arch()]}`;
+  }
+
+  if (customTFLibUri !== undefined) {
+    return customTFLibUri;
+  }
+
+  if (ALL_SUPPORTED_COMBINATION.indexOf(system) === -1) {
+    throw new Error(`Unsupported system: ${libType}-${platform}-${os.arch()}`);
   }
-  return targetUri;
+
+  return `${BASE_URI}${system}-${LIBTENSORFLOW_VERSION}.${PLATFORM_EXTENSION}`;
 }
 
 /**
@@ -156,27 +153,33 @@ async function downloadLibtensorflow(callback) {
  * Calls node-gyp for Node.js Tensorflow binding after lib is downloaded.
  */
 async function build() {
+  // Load package.json file
+  setPackageJsonFile();
+  // Update addon name in package.json file
+  const origBinary = JSON.parse(JSON.stringify(packageJsonFile['binary']));
+  updateAddonName();
   console.error('* Building TensorFlow Node.js bindings');
-  cp.exec('node-pre-gyp install --fallback-to-build', (err) => {
+  let buildOption = '--fallback-to-build';
+  if (customTFLibUri !== undefined && customAddon === undefined) {
+    // Has custom tensorflow shared libs but no addon. Then build it from source
+    buildOption = '--build-from-source';
+  }
+  cp.exec(`node-pre-gyp install ${buildOption}`, (err) => {
     if (err) {
       console.log('node-pre-gyp install failed with error: ' + err);
     }
     if (platform === 'win32') {
       // Move libtensorflow to module path, where tfjs_binding.node locates.
       cp.exec('node scripts/deps-stage.js symlink ' + modulePath);
     }
-    revertAddonName();
+    revertAddonName(origBinary);
   });
 }
 
 /**
  * Ensures libtensorflow requirements are met for building the binding.
  */
 async function run() {
-  // Load package.json file
-  setPackageJsonFile();
-  // Update addon name in package.json file
-  await updateAddonName();
   // First check if deps library exists:
   if (forceDownload !== 'download' && await exists(depsLibTensorFlowPath)) {
     // Library has already been downloaded, then compile and simlink: