Refactor IO-related code to use the new LoadOptions interface (tensorflow#1554)

- Refactor several IO-related functions to use the new `LoadOptions` interface.
- Move `monitorProgress` to the io-specific source directory.

DEV

Author: caisq

src/io/browser_http.ts

@@ -25,7 +25,7 @@ import {assert} from '../util';
 
 import {concatenateArrayBuffers, getModelArtifactsInfoForJSON} from './io_utils';
 import {IORouter, IORouterRegistry} from './router_registry';
-import {IOHandler, ModelArtifacts, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
+import {IOHandler, LoadOptions, ModelArtifacts, OnProgressCallback, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
 import {loadWeightsAsArrayBuffer} from './weights_loader';
 
 const OCTET_STREAM_MIME_TYPE = 'application/octet-stream';
@@ -41,30 +41,37 @@ export class BrowserHTTPRequest implements IOHandler {
 
   static readonly URL_SCHEME_REGEX = /^https?:\/\//;
 
-  constructor(
-      path: string|string[], requestInit?: RequestInit,
-      private readonly weightPathPrefix?: string, fetchFunc?: Function,
-      private readonly onProgress?: Function) {
-    if (fetchFunc == null) {
+  private readonly weightPathPrefix: string;
+  private readonly onProgress: OnProgressCallback;
+
+  constructor(path: string|string[], loadOptions?: LoadOptions) {
+    if (loadOptions == null) {
+      loadOptions = {};
+    }
+    this.weightPathPrefix = loadOptions.weightPathPrefix;
+    this.onProgress = loadOptions.onProgress;
+
+    if (loadOptions.fetchFunc == null) {
       if (typeof fetch === 'undefined') {
         throw new Error(
             'browserHTTPRequest is not supported outside the web browser ' +
             'without a fetch polyfill.');
       }
       // Make sure fetch is always bound to window (the
       // original object) when available.
-      fetchFunc = fetch.bind(typeof window === 'undefined' ? null : window);
+      loadOptions.fetchFunc =
+          fetch.bind(typeof window === 'undefined' ? null : window);
     } else {
       assert(
-          typeof fetchFunc === 'function',
+          typeof loadOptions.fetchFunc === 'function',
           () => 'Must pass a function that matches the signature of ' +
               '`fetch` (see ' +
               'https://developer.mozilla.org/en-US/docs/Web/API/Fetch_API)');
     }
 
     this.fetchFunc = (path: string, requestInits: RequestInit) => {
       // tslint:disable-next-line:no-any
-      return fetchFunc(path, requestInits).catch((error: any) => {
+      return loadOptions.fetchFunc(path, requestInits).catch((error: any) => {
         throw new Error(`Request for ${path} failed due to error: ${error}`);
       });
     };
@@ -83,11 +90,12 @@ export class BrowserHTTPRequest implements IOHandler {
     }
     this.path = path;
 
-    if (requestInit != null && requestInit.body != null) {
+    if (loadOptions.requestInit != null &&
+        loadOptions.requestInit.body != null) {
       throw new Error(
           'requestInit is expected to have no pre-existing body, but has one.');
     }
-    this.requestInit = requestInit || {};
+    this.requestInit = loadOptions.requestInit || {};
   }
 
   async save(modelArtifacts: ModelArtifacts): Promise<SaveResult> {
@@ -236,8 +244,11 @@ export class BrowserHTTPRequest implements IOHandler {
         fetchURLs.push(pathPrefix + path + suffix);
       });
     });
-    const buffers = await loadWeightsAsArrayBuffer(
-        fetchURLs, this.requestInit, this.getFetchFunc(), this.onProgress);
+    const buffers = await loadWeightsAsArrayBuffer(fetchURLs, {
+      requestInit: this.requestInit,
+      fetchFunc: this.getFetchFunc(),
+      onProgress: this.onProgress
+    });
     return [weightSpecs, concatenateArrayBuffers(buffers)];
   }
 
@@ -278,7 +289,7 @@ export function isHTTPScheme(url: string): boolean {
 }
 
 export const httpRequestRouter: IORouter =
-    (url: string|string[], onProgress?: Function) => {
+    (url: string|string[], onProgress?: OnProgressCallback) => {
       if (typeof fetch === 'undefined') {
         // browserHTTPRequest uses `fetch`, if one wants to use it in node.js
         // they have to setup a global fetch polyfill.
@@ -291,7 +302,7 @@ export const httpRequestRouter: IORouter =
           isHTTP = isHTTPScheme(url);
         }
         if (isHTTP) {
-          return browserHTTPRequest(url, null, null, null, onProgress);
+          return browserHTTPRequest(url, {onProgress});
         }
       }
       return null;
@@ -434,17 +445,17 @@ IORouterRegistry.registerLoadRouter(httpRequestRouter);
  * topology (filename: 'model.json') and the weights of the model (filename:
  * 'model.weights.bin') will be appended to the body. If `requestInit` has a
  * `body`, an Error will be thrown.
- * @param weightPathPrefix Optional, this specifies the path prefix for weight
- *   files, by default this is calculated from the path param.
- * @param fetchFunc Optional, custom `fetch` function. E.g., in Node.js,
- *   the `fetch` from node-fetch can be used here.
- * @param onProgress Optional, progress callback function, fired periodically
- *   before the load is completed.
+ * @param loadOptions Optional configuration for the loading. It includes the
+ *   following fields:
+ *   - weightPathPrefix Optional, this specifies the path prefix for weight
+ *     files, by default this is calculated from the path param.
+ *   - fetchFunc Optional, custom `fetch` function. E.g., in Node.js,
+ *     the `fetch` from node-fetch can be used here.
+ *   - onProgress Optional, progress callback function, fired periodically
+ *     before the load is completed.
  * @returns An instance of `IOHandler`.
  */
 export function browserHTTPRequest(
-    path: string|string[], requestInit?: RequestInit, weightPathPrefix?: string,
-    fetchFunc?: Function, onProgress?: Function): IOHandler {
-  return new BrowserHTTPRequest(
-      path, requestInit, weightPathPrefix, fetchFunc, onProgress);
+    path: string|string[], loadOptions?: LoadOptions): IOHandler {
+  return new BrowserHTTPRequest(path, loadOptions);
 }

src/io/browser_http_test.ts

@@ -303,9 +303,11 @@ describeWithFlags('browserHTTPRequest-save', CHROME_ENVS, () => {
   it('Save topology and weights, PUT method, extra headers', (done) => {
     const testStartDate = new Date();
     const handler = tf.io.browserHTTPRequest('model-upload-test', {
-      method: 'PUT',
-      headers:
-          {'header_key_1': 'header_value_1', 'header_key_2': 'header_value_2'}
+      requestInit: {
+        method: 'PUT',
+        headers:
+            {'header_key_1': 'header_value_1', 'header_key_2': 'header_value_2'}
+      }
     });
     handler.save(artifacts1)
         .then(saveResult => {
@@ -392,7 +394,9 @@ describeWithFlags('browserHTTPRequest-save', CHROME_ENVS, () => {
 
   it('Existing body leads to Error', () => {
     expect(() => tf.io.browserHTTPRequest('model-upload-test', {
-      body: 'existing body'
+      requestInit: {
+        body: 'existing body'
+      }
     })).toThrowError(/requestInit is expected to have no pre-existing body/);
   });
 
@@ -518,8 +522,11 @@ describeWithFlags('browserHTTPRequest-load', BROWSER_ENVS, () => {
           },
           requestInits);
 
-      const handler = tf.io.browserHTTPRequest(
-          './model.json', {headers: {'header_key_1': 'header_value_1'}});
+      const handler = tf.io.browserHTTPRequest('./model.json', {
+        requestInit: {
+          headers: {'header_key_1': 'header_value_1'}
+        }
+      });
       const modelArtifacts = await handler.load();
       expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
       expect(modelArtifacts.weightSpecs).toEqual(weightManifest1[0].weights);
@@ -899,8 +906,8 @@ describeWithFlags('browserHTTPRequest-load', BROWSER_ENVS, () => {
           requestInits);
 
       const handler = tf.io.browserHTTPRequest(
-          ['./model.pb', './weights_manifest.json'],
-          {headers: {'header_key_1': 'header_value_1'}});
+          ['./model.pb', './weights_manifest.json'], {
+            requestInit: {headers: {'header_key_1': 'header_value_1'}}});
       const modelArtifacts = await handler.load();
       expect(modelArtifacts.modelTopology).toEqual(modelData);
       expect(modelArtifacts.weightSpecs).toEqual(weightManifest1[0].weights);
@@ -1089,7 +1096,8 @@ describeWithFlags('browserHTTPRequest-load', BROWSER_ENVS, () => {
              },
              requestInits);
          const handler = tf.io.browserHTTPRequest(
-             ['path1/model.pb', 'path2/weights_manifest.json'], {}, 'path3/');
+             ['path1/model.pb', 'path2/weights_manifest.json'],
+             {weightPathPrefix: 'path3/'});
          const modelArtifacts = await handler.load();
          expect(modelArtifacts.modelTopology).toEqual(modelData);
          expect(modelArtifacts.weightSpecs)
@@ -1171,7 +1179,10 @@ describeWithFlags('browserHTTPRequest-load', BROWSER_ENVS, () => {
     }
 
     const handler = tf.io.browserHTTPRequest(
-        './model.json', {credentials: 'include'}, null, customFetch);
+        './model.json', {
+          requestInit: {credentials: 'include'},
+          fetchFunc: customFetch
+        });
     const modelArtifacts = await handler.load();
     expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
     expect(modelArtifacts.weightSpecs).toEqual(weightManifest1[0].weights);

src/io/progress.ts

@@ -0,0 +1,73 @@
+/**
+ * @license
+ * Copyright 2019 Google Inc. 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.
+ * =============================================================================
+ */
+
+import {assert} from '../util';
+
+import {OnProgressCallback} from './types';
+
+/**
+ * Monitor Promise.all progress, fire onProgress callback function.
+ *
+ * @param promises Promise list going to be monitored
+ * @param onProgress Callback function. Fired when a promise resolved.
+ * @param startFraction Optional fraction start. Default to 0.
+ * @param endFraction Optional fraction end. Default to 1.
+ */
+export function monitorPromisesProgress(
+    promises: Array<Promise<{}|void>>, onProgress: OnProgressCallback,
+    startFraction?: number, endFraction?: number) {
+  checkPromises(promises);
+  startFraction = startFraction == null ? 0 : startFraction;
+  endFraction = endFraction == null ? 1 : endFraction;
+  checkFraction(startFraction, endFraction);
+  let resolvedPromise = 0;
+
+  const registerMonitor = (promise: Promise<{}>) => {
+    promise.then(value => {
+      const fraction = startFraction +
+          ++resolvedPromise / promises.length * (endFraction - startFraction);
+      // pass fraction as parameter to callback function.
+      onProgress(fraction);
+      return value;
+    });
+    return promise;
+  };
+
+  function checkPromises(promises: Array<Promise<{}|void>>): void {
+    assert(
+        promises != null && Array.isArray(promises) && promises.length > 0,
+        () => 'promises must be a none empty array');
+  }
+
+  function checkFraction(startFraction: number, endFraction: number): void {
+    assert(
+        startFraction >= 0 && startFraction <= 1,
+        () => `Progress fraction must be in range [0, 1], but ` +
+            `got startFraction ${startFraction}`);
+    assert(
+        endFraction >= 0 && endFraction <= 1,
+        () => `Progress fraction must be in range [0, 1], but ` +
+            `got endFraction ${endFraction}`);
+    assert(
+        endFraction >= startFraction,
+        () => `startFraction must be no more than endFraction, but ` +
+            `got startFraction ${startFraction} and endFraction ` +
+            `${endFraction}`);
+  }
+
+  return Promise.all(promises.map(registerMonitor));
+}

src/io/progress_test.ts

@@ -0,0 +1,86 @@
+/**
+ * @license
+ * Copyright 2019 Google Inc. 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.
+ * =============================================================================
+ */
+
+import {monitorPromisesProgress} from './progress';
+
+describe('util.monitorPromisesProgress', () => {
+  it('Default progress from 0 to 1', (done) => {
+    const expectFractions: number[] = [0.25, 0.50, 0.75, 1.00];
+    const fractionList: number[] = [];
+    const tasks = Array(4).fill(0).map(() => {
+      return Promise.resolve();
+    });
+    monitorPromisesProgress(tasks, (progress: number) => {
+      fractionList.push(parseFloat(progress.toFixed(2)));
+    }).then(() => {
+      expect(fractionList).toEqual(expectFractions);
+      done();
+    });
+  });
+
+  it('Progress with pre-defined range', (done) => {
+    const startFraction = 0.2;
+    const endFraction = 0.8;
+    const expectFractions: number[] = [0.35, 0.50, 0.65, 0.80];
+    const fractionList: number[] = [];
+    const tasks = Array(4).fill(0).map(() => {
+      return Promise.resolve();
+    });
+    monitorPromisesProgress(tasks, (progress: number) => {
+      fractionList.push(parseFloat(progress.toFixed(2)));
+    }, startFraction, endFraction).then(() => {
+      expect(fractionList).toEqual(expectFractions);
+      done();
+    });
+  });
+
+  it('throws error when progress fraction is out of range', () => {
+    expect(() => {
+      const startFraction = -1;
+      const endFraction = 1;
+      const tasks = Array(4).fill(0).map(() => {
+        return Promise.resolve();
+      });
+      monitorPromisesProgress(
+          tasks, (progress: number) => {}, startFraction, endFraction);
+    }).toThrowError();
+  });
+
+  it('throws error when startFraction more than endFraction', () => {
+    expect(() => {
+      const startFraction = 0.8;
+      const endFraction = 0.2;
+      const tasks = Array(4).fill(0).map(() => {
+        return Promise.resolve();
+      });
+      monitorPromisesProgress(
+          tasks, (progress: number) => {}, startFraction, endFraction);
+    }).toThrowError();
+  });
+
+  it('throws error when promises is null', () => {
+    expect(() => {
+      monitorPromisesProgress(null, (progress: number) => {});
+    }).toThrowError();
+  });
+
+  it('throws error when promises is empty array', () => {
+    expect(() => {
+      monitorPromisesProgress([], (progress: number) => {});
+    }).toThrowError();
+  });
+});