[wasm] Cache wasm files when running Karma tests

tfjs-backend-wasm/src/BUILD.bazel

@@ -22,6 +22,7 @@ package(default_visibility = ["//visibility:public"])
 TEST_SRCS = [
     "**/*_test.ts",
     "test_node.ts",
+    "test_util.ts",
 ]
 
 TEST_ENTRYPOINTS = [

tfjs-backend-wasm/src/index_test.ts

@@ -22,8 +22,7 @@ import {ALL_ENVS, BROWSER_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/d
 
 import {init, resetWasmPath} from './backend_wasm';
 import {BackendWasm, setWasmPath, setWasmPaths} from './index';
-
-const VALID_PREFIX = '/base/tfjs/tfjs-backend-wasm/wasm-out/';
+import {VALID_PREFIX, setupCachedWasmPaths} from './test_util';
 
 /**
  * Tests specific to the wasm backend. The name of these tests must start with
@@ -79,9 +78,8 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
     spyOn(console, 'log');
   });
 
-  afterEach(() => {
-    resetWasmPath();
-    setWasmPaths(VALID_PREFIX);
+  afterEach(async () => {
+    await setupCachedWasmPaths();
     removeBackend('wasm-test');
   });
 
@@ -100,6 +98,7 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
   it('backend init fails when setWasmPaths is called with ' +
          'an invalid prefix',
      async () => {
+       resetWasmPath();
        setWasmPaths('invalid/prefix/');
        let wasmPath: string;
        const realFetch = fetch;
@@ -114,6 +113,7 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
   it('backend init fails when setWasmPaths is called with ' +
          'an invalid fileMap',
      async () => {
+       resetWasmPath();
        setWasmPaths({
          'tfjs-backend-wasm.wasm': 'invalid/path',
          'tfjs-backend-wasm-simd.wasm': 'invalid/path',
@@ -154,6 +154,7 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
 
     it('backend init works when the path is valid', async () => {
       const usePlatformFetch = true;
+      resetWasmPath();
       setWasmPaths(VALID_PREFIX, usePlatformFetch);
       let wasmPath: string;
       fetchSpy.and.callFake((path: string) => {

tfjs-backend-wasm/src/setup_test.ts

@@ -16,21 +16,16 @@
  */
 
 // Import core for side effects (e.g. flag registration)
-import * as tf from '@tensorflow/tfjs-core';
+import '@tensorflow/tfjs-core';
 // tslint:disable-next-line:no-imports-from-dist
 import '@tensorflow/tfjs-core/dist/public/chained_ops/register_all_chained_ops';
 // tslint:disable-next-line: no-imports-from-dist
 import '@tensorflow/tfjs-core/dist/register_all_gradients';
 // Register the wasm backend.
-import * as wasmBackend from './index';
+import './index';
 // tslint:disable-next-line: no-imports-from-dist
 import {setTestEnvs, setupTestFilters, TestFilter} from '@tensorflow/tfjs-core/dist/jasmine_util';
-
-// Bazel's karma_web_test does not support setting proxies for loaded files,
-// so set the wasm path to where it serves them.
-if (tf.device_util.isBrowser()) {
-  wasmBackend.setWasmPaths('/base/tfjs/tfjs-backend-wasm/wasm-out/');
-}
+import {setupCachedWasmPaths} from './test_util';
 
 setTestEnvs([{name: 'test-wasm', backendName: 'wasm', isDataSync: true}]);
 
@@ -414,6 +409,8 @@ const customInclude = (testName: string) => {
 };
 setupTestFilters(TEST_FILTERS, customInclude);
 
+beforeAll(setupCachedWasmPaths);
+
 // Import and run all the tests from core.
 // tslint:disable-next-line:no-imports-from-dist
 // tslint:disable-next-line:no-require-imports

tfjs-backend-wasm/src/test_util.ts

@@ -0,0 +1,40 @@
+import * as tf from '@tensorflow/tfjs-core';
+import {resetWasmPath} from './backend_wasm';
+import {setWasmPaths} from './index';
+
+export const VALID_PREFIX = '/base/tfjs/tfjs-backend-wasm/wasm-out/';
+
+async function cacheUrl(url: string): Promise<string> {
+  const blob = await (await fetch(url)).blob();
+  return URL.createObjectURL(blob);
+}
+
+let cachedUrlsPromises: Array<Promise<readonly [
+  'tfjs-backend-wasm.wasm'
+    | 'tfjs-backend-wasm-simd.wasm'
+    | 'tfjs-backend-wasm-threaded-simd.wasm', string]>>;
+
+if (tf.device_util.isBrowser()) {
+  cachedUrlsPromises = ([
+    'tfjs-backend-wasm.wasm',
+    'tfjs-backend-wasm-simd.wasm',
+    'tfjs-backend-wasm-threaded-simd.wasm',
+  ] as const).map(async name => {
+    return [name, await cacheUrl(VALID_PREFIX + name)] as const;
+  });
+}
+
+/**
+ * Set the wasm paths to the blob URLs that were loaded above. This is useful
+ * for Karma tests because it prevents Karma from loading the same WASM files
+ * every time the backend is reset (each `describe` block). This saves ~200MB
+ * of network and should reduce flakiness due to network slowness / instability.
+ */
+export async function setupCachedWasmPaths() {
+  resetWasmPath();
+  if (tf.device_util.isBrowser()) {
+    const cachedUrlsList = await Promise.all(cachedUrlsPromises);
+    const cachedUrls = Object.fromEntries(cachedUrlsList);
+    setWasmPaths(cachedUrls);
+  }
+}