apache--tvm/web/src/rpc_server.ts

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 { SizeOf, TypeIndex } from "./ctypes";
import { assert, StringToUint8Array, Uint8ArrayToString } from "./support";
import { detectGPUDevice, GPUDeviceDetectOutput } from "./webgpu";
import * as compact from "./compact";
import * as runtime from "./runtime";
import { Disposable } from "./types";

enum RPCServerState {
  InitHeader,
  InitHeaderKey,
  InitServer,
  WaitForCallback,
  ReceivePacketHeader,
  ReceivePacketBody,
}

/** RPC magic header */
const RPC_MAGIC = 0xff271;

/**
 * An utility class to read from binary bytes.
 */
class ByteStreamReader {
  offset = 0;
  bytes: Uint8Array;

  constructor(bytes: Uint8Array) {
    this.bytes = bytes;
  }

  readU32(): number {
    const i = this.offset;
    const b = this.bytes;
    const val = b[i] | (b[i + 1] << 8) | (b[i + 2] << 16) | (b[i + 3] << 24);
    this.offset += 4;
    return val;
  }

  readU64(): number {
    const val = this.readU32();
    this.offset += 4;
    return val;
  }

  readByteArray(): Uint8Array {
    const len = this.readU64();
    assert(this.offset + len <= this.bytes.byteLength);
    const ret = new Uint8Array(len);
    ret.set(this.bytes.slice(this.offset, this.offset + len));
    this.offset += len;
    return ret;
  }
}

/**
 * A websocket based RPC
 */
export class RPCServer {
  url: string;
  key: string;
  socket: WebSocket;
  state: RPCServerState = RPCServerState.InitHeader;
  logger: (msg: string) => void;
  getImports: () => Record<string, unknown>;
  private tensorCacheUrl: string;
  private tensorCacheDevice: string;
  private initProgressCallback?: runtime.InitProgressCallback;
  private asyncOnServerLoad?: (inst: runtime.Instance) => Promise<void>;
  private pendingSend: Promise<void> = Promise.resolve();
  private name: string;
  private inst?: runtime.Instance = undefined;
  private globalObjects: Array<Disposable> = [];
  private serverRecvData?: (header: Uint8Array, body: Uint8Array) => void;
  private currPacketHeader?: Uint8Array;
  private currPacketLength = 0;
  private remoteKeyLength = 0;
  private pendingBytes = 0;
  private buffredBytes = 0;
  private messageQueue: Array<Uint8Array> = [];

  constructor(
    url: string,
    key: string,
    getImports: () => Record<string, unknown>,
    logger: (msg: string) => void = console.log,
    tensorCacheUrl = "",
    tensorCacheDevice = "cpu",
    initProgressCallback: runtime.InitProgressCallback | undefined = undefined,
    asyncOnServerLoad: ((inst: runtime.Instance) => Promise<void>) | undefined = undefined,
  ) {
    this.url = url;
    this.key = key;
    this.name = "WebSocketRPCServer[" + this.key + "]: ";
    this.getImports = getImports;
    this.logger = logger;
    this.tensorCacheUrl = tensorCacheUrl;
    this.tensorCacheDevice = tensorCacheDevice;
    this.initProgressCallback = initProgressCallback;
    this.asyncOnServerLoad = asyncOnServerLoad;
    this.checkLittleEndian();
    this.socket = compact.createWebSocket(url);
    this.socket.binaryType = "arraybuffer";

    this.socket.addEventListener("open", (event: Event) => {
      return this.onOpen(event);
    });
    this.socket.addEventListener("message", (event: MessageEvent) => {
      return this.onMessage(event);
    });
    this.socket.addEventListener("close", (event: CloseEvent) => {
      return this.onClose(event);
    });
  }

  private onClose(_event: CloseEvent): void {
    if (this.inst !== undefined) {
      this.globalObjects.forEach(obj => {
        obj.dispose();
      });
      this.log(this.inst.runtimeStatsText());
      this.inst.dispose();
    }
    if (this.state === RPCServerState.ReceivePacketHeader) {
      this.log("Closing the server in clean state");
      this.log("Automatic reconnecting..");
      new RPCServer(
        this.url, this.key, this.getImports, this.logger,
        this.tensorCacheUrl, this.tensorCacheDevice,
        this.initProgressCallback, this.asyncOnServerLoad);
    } else {
      this.log("Closing the server, final state=" + this.state);
    }
  }

  private onOpen(_event: Event): void {
    // Send the headers
    let bkey = StringToUint8Array("server:" + this.key);
    bkey = bkey.slice(0, bkey.length - 1);
    const intbuf = new Int32Array(1);
    intbuf[0] = RPC_MAGIC;
    this.socket.send(intbuf);
    intbuf[0] = bkey.length;
    this.socket.send(intbuf);
    this.socket.send(bkey);
    this.log("connected...");
    // request bytes: magic + keylen
    this.requestBytes(SizeOf.I32 + SizeOf.I32);
    this.state = RPCServerState.InitHeader;
  }

  /** Handler for raw message. */
  private onMessage(event: MessageEvent): void {
    const buffer = event.data;
    this.buffredBytes += buffer.byteLength;
    this.messageQueue.push(new Uint8Array(buffer));
    this.processEvents();
  }
  /** Process ready events. */
  private processEvents(): void {
    while (this.buffredBytes >= this.pendingBytes && this.pendingBytes != 0) {
      this.onDataReady();
    }
  }
  /** State machine to handle each request */
  private onDataReady(): void {
    switch (this.state) {
      case RPCServerState.InitHeader: {
        this.handleInitHeader();
        break;
      }
      case RPCServerState.InitHeaderKey: {
        this.handleInitHeaderKey();
        break;
      }
      case RPCServerState.ReceivePacketHeader: {
        this.currPacketHeader = this.readFromBuffer(SizeOf.I64);
        const reader = new ByteStreamReader(this.currPacketHeader);
        this.currPacketLength = reader.readU64();
        assert(this.pendingBytes === 0);
        this.requestBytes(this.currPacketLength);
        this.state = RPCServerState.ReceivePacketBody;
        break;
      }
      case RPCServerState.ReceivePacketBody: {
        const body = this.readFromBuffer(this.currPacketLength);
        assert(this.pendingBytes === 0);
        assert(this.currPacketHeader !== undefined);
        this.onPacketReady(this.currPacketHeader, body);
        break;
      }
      case RPCServerState.WaitForCallback: {
        assert(this.pendingBytes === 0);
        break;
      }
      default: {
        throw new Error("Cannot handle state " + this.state);
      }
    }
  }

  private onPacketReady(header: Uint8Array, body: Uint8Array): void {
    if (this.inst === undefined) {
      // initialize server.
      const reader = new ByteStreamReader(body);
      const code = reader.readU32();
      const ver = Uint8ArrayToString(reader.readByteArray());
      const nargs = reader.readU32();

      // nargs=0 means no session_constructor_args (LocalSession request).
      // WASM RPC requires ["rpc.WasmSession", wasm_binary]. Wait for proper init.
      if (nargs === 0) {
        this.log("Received LocalSession init (nargs=0), waiting for WasmSession init...");
        this.requestBytes(SizeOf.I64);
        this.state = RPCServerState.ReceivePacketHeader;
        return;
      }

      const args = [];
      for (let i = 0; i < nargs; ++i) {
        const typeIndex = reader.readU32();
        if (typeIndex === TypeIndex.kTVMFFIRawStr) {
          const str = Uint8ArrayToString(reader.readByteArray());
          args.push(str);
        } else if (typeIndex === TypeIndex.kTVMFFIStr) {
          reader.readU32(); // skip duplicate type_index
          const str = Uint8ArrayToString(reader.readByteArray());
          args.push(str);
        } else if (typeIndex === TypeIndex.kTVMFFIByteArrayPtr) {
          args.push(reader.readByteArray());
        } else if (typeIndex === TypeIndex.kTVMFFIBytes) {
          reader.readU32(); // skip duplicate type_index
          args.push(reader.readByteArray());
        } else {
          throw new Error("cannot support type index " + typeIndex);
        }
      }
      this.onInitServer(args, header, body);
    } else {
      assert(this.serverRecvData !== undefined);
      this.serverRecvData(header, body);
      this.requestBytes(SizeOf.I64);
      this.state = RPCServerState.ReceivePacketHeader;
    }
  }

  /** Event handler during server initialization. */
  private onInitServer(
    args: Array<unknown>,
    header: Uint8Array,
    body: Uint8Array
  ): void {
    // start the server
    assert(args[0] === "rpc.WasmSession");
    assert(this.pendingBytes === 0);

    const asyncInitServer = async (): Promise<void> => {
      assert(args[1] instanceof Uint8Array);
      const inst = await runtime.instantiate(
        args[1].buffer as ArrayBuffer,
        this.getImports(),
        this.logger
      );

      try {
        const output: GPUDeviceDetectOutput | undefined = await detectGPUDevice();
        if (output !== undefined) {
          const label = "WebGPU: "+ output.adapterInfo.description;
          this.log("Initialize GPU device: " + label);
          inst.initWebGPU(output.device);
        } else {
          this.log("Cannot find WebGPU device in the env");
        }
      } catch (err) {
        this.log("Cannnot initialize WebGPU, " + err.toString());
      }

      this.inst = inst;
      // begin scope to allow handling of objects
      this.inst.beginScope();
      if (this.initProgressCallback !== undefined) {
        this.inst.registerInitProgressCallback(this.initProgressCallback);
      }

      if (this.tensorCacheUrl.length != 0) {
        if (this.tensorCacheDevice === "cpu") {
          await this.inst.fetchTensorCache(this.tensorCacheUrl, this.inst.cpu());
        } else {
          assert(this.tensorCacheDevice === "webgpu");
          await this.inst.fetchTensorCache(this.tensorCacheUrl, this.inst.webgpu());
        }
      }

      assert(this.inst !== undefined);
      if (this.asyncOnServerLoad !== undefined) {
        await this.asyncOnServerLoad(this.inst);
      }
      const fcreate = this.inst.getGlobalFunc("rpc.CreateEventDrivenServer");
      const messageHandler = fcreate(
        (cbytes: Uint8Array): runtime.Scalar => {
          assert(this.inst !== undefined);
          if (this.socket.readyState === 1) {
            // WebSocket will automatically close the socket
            // if we burst send data that exceeds its internal buffer
            // wait a bit before we send next one.
            const sendDataWithCongestionControl = async (): Promise<void> => {
              const packetSize = 4 << 10;
              const maxBufferAmount = 4 * packetSize;
              const waitTimeMs = 20;
              for (
                let offset = 0;
                offset < cbytes.length;
                offset += packetSize
              ) {
                const end = Math.min(offset + packetSize, cbytes.length);
                while (this.socket.bufferedAmount >= maxBufferAmount) {
                  await new Promise((r) => setTimeout(r, waitTimeMs));
                }
                this.socket.send(cbytes.slice(offset, end));
              }
            };
            // Chain up the pending send so that the async send is always in-order.
            this.pendingSend = this.pendingSend.then(
              sendDataWithCongestionControl
            );
            // Directly return since the data are "sent" from the caller's pov.
            return this.inst.scalar(cbytes.length, "int32");
          } else {
            return this.inst.scalar(0, "int32");
          }
        },
        this.name,
        this.key
      );
      // message handler should persist across RPC runs
      this.globalObjects.push(
        this.inst.detachFromCurrentScope(messageHandler)
      );
      const writeFlag = this.inst.scalar(3, "int32");

      this.serverRecvData = (header: Uint8Array, body: Uint8Array): void => {
        if (messageHandler(header, writeFlag) === 0) {
          this.socket.close();
        }
        if (messageHandler(body, writeFlag) === 0) {
          this.socket.close();
        }
      };

      // Forward the same init sequence to the wasm RPC.
      // The RPC will look for "rpc.wasmSession"
      // and we will redirect it to the correct local session.
      // register the callback to redirect the session to local.
      const flocal = this.inst.getGlobalFunc("wasm.LocalSession");
      const localSession = flocal();
      assert(localSession instanceof runtime.Module);

      this.inst.registerFunc(
        "rpc.WasmSession",
        (_args: unknown): runtime.Module => {
          return localSession;
        }
      );
      messageHandler(header, writeFlag);
      messageHandler(body, writeFlag);

      this.log("Finish initializing the Wasm Server..");
      this.requestBytes(SizeOf.I64);
      this.state = RPCServerState.ReceivePacketHeader;
      // call process events in case there are bufferred data.
      this.processEvents();
      // recycle all values.
      this.inst.endScope();
    };

    this.state = RPCServerState.WaitForCallback;
    asyncInitServer();
  }

  private log(msg: string): void {
    this.logger(this.name + msg);
  }

  private handleInitHeader(): void {
    const reader = new ByteStreamReader(this.readFromBuffer(SizeOf.I32 * 2));
    const magic = reader.readU32();
    if (magic === RPC_MAGIC + 1) {
      throw new Error("key: " + this.key + " has already been used in proxy");
    } else if (magic === RPC_MAGIC + 2) {
      throw new Error("RPCProxy do not have matching client key " + this.key);
    }
    assert(magic === RPC_MAGIC, this.url + " is not an RPC Proxy");
    this.remoteKeyLength = reader.readU32();
    assert(this.pendingBytes === 0);
    this.requestBytes(this.remoteKeyLength);
    this.state = RPCServerState.InitHeaderKey;
  }

  private handleInitHeaderKey(): void {
    const remoteKey = Uint8ArrayToString(
      this.readFromBuffer(this.remoteKeyLength)
    );
    assert(this.pendingBytes === 0);
    this.requestBytes(SizeOf.I64);
    this.state = RPCServerState.ReceivePacketHeader;
  }

  private checkLittleEndian(): void {
    const a = new ArrayBuffer(4);
    const b = new Uint8Array(a);
    const c = new Uint32Array(a);
    b[0] = 0x11;
    b[1] = 0x22;
    b[2] = 0x33;
    b[3] = 0x44;
    assert(c[0] === 0x44332211, "RPCServer little endian to work");
  }

  private requestBytes(nbytes: number): void {
    this.pendingBytes += nbytes;
  }

  private readFromBuffer(nbytes: number): Uint8Array {
    const ret = new Uint8Array(nbytes);
    let ptr = 0;
    while (ptr < nbytes) {
      assert(this.messageQueue.length != 0);
      const nleft = nbytes - ptr;
      if (this.messageQueue[0].byteLength <= nleft) {
        const buffer = this.messageQueue.shift() as Uint8Array;
        ret.set(buffer, ptr);
        ptr += buffer.byteLength;
      } else {
        const buffer = this.messageQueue[0];
        ret.set(buffer.slice(0, nleft), ptr);
        this.messageQueue[0] = buffer.slice(nleft, buffer.byteLength);
        ptr += nleft;
      }
    }
    this.buffredBytes -= nbytes;
    this.pendingBytes -= nbytes;
    return ret;
  }
}