@tensorflow/tfjs-node/scripts/install.js

This repository provides native TensorFlow execution in backend JavaScript applications under the Node.js runtime, accelerated by the TensorFlow C binary under the hood. It provides the same API as [TensorFlow.js](https://js.tensorflow.org/api/latest/).

/**
 * @license
 * Copyright 2019 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.
 * =============================================================================
 */

const fs = require('fs');
let path = require('path');
const rimraf = require('rimraf');
const util = require('util');
const cp = require('child_process');
const os = require('os');
const {
  depsPath,
  depsLibPath,
  depsLibTensorFlowPath,
  LIBTENSORFLOW_VERSION,
  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');

const exists = util.promisify(fs.exists);
const mkdir = util.promisify(fs.mkdir);
const rename = util.promisify(fs.rename);
const rimrafPromise = util.promisify(rimraf);

const CDN_STORAGE = process.env.TFJS_NODE_CDN_STORAGE ||
    process.env.npm_config_TFJS_NODE_CDN_STORAGE || process.env.CDN_STORAGE;
const BASE_HOST = CDN_STORAGE || 'https://storage.googleapis.com/';
const BASE_URI = process.env.TFJS_NODE_BASE_URI ||
    process.env.npm_config_TFJS_NODE_BASE_URI ||
    `${BASE_HOST}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];
let system = `${libType}-${PLATFORM_MAPPING[platform]}-` +
    `${ARCH_MAPPING[os.arch()]}`;
let forceDownload = process.argv[3] === undefined ? undefined : process.argv[3];

let packageJsonFile;

function setPackageJsonFile() {
  packageJsonFile =
      JSON.parse(fs.readFileSync(`${__dirname}/../package.json`).toString());
}

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);
}

function revertAddonName(orig) {
  packageJsonFile['binary'] = orig;
  const stringFile = JSON.stringify(packageJsonFile, null, 2).concat('\n');
  fs.writeFileSync((`${__dirname}/../package.json`), stringFile);
}

/**
 * Returns the libtensorflow hosted path of the current platform.
 */
function getPlatformLibtensorflowUri() {
  // Exception for mac+gpu user
  if (platform === 'darwin') {
    if (os.arch() === 'arm64') {
      return `${BASE_HOST}tf-builds/libtensorflow_r2_7_darwin_arm64_cpu.tar.gz`;
    }
    system = `cpu-${PLATFORM_MAPPING[platform]}-${ARCH_MAPPING[os.arch()]}`;
  }

  if (customTFLibUri !== undefined) {
    return customTFLibUri;
  }

  if (platform === 'linux') {
    if (os.arch() === 'arm') {
      return `${BASE_HOST}tf-builds/libtensorflow_r2_5_linux_arm7l.tar.gz`;
    } else if (os.arch() === 'arm64') {
      return `${BASE_HOST}tf-builds/libtensorflow_r2_7_linux_arm64.tar.gz`;
    }
  }

  if (ALL_SUPPORTED_COMBINATION.indexOf(system) === -1) {
    throw new Error(`Unsupported system: ${libType}-${platform}-${os.arch()}`);
  }

  return `${BASE_URI}${system}-${LIBTENSORFLOW_VERSION}.${PLATFORM_EXTENSION}`;
}

/**
 * Ensures a directory exists, creates as needed.
 */
async function ensureDir(dirPath) {
  if (!await exists(dirPath)) {
    await mkdir(dirPath);
  }
}

/**
 * Deletes the deps directory if it exists, and creates a fresh deps folder.
 */
async function cleanDeps() {
  if (await exists(depsPath)) {
    await rimrafPromise(depsPath);
  }
  await mkdir(depsPath);
}

/**
 * Downloads libtensorflow and notifies via a callback when unpacked.
 */
async function downloadLibtensorflow(callback) {
  // Ensure dependencies staged directory is available:
  await ensureDir(depsPath);

  console.warn('* Downloading libtensorflow');
  console.log(getPlatformLibtensorflowUri());
  resources.downloadAndUnpackResource(
      getPlatformLibtensorflowUri(), depsPath, async () => {
        if (platform === 'win32') {
          // Some windows libtensorflow zip files are missing structure and the
          // eager headers. Check, restructure, and download resources as
          // needed.
          if (!await exists(depsLibTensorFlowPath)) {
            // Verify that tensorflow.dll exists
            const libtensorflowDll = path.join(depsPath, 'tensorflow.dll');
            if (!await exists(libtensorflowDll)) {
              throw new Error('Could not find libtensorflow.dll');
            }

            await ensureDir(depsLibPath);
            await rename(libtensorflowDll, depsLibTensorFlowPath);
          }
        }
        // No other work is required on other platforms.
        if (callback !== undefined) {
          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');
  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);
      process.exit(1);
    }
    if (platform === 'win32') {
      // Move libtensorflow to module path, where tfjs_binding.node locates.
      cp.exec('node scripts/deps-stage.js symlink ' + modulePath, (error) => {
        if (error) {
          console.error('symlink ' + modulePath + ' failed: ', error);
          process.exit(1);
        }
      });
    }
    revertAddonName(origBinary);
  });
}

/**
 * Ensures libtensorflow requirements are met for building the binding.
 */
async function run() {
  // First check if deps library exists:
  if (forceDownload !== 'download' && await exists(depsLibTensorFlowPath)) {
    // Library has already been downloaded, then compile and simlink:
    await build();
  } else {
    // Library has not been downloaded, download, then compile and symlink:
    await cleanDeps();
    await downloadLibtensorflow(build);
  }
}

run();