[tfjs-node] Support external binaries

Modify the install.js to support external binaries while installing
npm package. It reads the `custom-binary.json` configuration file to
get URLs for the following resources:
- 'tf-lib': libtensorflow and libtensorflow_framework shared libraries.
  Its value is a URL that points to a `tar.gz` or `zip` and
  contains the shared libaries and header files.
- 'addon': pre-compiled node binding. Its value is a object contains
  three values: `host`, `remote_path` and `package_name`. These three
  values are concatenated into a URL that points to the
  pre-compiled node binding.

If you only provide `tf-lib` and no `addon`. It compiles the node
binding while installing the npm package.

Signed-off-by: Yihong Wang <yh.wang@ibm.com>

Author: yhwang

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

@@ -15,39 +15,61 @@
  * =============================================================================
  */
 const os = require('os');
+const fs = require('fs');
+const join = require('path').join;
 const name = require('../package.json').name;
 const version = require('../package.json').version;
 
 const platform = os.platform();
+const CUSTOM_BINARY_FILENAME = 'custom-binary.json';
+const customBinaries = loadCustomBinary();
 
-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`;
-
-let addonName;
+/** 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'
+];
 
-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.
 console.log(addonName);
 
+function loadCustomBinary() {
+  const cfg = join(__dirname, CUSTOM_BINARY_FILENAME);
+  return fs.existsSync(cfg) ? require(cfg) : {};
+}
+
+function getCustomBinary(name) {
+  return customBinaries[name];
+}
+
 module.exports = {
-  addonName: addonName
+  addonName: addonName,
+  customTFLibUri: customBinaries['tf-lib'],
+  customAddon: customBinaries['addon'],
+  ARCH_MAPPING,
+  PLATFORM_MAPPING,
+  PLATFORM_EXTENSION,
+  ALL_SUPPORTED_COMBINATION
 };

tfjs-node/scripts/install.js

@@ -30,23 +30,32 @@ const {
   modulePath
 } = require('./deps-constants.js');
 const resources = require('./resources');
-const {addonName} = require('./get-addon-name.js');
+const {
+  addonName,
+  customTFLibUri,
+  customAddon,
+  PLATFORM_MAPPING,
+  ARCH_MAPPING,
+  PLATFORM_EXTENSION,
+  ALL_SUPPORTED_COMBINATION,
+} = require('./get-addon-name.js');
 
 const exists = util.promisify(fs.exists);
 const mkdir = util.promisify(fs.mkdir);
 const rename = util.promisify(fs.rename);
 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`;
+  'https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-';
 
 const platform = os.platform();
+// Use windows path
+if (platform === 'win32') {
+  path = path.win32;
+}
 let libType = process.argv[2] === undefined ? 'cpu' : process.argv[2];
+const system = `${libType}-${PLATFORM_MAPPING[platform]}-` +
+    `${ARCH_MAPPING[os.arch()]}`;
 let forceDownload = process.argv[3] === undefined ? undefined : process.argv[3];
 
 let packageJsonFile;
@@ -56,14 +65,20 @@ async function setPackageJsonFile() {
       JSON.parse(fs.readFileSync(`${__dirname}/../package.json`).toString());
 }
 
-async function updateAddonName() {
-  packageJsonFile['binary']['package_name'] = addonName;
+function updateAddonName() {
+  const origBinary = JSON.parse(JSON.stringify(packageJsonFile['binary']));
+  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);
+  return origBinary;
 }
 
-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 +87,16 @@ 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}`);
+
+  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}`;
 }
 
 /**
@@ -157,15 +155,19 @@ async function downloadLibtensorflow(callback) {
  */
 async function build() {
   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();
   });
 }
 
@@ -176,7 +178,7 @@ async function run() {
   // Load package.json file
   setPackageJsonFile();
   // Update addon name in package.json file
-  await updateAddonName();
+  const origBinary = updateAddonName();
   // First check if deps library exists:
   if (forceDownload !== 'download' && await exists(depsLibTensorFlowPath)) {
     // Library has already been downloaded, then compile and simlink:
@@ -186,6 +188,7 @@ async function run() {
     await cleanDeps();
     await downloadLibtensorflow(build);
   }
+  revertAddonName(origBinary);
 }
 
 run();