Triton Inference Serverを​色々な​環境から​使いたい​: TypeScript (Node.js) 編

# https://github.com/tuxedocat/triton-client-polyglot-example/blob/c0cb5e2e76170d8f8c3599f9c9c0d4d2acb011ea/client/typescript/package.json#L15
git clone https://github.com/triton-inference-server/common -b r22.02 ./upstream
mkdir proto && mv ./upstream/protobuf/*.proto ./proto/

// package.json
// https://github.com/tuxedocat/triton-client-polyglot-example/blob/c0cb5e2e76170d8f8c3599f9c9c0d4d2acb011ea/client/typescript/package.json#L29-L44
  "dependencies": {
    "@grpc/grpc-js": "^1.6.2",
    "@grpc/proto-loader": "^0.6.9"
  },
  "devDependencies": {
    "@types/node": "^14.0",
    "grpc-tools": "^1.11.0",
    "ts-node": "*",
    "typescript": "^4.3.5",
  },
  ...

./node_modules/bin/proto-loader-gen-types \
  --longs=String \
  --enums=String \
  --defaults \
  --oneofs \
  --keepCase \
  --grpcLib=@grpc/grpc-js \
  --outDir=./generated \
  proto/*.proto"

generated/
├── inference
│   ├── たくさん...
├── grpc_service.ts
└── model_config.ts

import { loadPackageDefinition } from '@grpc/grpc-js'
import { loadSync } from '@grpc/proto-loader'
import { ProtoGrpcType } from './generated/grpc_service'

const PROTO_IMPORT_PATH = __dirname + '/proto'
const PROTO_PATH = PROTO_IMPORT_PATH + '/grpc_service.proto'
const packageDefinition = loadSync(PROTO_PATH, {
  includeDirs: [PROTO_IMPORT_PATH],
  keepCase: true,
  longs: String,
  enums: String,
  defaults: true,
  oneofs: true,
})
const loadedPackageDefinition = loadPackageDefinition(packageDefinition) as unknown as ProtoGrpcType
export const triton = loadedPackageDefinition.inference

// https://github.com/tuxedocat/triton-client-polyglot-example/blob/c0cb5e2e76170d8f8c3599f9c9c0d4d2acb011ea/client/typescript/src/index.ts#L39-L77
/**
 * Promisified API subset of triton.GRPCInferenceServiceClient
 * Notes: Only unary ops are supported e.g. no `modelStreamInfer`, due to limitation of util.promisify.
 * */
export class TritonCoreAPI {
  private client: GRPCInferenceServiceClient
  constructor(endpoint: string, creds?: ChannelCredentials) {
    this.client = new triton.GRPCInferenceService(endpoint, creds ?? credentials.createInsecure())
  }
  public async serverReady(req: ServerReadyRequest): Promise<ServerReadyResponse> {
    return promisify<ServerReadyRequest>(this.client.serverReady.bind(this.client))(
      req
    ) as unknown as ServerReadyResponse
  }
  public async serverLive(req: ServerLiveRequest): Promise<ServerLiveResponse> {
    return promisify<ServerLiveRequest>(this.client.serverLive.bind(this.client))(
      req
    ) as unknown as ServerLiveResponse
  }
  public async modelReady(req: ModelReadyRequest): Promise<ModelReadyResponse> {
    return promisify<ModelReadyRequest>(this.client.modelReady.bind(this.client))(
      req
    ) as unknown as ModelReadyResponse
  }
  public async modelConfig(req: ModelConfigRequest): Promise<ModelConfigResponse> {
    return promisify<ModelConfigRequest>(this.client.modelConfig.bind(this.client))(
      req
    ) as unknown as ModelConfigResponse
  }
  public async modelMetadata(req: ModelMetadataRequest): Promise<ModelMetadataResponse> {
    return promisify<ModelMetadataRequest>(this.client.modelMetadata.bind(this.client))(
      req
    ) as unknown as ModelMetadataResponse
  }
  public async modelStatistics(req: ModelStatisticsRequest): Promise<ModelStatisticsResponse> {
    return promisify<ModelStatisticsRequest>(this.client.modelStatistics.bind(this.client))(
      req
    ) as unknown as ModelStatisticsResponse
  }
  public async modelInfer(req: ModelInferRequest): Promise<ModelInferResponse> {
    return promisify<ModelInferRequest>(this.client.modelInfer.bind(this.client))(
      req
    ) as unknown as ModelInferResponse
  }
}

test('client can call liveness endpoint', async () => {
  tritonClient = new TritonCoreAPI(ENDPOINT, credentials.createInsecure())
  const serverLiveResponse = (await tritonClient.serverLive({})) as ServerLiveResponse
  expect(serverLiveResponse).toBeDefined()
  expect(serverLiveResponse.live).toBe(true)
})

// Utility
const toUTF8Buffer = (s: string): Buffer => {
  return Buffer.from(new TextEncoder().encode(s).buffer)
}

// Tensor
const inputText = {
  name: 'TEXT',
  datatype: 'BYTES',
  shape: [1, 1],
  contents: {
    bytes_contents: [toUTF8Buffer(someString)],
  },
} as InferTensorContents

// Request
const inferRequest = {
  model_name: 'model_name',
  inputs: [inputText],
  outputs: [{ name: 'OUTPUT_NAME' }],
} as ModelInferRequest

import { Sharp } from 'sharp'
import {
  ModelInferRequest,
  _inference_ModelInferRequest_InferInputTensor as InferInputTensor,
} from './generated/inference/ModelInferRequest'
import { InferTensorContents } from './generated/inference/InferTensorContents'

/**
 * Convert image buffer to unsigned int array
 */
fromImageBuffer(imgBuf: ArrayBufferLike): number[] {
    return Array.from(new Uint8ClampedArray(imgBuf))
}

// Tensor
const { data, info } = await img.removeAlpha().raw().toBuffer({ resolveWithObject: true })
const imgBuf = data.buffer
const inputImage = {
    name,
    datatype: 'UINT8',
    shape: [1, info.height, info.width, info.channels],
    contents: {
      uint_contents: fromImageBuffer(imgBuf),
    } as InferTensorContents,
} as InferInputTensor

// Request
const inferRequest = {
    model_name: 'image_model',
    inputs: [inputImage],
    outputs: [{ name: 'OUTPUT' }],
} as ModelInferRequest

interface TritonInferenceOutput {
  name: string
  value: number
}
getFloats(response: ModelInferResponse): TritonInferenceOutput[] {
    const outputDefs = response.outputs as _inference_ModelInferResponse_InferOutputTensor[]
    const outputBuffer = response.raw_output_contents as Buffer[]
    const parsedFp32 = outputDefs.map((e, i) => {
        return { name: e.name ?? '', value: outputBuffer[i].readFloatLE(0) }
    })
    return parsedFp32 ?? []
}
const outputs = getFloats((await tritonClient.modelInfer(inferRequest)) as ModelInferResponse)