Skip to content

GitLab

Unverified Commit b39850f9 authored by liuzhe-lz's avatar liuzhe-lz Committed by GitHub
Browse files

WebSocket (step 2) - TS server (#4808)

parent 05c7d6e9
Hide whitespace changes
Inline Side-by-side
Showing with 516 additions and 159 deletions
ts/nni_manager/core/ipcInterface.ts
View file @ b39850f9
// Copyright (c) Microsoft Corporation. // Copyright (c) Microsoft Corporation.
// Licensed under the MIT license. // Licensed under the MIT license.
import assert from 'assert'; export { IpcInterface } from './tuner_command_channel/common';
import { ChildProcess } from 'child_process'; export { createDispatcherInterface, createDispatcherPipeInterface, encodeCommand } from './tuner_command_channel/legacy';
import { EventEmitter } from 'events';
import net from 'net';
import { Readable, Writable } from 'stream';
import { NNIError } from '../common/errors';
import { getLogger, Logger } from '../common/log';
import { getLogDir } from '../common/utils';
import * as CommandType from './commands';
const ipcOutgoingFd: number = 3;
const ipcIncomingFd: number = 4;
/**
* Encode a command
* @param commandType a command type defined in 'core/commands'
* @param content payload of the command
* @returns binary command data
*/
function encodeCommand(commandType: string, content: string): Buffer {
const contentBuffer: Buffer = Buffer.from(content);
const contentLengthBuffer: Buffer = Buffer.from(contentBuffer.length.toString().padStart(14, '0'));
return Buffer.concat([Buffer.from(commandType), contentLengthBuffer, contentBuffer]);
}
/**
* Decode a command
* @param Buffer binary incoming data
* @returns a tuple of (success, commandType, content, remain)
* success: true if the buffer contains at least one complete command; otherwise false
* remain: remaining data after the first command
*/
function decodeCommand(data: Buffer): [boolean, string, string, Buffer] {
if (data.length < 8) {
return [false, '', '', data];
}
const commandType: string = data.slice(0, 2).toString();
const contentLength: number = parseInt(data.slice(2, 16).toString(), 10);
if (data.length < contentLength + 16) {
return [false, '', '', data];
}
const content: string = data.slice(16, contentLength + 16).toString();
const remain: Buffer = data.slice(contentLength + 16);
return [true, commandType, content, remain];
}
class IpcInterface {
private acceptCommandTypes: Set<string>;
private outgoingStream: Writable;
private incomingStream: Readable;
private eventEmitter: EventEmitter;
private readBuffer: Buffer;
private logger: Logger = getLogger('IpcInterface');
/**
* Construct a IPC proxy
* @param proc the process to wrap
* @param acceptCommandTypes set of accepted commands for this process
*/
constructor(outStream: Writable, inStream: Readable, acceptCommandTypes: Set<string>) {
this.acceptCommandTypes = acceptCommandTypes;
this.outgoingStream = outStream;
this.incomingStream = inStream;
this.eventEmitter = new EventEmitter();
this.readBuffer = Buffer.alloc(0);
this.incomingStream.on('data', (data: Buffer) => { this.receive(data); });
this.incomingStream.on('error', (error: Error) => { this.eventEmitter.emit('error', error); });
this.outgoingStream.on('error', (error: Error) => { this.eventEmitter.emit('error', error); });
}
/**
* Send a command to process
* @param commandType: a command type defined in 'core/commands'
* @param content: payload of command
*/
public sendCommand(commandType: string, content: string = ''): void {
this.logger.debug(`ipcInterface command type: [${commandType}], content:[${content}]`);
assert.ok(this.acceptCommandTypes.has(commandType));
try {
const data: Buffer = encodeCommand(commandType, content);
if (!this.outgoingStream.write(data)) {
this.logger.warning('Commands jammed in buffer!');
}
} catch (err) {
throw NNIError.FromError(
err,
`Dispatcher Error, please check this dispatcher log file for more detailed information: ${getLogDir()}/dispatcher.log . `
);
}
}
/**
* Add a command listener
* @param listener the listener callback
*/
public onCommand(listener: (commandType: string, content: string) => void): void {
this.eventEmitter.on('command', listener);
}
public onError(listener: (error: Error) => void): void {
this.eventEmitter.on('error', listener);
}
/**
* Deal with incoming data from process
* Invoke listeners for each complete command received, save incomplete command to buffer
* @param data binary incoming data
*/
private receive(data: Buffer): void {
this.readBuffer = Buffer.concat([this.readBuffer, data]);
while (this.readBuffer.length > 0) {
const [success, commandType, content, remain] = decodeCommand(this.readBuffer);
if (!success) {
break;
}
assert.ok(this.acceptCommandTypes.has(commandType));
this.eventEmitter.emit('command', commandType, content);
this.readBuffer = remain;
}
}
}
/**
* Create IPC proxy for tuner process
* @param process_ the tuner process
*/
function createDispatcherInterface(process: ChildProcess): IpcInterface {
const outStream = <Writable>process.stdio[ipcOutgoingFd];
const inStream = <Readable>process.stdio[ipcIncomingFd];
return new IpcInterface(outStream, inStream, new Set([...CommandType.TUNER_COMMANDS, ...CommandType.ASSESSOR_COMMANDS]));
}
function createDispatcherPipeInterface(pipePath: string): IpcInterface {
const client = net.createConnection(pipePath);
return new IpcInterface(client, client, new Set([...CommandType.TUNER_COMMANDS, ...CommandType.ASSESSOR_COMMANDS]));
}
export { IpcInterface, createDispatcherInterface, createDispatcherPipeInterface, encodeCommand, decodeCommand };
ts/nni_manager/core/nnimanager.ts
View file @ b39850f9
...@@ -188,7 +188,7 @@ class NNIManager implements Manager { ...@@ -188,7 +188,7 @@ class NNIManager implements Manager {
const dispatcherCommand: string = getMsgDispatcherCommand(config); const dispatcherCommand: string = getMsgDispatcherCommand(config);
this.log.debug(`dispatcher command: ${dispatcherCommand}`); this.log.debug(`dispatcher command: ${dispatcherCommand}`);
const checkpointDir: string = await this.createCheckpointDir(); const checkpointDir: string = await this.createCheckpointDir();
this.setupTuner(dispatcherCommand, undefined, 'start', checkpointDir); await this.setupTuner(dispatcherCommand, undefined, 'start', checkpointDir);
this.setStatus('RUNNING'); this.setStatus('RUNNING');
await this.storeExperimentProfile(); await this.storeExperimentProfile();
this.run().catch((err: Error) => { this.run().catch((err: Error) => {
...@@ -221,7 +221,7 @@ class NNIManager implements Manager { ...@@ -221,7 +221,7 @@ class NNIManager implements Manager {
const dispatcherCommand: string = getMsgDispatcherCommand(config); const dispatcherCommand: string = getMsgDispatcherCommand(config);
this.log.debug(`dispatcher command: ${dispatcherCommand}`); this.log.debug(`dispatcher command: ${dispatcherCommand}`);
const checkpointDir: string = await this.createCheckpointDir(); const checkpointDir: string = await this.createCheckpointDir();
this.setupTuner(dispatcherCommand, undefined, 'resume', checkpointDir); await this.setupTuner(dispatcherCommand, undefined, 'resume', checkpointDir);
const allTrialJobs: TrialJobInfo[] = await this.dataStore.listTrialJobs(); const allTrialJobs: TrialJobInfo[] = await this.dataStore.listTrialJobs();
...@@ -462,7 +462,7 @@ class NNIManager implements Manager { ...@@ -462,7 +462,7 @@ class NNIManager implements Manager {
} }
} }
private setupTuner(command: string, cwd: string | undefined, mode: 'start' | 'resume', dataDirectory: string): void { private async setupTuner(command: string, cwd: string | undefined, mode: 'start' | 'resume', dataDirectory: string): Promise<void> {
if (this.dispatcher !== undefined) { if (this.dispatcher !== undefined) {
return; return;
} }
...@@ -488,7 +488,7 @@ class NNIManager implements Manager { ...@@ -488,7 +488,7 @@ class NNIManager implements Manager {
const newEnv = Object.assign({}, process.env, nniEnv); const newEnv = Object.assign({}, process.env, nniEnv);
const tunerProc: ChildProcess = getTunerProc(command, stdio, newCwd, newEnv); const tunerProc: ChildProcess = getTunerProc(command, stdio, newCwd, newEnv);
this.dispatcherPid = tunerProc.pid!; this.dispatcherPid = tunerProc.pid!;
this.dispatcher = createDispatcherInterface(tunerProc); this.dispatcher = await createDispatcherInterface(tunerProc);
return; return;
} }
......
ts/nni_manager/core/tuner_command_channel/common.ts 0 → 100644
View file @ b39850f9
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
export interface IpcInterface {
sendCommand(commandType: string, content?: string): void;
onCommand(listener: (commandType: string, content: string) => void): void;
onError(listener: (error: Error) => void): void;
}
ts/nni_manager/core/tuner_command_channel/index.ts 0 → 100644
View file @ b39850f9
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
export { getWebSocketChannel, serveWebSocket } from './websocket_channel';
ts/nni_manager/core/tuner_command_channel/legacy.ts 0 → 100644
View file @ b39850f9
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
import assert from 'assert';
import { ChildProcess } from 'child_process';
import { EventEmitter } from 'events';
import net from 'net';
import { Readable, Writable } from 'stream';
import { NNIError } from '../../common/errors';
import { getLogger, Logger } from '../../common/log';
import { getLogDir } from '../../common/utils';
import * as CommandType from '../commands';
import type { IpcInterface } from './common';
const ipcOutgoingFd: number = 3;
const ipcIncomingFd: number = 4;
/**
* Encode a command
* @param commandType a command type defined in 'core/commands'
* @param content payload of the command
* @returns binary command data
*/
function encodeCommand(commandType: string, content: string): Buffer {
const contentBuffer: Buffer = Buffer.from(content);
const contentLengthBuffer: Buffer = Buffer.from(contentBuffer.length.toString().padStart(14, '0'));
return Buffer.concat([Buffer.from(commandType), contentLengthBuffer, contentBuffer]);
}
/**
* Decode a command
* @param Buffer binary incoming data
* @returns a tuple of (success, commandType, content, remain)
* success: true if the buffer contains at least one complete command; otherwise false
* remain: remaining data after the first command
*/
function decodeCommand(data: Buffer): [boolean, string, string, Buffer] {
if (data.length < 8) {
return [false, '', '', data];
}
const commandType: string = data.slice(0, 2).toString();
const contentLength: number = parseInt(data.slice(2, 16).toString(), 10);
if (data.length < contentLength + 16) {
return [false, '', '', data];
}
const content: string = data.slice(16, contentLength + 16).toString();
const remain: Buffer = data.slice(contentLength + 16);
return [true, commandType, content, remain];
}
class LegacyIpcInterface implements IpcInterface {
private acceptCommandTypes: Set<string>;
private outgoingStream: Writable;
private incomingStream: Readable;
private eventEmitter: EventEmitter;
private readBuffer: Buffer;
private logger: Logger = getLogger('IpcInterface');
/**
* Construct a IPC proxy
* @param proc the process to wrap
* @param acceptCommandTypes set of accepted commands for this process
*/
constructor(outStream: Writable, inStream: Readable, acceptCommandTypes: Set<string>) {
this.acceptCommandTypes = acceptCommandTypes;
this.outgoingStream = outStream;
this.incomingStream = inStream;
this.eventEmitter = new EventEmitter();
this.readBuffer = Buffer.alloc(0);
this.incomingStream.on('data', (data: Buffer) => { this.receive(data); });
this.incomingStream.on('error', (error: Error) => { this.eventEmitter.emit('error', error); });
this.outgoingStream.on('error', (error: Error) => { this.eventEmitter.emit('error', error); });
}
/**
* Send a command to process
* @param commandType: a command type defined in 'core/commands'
* @param content: payload of command
*/
public sendCommand(commandType: string, content: string = ''): void {
this.logger.debug(`ipcInterface command type: [${commandType}], content:[${content}]`);
assert.ok(this.acceptCommandTypes.has(commandType));
try {
const data: Buffer = encodeCommand(commandType, content);
if (!this.outgoingStream.write(data)) {
this.logger.warning('Commands jammed in buffer!');
}
} catch (err) {
throw NNIError.FromError(
err,
`Dispatcher Error, please check this dispatcher log file for more detailed information: ${getLogDir()}/dispatcher.log . `
);
}
}
/**
* Add a command listener
* @param listener the listener callback
*/
public onCommand(listener: (commandType: string, content: string) => void): void {
this.eventEmitter.on('command', listener);
}
public onError(listener: (error: Error) => void): void {
this.eventEmitter.on('error', listener);
}
/**
* Deal with incoming data from process
* Invoke listeners for each complete command received, save incomplete command to buffer
* @param data binary incoming data
*/
private receive(data: Buffer): void {
this.readBuffer = Buffer.concat([this.readBuffer, data]);
while (this.readBuffer.length > 0) {
const [success, commandType, content, remain] = decodeCommand(this.readBuffer);
if (!success) {
break;
}
assert.ok(this.acceptCommandTypes.has(commandType));
this.eventEmitter.emit('command', commandType, content);
this.readBuffer = remain;
}
}
}
/**
* Create IPC proxy for tuner process
* @param process_ the tuner process
*/
async function createDispatcherInterface(process: ChildProcess): Promise<IpcInterface> {
const outStream = <Writable>process.stdio[ipcOutgoingFd];
const inStream = <Readable>process.stdio[ipcIncomingFd];
return new LegacyIpcInterface(outStream, inStream, new Set([...CommandType.TUNER_COMMANDS, ...CommandType.ASSESSOR_COMMANDS]));
}
function createDispatcherPipeInterface(pipePath: string): IpcInterface {
const client = net.createConnection(pipePath);
return new LegacyIpcInterface(client, client, new Set([...CommandType.TUNER_COMMANDS, ...CommandType.ASSESSOR_COMMANDS]));
}
export { createDispatcherInterface, createDispatcherPipeInterface, encodeCommand, decodeCommand };
ts/nni_manager/core/tuner_command_channel/shim.ts 0 → 100644
View file @ b39850f9
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
import type { IpcInterface } from './common';
import { WebSocketChannel, getWebSocketChannel } from './websocket_channel';
export async function createDispatcherInterface(): Promise<IpcInterface> {
const ipcInterface = new WsIpcInterface();
await ipcInterface.init();
return ipcInterface;
}
class WsIpcInterface implements IpcInterface {
private channel: WebSocketChannel = getWebSocketChannel();
public async init(): Promise<void> {
await this.channel.init();
}
public sendCommand(commandType: string, content: string = ''): void {
if (commandType !== 'PI') { // ping is handled with WebSocket protocol
this.channel.sendCommand(commandType + content);
if (commandType === 'TE') {
this.channel.shutdown();
}
}
}
public onCommand(listener: (commandType: string, content: string) => void): void {
this.channel.onCommand((command: string) => {
listener(command.slice(0, 2), command.slice(2));
});
}
public onError(listener: (error: Error) => void): void {
this.channel.onError(listener);
}
}
ts/nni_manager/core/tuner_command_channel/websocket_channel.ts 0 → 100644
View file @ b39850f9
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
/**
* The IPC channel between NNI manager and tuner.
*
* TODO:
* 1. Merge with environment service's WebSocket channel.
* 2. Split import data command to avoid extremely long message.
* 3. Refactor message format.
**/
import assert from 'assert/strict';
import { EventEmitter } from 'events';
import { Deferred } from 'ts-deferred';
import type WebSocket from 'ws';
import { Logger, getLogger } from 'common/log';
const logger: Logger = getLogger('tuner_command_channel.WebSocketChannel');
export interface WebSocketChannel {
init(): Promise<void>;
shutdown(): Promise<void>;
sendCommand(command: string): void; // maybe this should return Promise<void>
onCommand(callback: (command: string) => void): void;
onError(callback: (error: Error) => void): void;
}
/**
* Get the singleton tuner command channel.
* Remember to invoke ``await channel.init()`` before doing anything else.
**/
export function getWebSocketChannel(): WebSocketChannel {
return channelSingleton;
}
/**
* The callback to serve WebSocket connection request. Used by REST server module.
* It should only be invoked once, or an error will be raised.
*
* Type hint of express-ws is somewhat problematic. Don't want to waste time on it so use `any`.
**/
export function serveWebSocket(ws: WebSocket): void {
channelSingleton.setWebSocket(ws);
}
class WebSocketChannelImpl implements WebSocketChannel {
private deferredInit: Deferred<void> = new Deferred<void>();
private emitter: EventEmitter = new EventEmitter();
private heartbeatTimer!: NodeJS.Timer;
private serving: boolean = false;
private waitingPong: boolean = false;
private ws!: WebSocket;
public setWebSocket(ws: WebSocket): void {
if (this.ws !== undefined) {
logger.error('A second client is trying to connect');
ws.close(4030, 'Already serving a tuner.');
return;
}
logger.debug('Connected.');
this.serving = true;
this.ws = ws;
ws.on('close', () => { this.handleError(new Error('tuner_command_channel: Tuner closed connection')); });
ws.on('error', this.handleError.bind(this));
ws.on('message', this.receive.bind(this));
ws.on('pong', () => { this.waitingPong = false; });
this.heartbeatTimer = setInterval(this.heartbeat.bind(this), heartbeatInterval);
this.deferredInit.resolve();
}
public init(): Promise<void> {
logger.debug(this.ws === undefined ? 'Waiting connection...' : 'Initialized.');
return this.deferredInit.promise;
}
public async shutdown(): Promise<void> {
if (this.ws === undefined) {
return;
}
clearInterval(this.heartbeatTimer);
this.serving = false;
this.emitter.removeAllListeners();
}
public sendCommand(command: string): void {
assert.ok(this.ws !== undefined);
logger.debug('Sending', command);
this.ws.send(command);
if (this.ws.bufferedAmount > command.length + 1000) {
logger.warning('Sending too fast! Try to reduce the frequency of intermediate results.');
}
}
public onCommand(callback: (command: string) => void): void {
this.emitter.on('command', callback);
}
public onError(callback: (error: Error) => void): void {
this.emitter.on('error', callback);
}
private heartbeat(): void {
if (this.waitingPong) {
this.ws.terminate(); // this will trigger "close" event
this.handleError(new Error('tuner_command_channel: Tuner loses responsive'));
}
this.waitingPong = true;
this.ws.ping();
}
private receive(data: Buffer, _isBinary: boolean): void {
logger.debug('Received', data);
this.emitter.emit('command', data.toString());
}
private handleError(error: Error): void {
if (!this.serving) {
logger.debug('Silent error:', error);
return;
}
logger.error('Error:', error);
clearInterval(this.heartbeatTimer);
this.emitter.emit('error', error);
this.serving = false;
}
}
const channelSingleton: WebSocketChannelImpl = new WebSocketChannelImpl();
let heartbeatInterval: number = 5000;
export namespace UnitTestHelpers {
export function setHeartbeatInterval(ms: number): void {
heartbeatInterval = ms;
}
}
ts/nni_manager/package.json
View file @ b39850f9
...@@ -16,6 +16,7 @@ ...@@ -16,6 +16,7 @@
"child-process-promise": "^2.2.1", "child-process-promise": "^2.2.1",
"express": "^4.17.2", "express": "^4.17.2",
"express-joi-validator": "^2.0.1", "express-joi-validator": "^2.0.1",
"express-ws": "^5.0.2",
"http-proxy": "^1.18.1", "http-proxy": "^1.18.1",
"ignore": "^5.1.8", "ignore": "^5.1.8",
"js-base64": "^3.6.1", "js-base64": "^3.6.1",
...@@ -32,13 +33,14 @@ ...@@ -32,13 +33,14 @@
"ts-deferred": "^1.0.4", "ts-deferred": "^1.0.4",
"typescript-ioc": "^1.2.6", "typescript-ioc": "^1.2.6",
"typescript-string-operations": "^1.4.1", "typescript-string-operations": "^1.4.1",
"ws": "^7.4.6", "ws": "^8.5.0",
"yargs": "^17.3.1" "yargs": "^17.3.1"
}, },
"devDependencies": { "devDependencies": {
"@types/chai": "^4.2.18", "@types/chai": "^4.2.18",
"@types/chai-as-promised": "^7.1.0", "@types/chai-as-promised": "^7.1.0",
"@types/express": "^4.17.2", "@types/express": "^4.17.2",
"@types/express-ws": "^3.0.1",
"@types/glob": "^7.1.3", "@types/glob": "^7.1.3",
"@types/http-proxy": "^1.17.7", "@types/http-proxy": "^1.17.7",
"@types/js-base64": "^3.3.1", "@types/js-base64": "^3.3.1",
...@@ -54,7 +56,7 @@ ...@@ -54,7 +56,7 @@
"@types/stream-buffers": "^3.0.3", "@types/stream-buffers": "^3.0.3",
"@types/tar": "^4.0.4", "@types/tar": "^4.0.4",
"@types/tmp": "^0.2.0", "@types/tmp": "^0.2.0",
"@types/ws": "^7.4.4", "@types/ws": "^8.5.3",
"@types/yargs": "^17.0.8", "@types/yargs": "^17.0.8",
"@typescript-eslint/eslint-plugin": "^2.10.0", "@typescript-eslint/eslint-plugin": "^2.10.0",
"@typescript-eslint/parser": "^4.26.0", "@typescript-eslint/parser": "^4.26.0",
......
ts/nni_manager/rest_server/index.ts
View file @ b39850f9
...@@ -26,11 +26,13 @@ import type { AddressInfo } from 'net'; ...@@ -26,11 +26,13 @@ import type { AddressInfo } from 'net';
import path from 'path'; import path from 'path';
import express, { Request, Response, Router } from 'express'; import express, { Request, Response, Router } from 'express';
import expressWs from 'express-ws';
import httpProxy from 'http-proxy'; import httpProxy from 'http-proxy';
import { Deferred } from 'ts-deferred'; import { Deferred } from 'ts-deferred';
import globals from 'common/globals'; import globals from 'common/globals';
import { Logger, getLogger } from 'common/log'; import { Logger, getLogger } from 'common/log';
import * as tunerCommandChannel from 'core/tuner_command_channel';
import { createRestHandler } from './restHandler'; import { createRestHandler } from './restHandler';
const logger: Logger = getLogger('RestServer'); const logger: Logger = getLogger('RestServer');
...@@ -60,6 +62,8 @@ export class RestServer { ...@@ -60,6 +62,8 @@ export class RestServer {
logger.info(`Starting REST server at port ${this.port}, URL prefix: "/${this.urlPrefix}"`); logger.info(`Starting REST server at port ${this.port}, URL prefix: "/${this.urlPrefix}"`);
const app = express(); const app = express();
expressWs(app, undefined, { wsOptions: { maxPayload: 4 * 1024 * 1024 * 1024 }});
app.use('/' + this.urlPrefix, rootRouter()); app.use('/' + this.urlPrefix, rootRouter());
app.all('*', (_req: Request, res: Response) => { res.status(404).send(`Outside prefix "/${this.urlPrefix}"`); }); app.all('*', (_req: Request, res: Response) => { res.status(404).send(`Outside prefix "/${this.urlPrefix}"`); });
this.server = app.listen(this.port); this.server = app.listen(this.port);
...@@ -100,12 +104,15 @@ export class RestServer { ...@@ -100,12 +104,15 @@ export class RestServer {
* In fact experiments management should have a separate prefix and module. * In fact experiments management should have a separate prefix and module.
**/ **/
function rootRouter(): Router { function rootRouter(): Router {
const router = Router(); const router = Router() as expressWs.Router;
router.use(express.json({ limit: '50mb' })); router.use(express.json({ limit: '50mb' }));
/* NNI manager APIs */ /* NNI manager APIs */
router.use('/api/v1/nni', restHandlerFactory()); router.use('/api/v1/nni', restHandlerFactory());
/* WebSocket APIs */
router.ws('/tuner', (ws, _req, _next) => { tunerCommandChannel.serveWebSocket(ws); });
/* Download log files */ /* Download log files */
// The REST API path "/logs" does not match file system path "/log". // The REST API path "/logs" does not match file system path "/log".
// Here we use an additional router to workaround this problem. // Here we use an additional router to workaround this problem.
......
ts/nni_manager/test/core/ipcInterface.test.ts
View file @ b39850f9
...@@ -16,7 +16,7 @@ const receivedCommands: { [key: string]: string }[] = []; ...@@ -16,7 +16,7 @@ const receivedCommands: { [key: string]: string }[] = [];
let rejectCommandType: Error | undefined; let rejectCommandType: Error | undefined;
function runProcess(): Promise<Error | null> { async function runProcess(): Promise<Error | null> {
// the process is intended to throw error, do not reject // the process is intended to throw error, do not reject
const deferred: Deferred<Error | null> = new Deferred<Error | null>(); const deferred: Deferred<Error | null> = new Deferred<Error | null>();
...@@ -42,7 +42,7 @@ function runProcess(): Promise<Error | null> { ...@@ -42,7 +42,7 @@ function runProcess(): Promise<Error | null> {
}); });
// create IPC interface // create IPC interface
const dispatcher: IpcInterface = createDispatcherInterface(proc); const dispatcher: IpcInterface = await createDispatcherInterface(proc);
dispatcher.onCommand((commandType: string, content: string): void => { dispatcher.onCommand((commandType: string, content: string): void => {
receivedCommands.push({ commandType, content }); receivedCommands.push({ commandType, content });
}); });
......
ts/nni_manager/test/core/ipcInterfaceTerminate.test.ts
View file @ b39850f9
...@@ -14,7 +14,7 @@ let dispatcher: IpcInterface | undefined; ...@@ -14,7 +14,7 @@ let dispatcher: IpcInterface | undefined;
let procExit: boolean = false; let procExit: boolean = false;
let procError: boolean = false; let procError: boolean = false;
function startProcess(): void { async function startProcess(): Promise<void> {
// create fake assessor process // create fake assessor process
const stdio: StdioOptions = ['ignore', 'pipe', process.stderr, 'pipe', 'pipe']; const stdio: StdioOptions = ['ignore', 'pipe', process.stderr, 'pipe', 'pipe'];
...@@ -53,7 +53,7 @@ function startProcess(): void { ...@@ -53,7 +53,7 @@ function startProcess(): void {
}); });
// create IPC interface // create IPC interface
dispatcher = createDispatcherInterface(proc); dispatcher = await createDispatcherInterface(proc);
(<IpcInterface>dispatcher).onCommand((commandType: string, content: string): void => { (<IpcInterface>dispatcher).onCommand((commandType: string, content: string): void => {
console.log(commandType, content); console.log(commandType, content);
}); });
......
ts/nni_manager/test/core/tuner_command_channel.test.ts 0 → 100644
View file @ b39850f9
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
import assert from 'assert/strict';
import { setTimeout } from 'timers/promises';
import WebSocket from 'ws';
import { getWebSocketChannel, serveWebSocket } from 'core/tuner_command_channel';
import { UnitTestHelpers } from 'core/tuner_command_channel/websocket_channel';
UnitTestHelpers.setHeartbeatInterval(10); // for testError, must be set before serveWebSocket()
/* test cases */
// Start serving and let a client connect.
async function testInit(): Promise<void> {
server.on('connection', serveWebSocket);
startClient();
await getWebSocketChannel().init();
}
// Send commands from server to client.
async function testSend(): Promise<void> {
const channel = getWebSocketChannel();
channel.sendCommand(command1);
channel.sendCommand(command2);
await setTimeout(10);
assert.equal(clientReceived.length, 2);
assert.equal(clientReceived[0], command1);
assert.equal(clientReceived[1], command2);
}
// Send commands from client to server.
async function testReceive(): Promise<void> {
const channel = getWebSocketChannel();
channel.onCommand(command => { serverReceived.push(command); });
client.send(command1);
client.send(command2);
await setTimeout(10);
assert.equal(serverReceived.length, 2);
assert.deepEqual(serverReceived[0], command1);
assert.deepEqual(serverReceived[1], command2);
}
// Simulate client side crash.
async function testError(): Promise<void> {
const channel = getWebSocketChannel();
if (process.platform === 'darwin') {
// macOS does not raise the error in 30ms
// not a big problem and don't want to debug. ignore it.
channel.shutdown();
return;
}
channel.onError(error => { catchedError = error; });
// we have set heartbeat interval to 10ms, so pause for 30ms should make it timeout
client.pause();
await setTimeout(30);
assert.notEqual(catchedError, undefined);
client.resume();
}
// Clean up.
async function testShutdown(): Promise<void> {
const channel = getWebSocketChannel();
await channel.shutdown();
client.close();
server.close();
}
/* register */
describe('## tuner_command_channel ##', () => {
it('init', testInit);
it('send', testSend);
it('receive', testReceive);
it('catch error', testError);
it('shutdown', testShutdown);
});
/** helpers **/
const command1 = 'T_hello world';
const command2 = 'T_你好';
const commandPing = 'PI';
const server = new WebSocket.Server({ port: 0 });
let client!: WebSocket;
const serverReceived: string[] = [];
const clientReceived: string[] = [];
let catchedError: Error | undefined;
function startClient() {
const port = (server.address() as any).port;
client = new WebSocket(`ws://localhost:${port}`);
client.on('message', message => { clientReceived.push(message.toString()); });
}
ts/nni_manager/yarn.lock
View file @ b39850f9
...@@ -541,6 +541,15 @@ ...@@ -541,6 +541,15 @@
dependencies: dependencies:
"@types/node" "*" "@types/node" "*"
"@types/express-serve-static-core@*":
version "4.17.28"
resolved "https://registry.yarnpkg.com/@types/express-serve-static-core/-/express-serve-static-core-4.17.28.tgz#c47def9f34ec81dc6328d0b1b5303d1ec98d86b8"
integrity sha512-P1BJAEAW3E2DJUlkgq4tOL3RyMunoWXqbSCygWo5ZIWTjUgN1YnaXWW4VWl/oc8vs/XoYibEGBKP0uZyF4AHig==
dependencies:
"@types/node" "*"
"@types/qs" "*"
"@types/range-parser" "*"
"@types/express-serve-static-core@^4.17.18": "@types/express-serve-static-core@^4.17.18":
version "4.17.21" version "4.17.21"
resolved "https://registry.yarnpkg.com/@types/express-serve-static-core/-/express-serve-static-core-4.17.21.tgz#a427278e106bca77b83ad85221eae709a3414d42" resolved "https://registry.yarnpkg.com/@types/express-serve-static-core/-/express-serve-static-core-4.17.21.tgz#a427278e106bca77b83ad85221eae709a3414d42"
...@@ -550,6 +559,25 @@ ...@@ -550,6 +559,25 @@
"@types/qs" "*" "@types/qs" "*"
"@types/range-parser" "*" "@types/range-parser" "*"
"@types/express-ws@^3.0.1":
version "3.0.1"
resolved "https://registry.yarnpkg.com/@types/express-ws/-/express-ws-3.0.1.tgz#6fbf5dfdbeedd16479ccbeecbca63c14be26612e"
integrity sha512-VguRXzcpPBF0IggIGpUoM65cZJDfMQxoc6dKoCz1yLzcwcXW7ft60yhq3ygKhyEhEIQFtLrWjyz4AJ1qjmzCFw==
dependencies:
"@types/express" "*"
"@types/express-serve-static-core" "*"
"@types/ws" "*"
"@types/express@*":
version "4.17.13"
resolved "https://registry.yarnpkg.com/@types/express/-/express-4.17.13.tgz#a76e2995728999bab51a33fabce1d705a3709034"
integrity sha512-6bSZTPaTIACxn48l50SR+axgrqm6qXFIxrdAKaG6PaJk3+zuUr35hBlgT7vOmJcum+OEaIBLtHV/qloEAFITeA==
dependencies:
"@types/body-parser" "*"
"@types/express-serve-static-core" "^4.17.18"
"@types/qs" "*"
"@types/serve-static" "*"
"@types/express@^4.17.2": "@types/express@^4.17.2":
version "4.17.12" version "4.17.12"
resolved "https://registry.yarnpkg.com/@types/express/-/express-4.17.12.tgz#4bc1bf3cd0cfe6d3f6f2853648b40db7d54de350" resolved "https://registry.yarnpkg.com/@types/express/-/express-4.17.12.tgz#4bc1bf3cd0cfe6d3f6f2853648b40db7d54de350"
...@@ -813,10 +841,10 @@ ...@@ -813,10 +841,10 @@
resolved "https://registry.yarnpkg.com/@types/tough-cookie/-/tough-cookie-4.0.0.tgz#fef1904e4668b6e5ecee60c52cc6a078ffa6697d" resolved "https://registry.yarnpkg.com/@types/tough-cookie/-/tough-cookie-4.0.0.tgz#fef1904e4668b6e5ecee60c52cc6a078ffa6697d"
integrity sha512-I99sngh224D0M7XgW1s120zxCt3VYQ3IQsuw3P3jbq5GG4yc79+ZjyKznyOGIQrflfylLgcfekeZW/vk0yng6A== integrity sha512-I99sngh224D0M7XgW1s120zxCt3VYQ3IQsuw3P3jbq5GG4yc79+ZjyKznyOGIQrflfylLgcfekeZW/vk0yng6A==
"@types/ws@^7.4.4": "@types/ws@*", "@types/ws@^8.5.3":
version "7.4.4" version "8.5.3"
resolved "https://registry.yarnpkg.com/@types/ws/-/ws-7.4.4.tgz#93e1e00824c1de2608c30e6de4303ab3b4c0c9bc" resolved "https://registry.yarnpkg.com/@types/ws/-/ws-8.5.3.tgz#7d25a1ffbecd3c4f2d35068d0b283c037003274d"
integrity sha512-d/7W23JAXPodQNbOZNXvl2K+bqAQrCMwlh/nuQsPSQk6Fq0opHoPrUw43aHsvSbIiQPr8Of2hkFbnz1XBFVyZQ== integrity sha512-6YOoWjruKj1uLf3INHH7D3qTXwFfEsg1kf3c0uDdSBJwfa/llkwIjrAGV7j7mVgGNbzTQ3HiHKKDXl6bJPD97w==
dependencies: dependencies:
"@types/node" "*" "@types/node" "*"
...@@ -2094,6 +2122,13 @@ express-joi-validator@^2.0.1: ...@@ -2094,6 +2122,13 @@ express-joi-validator@^2.0.1:
extend "2.0.x" extend "2.0.x"
joi "6.x.x" joi "6.x.x"
express-ws@^5.0.2:
version "5.0.2"
resolved "https://registry.yarnpkg.com/express-ws/-/express-ws-5.0.2.tgz#5b02d41b937d05199c6c266d7cc931c823bda8eb"
integrity sha512-0uvmuk61O9HXgLhGl3QhNSEtRsQevtmbL94/eILaliEADZBHZOQUAiHFrGPrgsjikohyrmSG5g+sCfASTt0lkQ==
dependencies:
ws "^7.4.6"
express@^4.17.2: express@^4.17.2:
version "4.17.2" version "4.17.2"
resolved "https://registry.yarnpkg.com/express/-/express-4.17.2.tgz#c18369f265297319beed4e5558753cc8c1364cb3" resolved "https://registry.yarnpkg.com/express/-/express-4.17.2.tgz#c18369f265297319beed4e5558753cc8c1364cb3"
...@@ -5728,9 +5763,14 @@ ws@^6.2.1: ...@@ -5728,9 +5763,14 @@ ws@^6.2.1:
async-limiter "~1.0.0" async-limiter "~1.0.0"
ws@^7.4.6: ws@^7.4.6:
version "7.4.6" version "7.5.7"
resolved "https://registry.yarnpkg.com/ws/-/ws-7.4.6.tgz#5654ca8ecdeee47c33a9a4bf6d28e2be2980377c" resolved "https://registry.yarnpkg.com/ws/-/ws-7.5.7.tgz#9e0ac77ee50af70d58326ecff7e85eb3fa375e67"
integrity sha512-YmhHDO4MzaDLB+M9ym/mDA5z0naX8j7SIlT8f8z+I0VtzsRbekxEutHSme7NPS2qE8StCYQNUnfWdXta/Yu85A== integrity sha512-KMvVuFzpKBuiIXW3E4u3mySRO2/mCHSyZDJQM5NQ9Q9KHWHWh0NHgfbRMLLrceUK5qAL4ytALJbpRMjixFZh8A==
ws@^8.5.0:
version "8.5.0"
resolved "https://registry.yarnpkg.com/ws/-/ws-8.5.0.tgz#bfb4be96600757fe5382de12c670dab984a1ed4f"
integrity sha512-BWX0SWVgLPzYwF8lTzEy1egjhS4S4OEAHfsO8o65WOVsrnSRGaSiUaa9e0ggGlkMTtBlmOpEXiie9RUcBO86qg==
xml2js@0.2.8: xml2js@0.2.8:
version "0.2.8" version "0.2.8"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment