Merge pull request #26 from PAIR-code/iislucas-30-aug-2024-webworker-…

…fixes

fixed webworker cell calling and library

animated-transformer/src/app/web-colab/app.worker.ts

@@ -53,6 +53,8 @@ async function run() {
     t: t.toSerialised(),
     v,
   });
+
+  cell.finished();
 }
 
 run();

animated-transformer/src/app/web-colab/foo.ailab.ts

@@ -0,0 +1,55 @@
+/* Copyright 2023 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.
+==============================================================================*/
+/**
+ * This is a simple example (web) ailab. This provides an example of defining
+ * the types for a cell.
+ */
+
+import { SerializedGTensor } from 'src/lib/gtensor/gtensor';
+import { CellSpec } from '../../lib/weblab/cellspec';
+
+export type Name = string;
+export type TensorValue = {
+  t: SerializedGTensor<'a'>;
+  v: number;
+} | null;
+
+export type Globals = {
+  name: Name;
+  tensor: TensorValue;
+};
+
+export type GlobalValue<Name extends string> = { [Key in keyof Globals & Name]: Globals[Key] };
+
+const globals: Partial<Globals> = {
+  name: 'some silly fake initial name',
+};
+
+// export const exampleWorkerOp = {
+//   workerPath: './app.worker',
+//   inputs: ['name'] as const,
+//   outputs: ['t'] as const,
+// } as WorkerOp<'name', 't'>;
+
+export const exampleWorkerSpec = new CellSpec<GlobalValue<'name'>, GlobalValue<'tensor'>>(
+  'example app worker',
+  // 'src/lib/weblab/example.worker.js' as never as URL,
+  // Hack because angular dev builder does a regexp replacement, so we need the full string of
+  // new Worker(new URL('<literal path>', import.meta.url)) in order for dev server and prod
+  // build to correctly create these paths.
+  () => new Worker(new URL('./app.worker', import.meta.url)),
+  ['name'], // new URL('http://localhost:9876/_karma_webpack_/example.worker'),
+  ['tensor']
+);

animated-transformer/src/app/web-colab/web-colab.component.ts

@@ -16,6 +16,9 @@ limitations under the License.
 import { Component } from '@angular/core';
 import { GTensor, SerializedGTensor, makeScalar } from 'src/lib/gtensor/gtensor';
 import * as tf from '@tensorflow/tfjs';
+import { LabState } from 'src/lib/weblab/lab-state';
+import { LabEnv } from 'src/lib/weblab/lab-env';
+import { exampleWorkerSpec, Globals } from './foo.ailab';
 
 // Create a new
 async function onceOutput<T>(worker: Worker): Promise<T> {
@@ -34,11 +37,11 @@ async function onceOutput<T>(worker: Worker): Promise<T> {
   styleUrl: './web-colab.component.scss',
 })
 export class WebColabComponent {
-  public worker: Worker;
+  // public worker2: Worker;
 
   constructor() {
-    this.worker = new Worker(new URL('./app.worker', import.meta.url));
-    // this.foo();
+    // This works...
+    // this.worker2 = new Worker(new URL('./app.worker', import.meta.url));
   }
 
   async foo() {
@@ -55,18 +58,14 @@ export class WebColabComponent {
       console.error('We require webworkers. Sorry.');
       return;
     }
-    this.worker.postMessage('hello, are you there webworker?');
-    console.log('posted message');
-    // Create a new
-    const output = await onceOutput<{
-      data: { t: SerializedGTensor<'a'>; v: number };
-    }>(this.worker);
-    console.log('webworker completed');
-    console.log(output);
-    console.log(GTensor.fromSerialised(output.data.t).scalarDiv(makeScalar(3)).tensor.arraySync());
-    console.log(output.data.v);
-
-    // const myWorker = new Worker('worker.js');
+    const state = new LabState();
+    const env = new LabEnv<Globals>(state);
+    env.stateVars.name = 'initial fake name';
+    const outputs = await env.run(exampleWorkerSpec);
+    console.log(outputs);
+    console.log(
+      GTensor.fromSerialised(outputs.tensor!.t).scalarDiv(makeScalar(3)).tensor.arraySync()
+    );
   }
 
   async doOpen() {

animated-transformer/src/lib/seqtasks/tiny_worlds_generate.script.ts

@@ -13,7 +13,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-// Run with: ts-node src/lib/seqtasks/tiny_worlds.script.ts
+/*
+Run with:
+  npx ts-node src/lib/seqtasks/tiny_worlds_generate.script.ts
+*/
 
 import {
   TinyWorldTask,
@@ -38,16 +41,16 @@ import {
   console.log(' - output: ', JSON.stringify(example2.output.join('')));
 }
 
-{
-  const initConfig = { ...bayesianV1TinyWorldTaskConfig };
-  initConfig.maxInputLen = 400;
-  initConfig.maxOutputLen = 400;
-  const tinyWorld = new TinyWorldTask(initConfig);
-  const [example] = tinyWorld.exampleIter.takeOutN(1);
-  console.log('Example 1');
-  console.log('input: ', JSON.stringify(example.input.join('')));
-  console.log('output: ', JSON.stringify(example.output.join('')));
-  console.log('Example 1');
-  console.log('input: ', JSON.stringify(example.input.join('')));
-  console.log('output: ', JSON.stringify(example.output.join('')));
-}
+// {
+//   const initConfig = { ...bayesianV1TinyWorldTaskConfig };
+//   initConfig.maxInputLen = 400;
+//   initConfig.maxOutputLen = 400;
+//   const tinyWorld = new TinyWorldTask(initConfig);
+//   const [example] = tinyWorld.exampleIter.takeOutN(1);
+//   console.log('Example 1');
+//   console.log('input: ', JSON.stringify(example.input.join('')));
+//   console.log('output: ', JSON.stringify(example.output.join('')));
+//   console.log('Example 1');
+//   console.log('input: ', JSON.stringify(example.input.join('')));
+//   console.log('output: ', JSON.stringify(example.output.join('')));
+// }

animated-transformer/src/lib/seqtasks/tiny_worlds_train.script.ts

@@ -17,7 +17,7 @@ limitations under the License.
 Tiny Worlds, run with (gtensor-based) transformers.
 
 Run:
-npx ts-node src/lib/seqtasks/tiny_worlds.run_with_transformer.script.ts
+  npx ts-node src/lib/seqtasks/tiny_worlds_train.script.ts
 */
 
 import * as tf from '@tensorflow/tfjs-node';

animated-transformer/src/lib/weblab/cellspec.ts

@@ -38,7 +38,8 @@ export type InputPromises<S extends ValueStruct> = { [Key in keyof S]: Promise<S
 
 export class CellSpec<Inputs extends ValueStruct, Outputs extends ValueStruct> {
   constructor(
-    public workerPath: URL,
+    public name: string,
+    public createWorker: () => Worker,
     public inputs: (keyof Inputs)[],
     public outputs: (keyof Outputs)[]
   ) {}

animated-transformer/src/lib/weblab/example.ailab.ts

@@ -44,8 +44,9 @@ const globals: Partial<Globals> = {
 // } as WorkerOp<'name', 't'>;
 
 export const exampleWorkerSpec = new CellSpec<GlobalValue<'name'>, GlobalValue<'tensor'>>(
+  'an example cell',
   // 'src/lib/weblab/example.worker.js' as never as URL,
-  new URL('./example.worker', import.meta.url),
+  () => new Worker(new URL('./example.worker', import.meta.url)),
   ['name'], // new URL('http://localhost:9876/_karma_webpack_/example.worker'),
   ['tensor']
 );

animated-transformer/src/lib/weblab/example.worker.ts

@@ -39,6 +39,10 @@ async function run() {
     t: t.toSerialised(),
     v,
   });
+
+  console.log('worker going to finish...');
+  cell.finished();
+  console.log('worker finished.');
 }
 
 run();

animated-transformer/src/lib/weblab/lab-env.spec.ts

@@ -18,7 +18,7 @@ import { LabState } from './lab-state';
 
 import { exampleWorkerSpec, Globals } from './example.ailab';
 
-xdescribe('lab-env', () => {
+describe('lab-env', () => {
   const state = new LabState();
   beforeEach(async () => {});
 
@@ -28,8 +28,4 @@ xdescribe('lab-env', () => {
     const outputs = await env.run(exampleWorkerSpec);
     expect(outputs.tensor).toBeTruthy();
   });
-
-  it('ignoreme', () => {
-    expect(true).toBeTruthy();
-  });
 });

animated-transformer/src/lib/weblab/lab-env.ts

@@ -29,51 +29,71 @@ export class LabEnv<Globals extends ValueStruct> {
   inputFiles: Map<keyof Globals, FileSystemFileHandle> = new Map();
   stateVars: Partial<Globals> = {};
   metadata: Map<keyof Globals, ItemMetaData> = new Map();
+  runningCells: {
+    [name: string]: CellSpec<{}, {}>;
+  } = {};
 
-  constructor(public workerState: LabState) {}
+  constructor(public workerState: LabState) {
+    console.log('import.meta.url', import.meta.url.toString());
+  }
 
   async run<I extends keyof Globals & string, O extends keyof Globals & string>(
     op: CellSpec<SpecificValueStruct<I>, SpecificValueStruct<O>>
   ): Promise<{ [key in O]: Globals[O] }> {
+    this.runningCells[op.name] = op as CellSpec<{}, {}>;
+
     const outputs = {} as { [key in O]: Globals[O] };
     // Ensure inputs in memory.
     for (const inputName of op.inputs) {
       if (this.stateVars[inputName] === undefined) {
         const inputValue = await this.workerState.loadValue<Globals[O]>(inputName);
         if (!inputValue) {
-          throw new Error(`No state for op (${op.workerPath}) for input: ${inputName}`);
+          throw new Error(`No state for op (${op.createWorker}) for input: ${inputName}`);
         }
         this.stateVars[inputName] = inputValue.data;
       }
     }
 
-    // const workerUrl = new URL(op.workerUrl, import.meta.url);
-    // console.log('op.workerPath:', op.workerPath);
-    // console.log('import.meta.url:', import.meta.url);
-    // console.log('workerUrl:', workerUrl);
-    const worker = new Worker(op.workerPath, { type: 'module' });
+    let resolveWithOutputFn: (output: { [key in O]: Globals[O] }) => void;
+    const onceFinished = new Promise<{ [key in O]: Globals[O] }>((resolve, reject) => {
+      resolveWithOutputFn = resolve;
+    });
+
+    const worker = op.createWorker();
+    // const worker = new Worker(op.workerPath, { type: 'module' });
     // console.log(worker);
     worker.onmessage = ({ data }) => {
+      console.log('main thread got worker.onmessage', data);
       const messageFromWorker = data as FromWorkerMessage;
       switch (messageFromWorker.kind) {
-        case 'finished':
-          worker.terminate();
-          break;
         case 'requestInput':
-          worker.postMessage(this.stateVars[messageFromWorker.name]);
+          console.log(
+            'this.stateVars[messageFromWorker.name]',
+            this.stateVars[messageFromWorker.name]
+          );
+          worker.postMessage({
+            kind: 'providingInput',
+            name: messageFromWorker.name,
+            inputData: this.stateVars[messageFromWorker.name],
+          });
+          break;
+        // only called when the webworker is really finished.
+        case 'finished':
+          delete this.runningCells[op.name];
+          resolveWithOutputFn(outputs);
           break;
         case 'providingOutput':
           const outputName = messageFromWorker.name as O;
           outputs[outputName] = messageFromWorker.outputData as Globals[O];
           break;
         default:
-          console.error('unknown worker message: ', data);
+          console.error('main thread go unknown worker message: ', data);
           break;
       }
     };
 
     // worker.onmessage(() => {});
 
-    return outputs;
+    return onceFinished;
   }
 }

animated-transformer/src/lib/weblab/workerlab.ts

@@ -6,10 +6,6 @@ import { InputPromises, ValueStruct, CellSpec as CellSpec } from './cellspec';
 
 export const space = new SignalSpace();
 
-function workerToMainMessage(m: FromWorkerMessage) {
-  postMessage(m);
-}
-
 const initInputs = {} as { [name: string]: unknown };
 // const recievedInputs = space.writable(initInputs);
 const inputResolvers = {} as { [name: string]: (value: unknown) => void };
@@ -20,12 +16,16 @@ addEventListener('message', ({ data }) => {
     initInputs[toWorkerMessage.name] = toWorkerMessage.inputData;
     if (toWorkerMessage.name in inputResolvers) {
       inputResolvers[toWorkerMessage.name](toWorkerMessage.inputData);
+    } else {
+      console.warn('got sent an input we do not know about: ', data);
     }
+  } else {
+    console.warn('unknown message from the main thread: ', data);
   }
 });
 
 export function onceGetInput<T>(name: string): Promise<T> {
-  workerToMainMessage({ kind: 'requestInput', name });
+  postMessage({ kind: 'requestInput', name });
   return new Promise<T>((resolve, reject) => {
     // TODO: consider allowing parent to send stuff before we ask for it..
     // this would just involved checking the inputResolvers here.
@@ -34,7 +34,7 @@ export function onceGetInput<T>(name: string): Promise<T> {
 }
 
 export function sendOutput<T>(name: string, outputData: T) {
-  workerToMainMessage({ kind: 'providingOutput', name, outputData });
+  postMessage({ kind: 'providingOutput', name, outputData });
 }
 
 // export class LabCell<Globals extends { [key: string]: any }, I extends string, O extends string> {
@@ -55,4 +55,9 @@ export class Cell<Input extends ValueStruct, Output extends ValueStruct> {
   output<Key extends keyof Output>(key: Key, value: Output[Key]) {
     sendOutput(key as string, value);
   }
+
+  finished() {
+    postMessage({ kind: 'finished' });
+    close();
+  }
 }