Fix YAML encoding and WebSocket reconnect UT (#5189)

docs/source/conf.py

@@ -114,6 +114,7 @@ linkcheck_ignore = [
     r'https://ml\.informatik\.uni-freiburg\.de/',
     r'https://docs\.nvidia\.com/deeplearning/',
     r'https://cla\.opensource\.microsoft\.com',
+    r'https://www\.docker\.com/',
 ]
 
 # Ignore all links located in release.rst

nni/experiment/config/base.py

@@ -137,7 +137,7 @@ class ConfigBase:
         cls
             An object of ConfigBase subclass.
         """
-        with open(path) as yaml_file:
+        with open(path, encoding='utf_8') as yaml_file:
             data = yaml.safe_load(yaml_file)
         if not isinstance(data, dict):
             raise TypeError(f'Conent of config file {path} is not a dict/object')

nni/tools/nnictl/common_utils.py

@@ -20,7 +20,7 @@ from .constants import ERROR_INFO, NORMAL_INFO, WARNING_INFO
 def get_yml_content(file_path):
     '''Load yaml file content'''
     try:
-        with open(file_path, 'r') as file:
+        with open(file_path, 'r', encoding='utf_8') as file:
             return yaml.safe_load(file)
     except yaml.scanner.ScannerError as err:
         print_error('yaml file format error!')

nni/tools/nnictl/config_utils.py

@@ -158,7 +158,7 @@ class Experiments:
     def write_file(self):
         '''save config to local file'''
         try:
-            with open(self.experiment_file, 'w') as file:
+            with open(self.experiment_file, 'w', encoding='utf_8') as file:
                 nni.dump(self.experiments, file, indent=4)
         except IOError as error:
             print('Error:', error)
@@ -168,7 +168,7 @@ class Experiments:
         '''load config from local file'''
         if os.path.exists(self.experiment_file):
             try:
-                with open(self.experiment_file, 'r') as file:
+                with open(self.experiment_file, 'r', encoding='utf_8') as file:
                     return nni.load(fp=file)
             except ValueError:
                 return {}

nni/tools/nnictl/launcher.py

@@ -31,7 +31,7 @@ def create_experiment(args):
         _logger.error(f'"{config_file}" is not a valid file.')
         exit(1)
 
-    with config_file.open() as config:
+    with config_file.open(encoding='utf_8') as config:
         config_content = yaml.safe_load(config)
 
     v1_platform = config_content.get('trainingServicePlatform')

nni/tools/package_utils/config_manager.py

@@ -63,7 +63,7 @@ def _load_custom_config():
     return [algo for algo in  _load_config_file(path) if not algo.is_builtin]
 
 def _load_config_file(path):
-    with open(path) as f:
+    with open(path, encoding='utf_8') as f:
         config = yaml.safe_load(f)
     algos = []
     for algo_type in ['tuner', 'assessor', 'advisor']:

ts/nni_manager/common/globals/unittest.ts

@@ -25,6 +25,10 @@ import path from 'path';
 import type { NniManagerArgs } from './arguments';
 import { NniPaths, createPaths } from './paths';
 import type { LogStream } from './log_stream';
+// Enforce ts-node to import `shutdown.ts`.
+// Without this line it might complain "log_1.getRobustLogger is not a function".
+// "Magic. Do not touch."
+import './shutdown';
 
 // copied from https://www.typescriptlang.org/docs/handbook/2/mapped-types.html
 type Mutable<Type> = {

ts/nni_manager/core/tuner_command_channel/websocket_channel.ts

@@ -142,6 +142,7 @@ class WebSocketChannelImpl implements WebSocketChannel {
         }
 
         this.serving = false;
+        this.waitingPong = false;
         clearInterval(this.heartbeatTimer);
 
         this.ws.off('close', this.handleWsClose);

ts/nni_manager/package.json

@@ -64,14 +64,14 @@
     "chai-as-promised": "^7.1.1",
     "eslint": "^7.28.0",
     "glob": "^7.1.7",
-    "mocha": "^10.0.0",
+    "mocha": "^10.1.0",
     "node-fetch": "<3.0.0",
     "npm": ">=8.11.0",
     "nyc": "^15.1.0",
     "request": "^2.88.2",
     "rmdir": "^1.2.0",
     "tmp": "^0.2.1",
-    "ts-node": "^10.0.0",
+    "ts-node": "^10.9.1",
     "typescript": "^4.3.2"
   },
   "resolutions": {

ts/nni_manager/test/core/tuner_command_channel.test.ts

@@ -6,74 +6,71 @@ import { setTimeout } from 'timers/promises';
 
 import WebSocket from 'ws';
 
+import { Deferred } from 'common/deferred';
 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()
+const heartbeatInterval: number = 10;
+
+// for testError, must be set before serveWebSocket()
+UnitTestHelpers.setHeartbeatInterval(heartbeatInterval);
 
 /* test cases */
 
 // Start serving and let a client connect.
 async function testInit(): Promise<void> {
+    const channel = getWebSocketChannel();
+    channel.onCommand(command => { serverReceived.push(command); });
+    channel.onError(error => { catchedError = error; });
+
     server.on('connection', serveWebSocket);
-    startClient();
-    await getWebSocketChannel().init();
+    client1 = new Client('client1');
+    await channel.init();
 }
 
 // Send commands from server to client.
-async function testSend(): Promise<void> {
+async function testSend(client: Client): Promise<void> {
     const channel = getWebSocketChannel();
 
     channel.sendCommand(command1);
     channel.sendCommand(command2);
-    await setTimeout(10);
+    await setTimeout(heartbeatInterval);
 
-    assert.equal(clientReceived.length, 2);
-    assert.equal(clientReceived[0], command1);
-    assert.equal(clientReceived[1], command2);
+    assert.deepEqual(client.received, [command1, command2]);
 }
 
 // Send commands from client to server.
-async function testReceive(): Promise<void> {
-    const channel = getWebSocketChannel();
-    channel.onCommand(command => { serverReceived.push(command); });
+async function testReceive(client: Client): Promise<void> {
+    serverReceived.length = 0;
 
-    client.send(command1);
-    client.send(command2);
-    await setTimeout(10);
+    client.ws.send(command2);
+    client.ws.send(command1);
+    await setTimeout(heartbeatInterval);
 
-    assert.equal(serverReceived.length, 2);
-    assert.deepEqual(serverReceived[0], command1);
-    assert.deepEqual(serverReceived[1], command2);
+    assert.deepEqual(serverReceived, [command2, command1]);
 }
 
 // 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();
+        client1.ws.terminate();
         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);
+    client1.ws.pause();
+    await setTimeout(heartbeatInterval * 3);
+    client1.ws.resume();
 
     assert.notEqual(catchedError, undefined);
-    client.resume();
 }
 
-// WebSocket might get broken in long experiments. Simulate reconnect.
+// If the client losses connection by accident but not crashed, it will reconnect.
 async function testReconnect(): Promise<void> {
-    client.close();
-    startClient();
-    testInit();
-    testSend();
+    client2 = new Client('client2');
+    await client2.deferred.promise;
 }
 
 // Clean up.
@@ -81,21 +78,24 @@ async function testShutdown(): Promise<void> {
     const channel = getWebSocketChannel();
     await channel.shutdown();
 
-    try {
-        client.close();
-    } catch (error) {
-        console.log('Error on clean up:', error);
-    }
+    client1.ws.close();
+    client2.ws.close();
     server.close();
 }
 
 /* register */
 describe('## tuner_command_channel ##', () => {
     it('init', testInit);
-    it('send', testSend);
-    it('receive', testReceive);
-    it('catch error', testError);
+
+    it('send', () => testSend(client1));
+    it('receive', () => testReceive(client1));
+
+    it('mock timeout', testError);
     it('reconnect', testReconnect);
+
+    it('send after reconnect', () => testSend(client2));
+    it('receive after reconnect', () => testReceive(client2));
+
     it('shutdown', testShutdown);
 });
 
@@ -103,17 +103,29 @@ describe('## tuner_command_channel ##', () => {
 
 const command1 = 'T_hello world';
 const command2 = 'T_你好';
-const commandPing = 'PI';
 
 const server = new WebSocket.Server({ port: 0 });
-let client!: WebSocket;
+let client1!: Client;
+let client2!: Client;
 
 const serverReceived: string[] = [];
-const clientReceived: string[] = [];
 let catchedError: Error | undefined;
 
-function startClient() {
+class Client {
+    name: string;
+    received: string[] = [];
+    ws!: WebSocket;
+    deferred: Deferred<void> = new Deferred();
+
+    constructor(name: string) {
+        this.name = name;
         const port = (server.address() as any).port;
-    client = new WebSocket(`ws://localhost:${port}`);
-    client.on('message', message => { clientReceived.push(message.toString()); });
+        this.ws = new WebSocket(`ws://localhost:${port}`);
+        this.ws.on('message', (data, _isBinary) => {
+            this.received.push(data.toString());
+        });
+        this.ws.on('open', () => {
+            this.deferred.resolve();
+        });
+    }
 }

ts/nni_manager/test/register.js

 require('ts-node/register');
 require('app-module-path/cwd');
+
+require('common/globals/unittest');

ts/nni_manager/yarn.lock

@@ -205,6 +205,13 @@
   resolved "https://registry.yarnpkg.com/@colors/colors/-/colors-1.5.0.tgz#bb504579c1cae923e6576a4f5da43d25f97bdbd9"
   integrity sha512-ooWCrlZP11i8GImSjTHYHLkvFDP48nS4+204nGb1RiX/WXYHmJA2III9/e2DWVabCESdW7hBAEzHRqUn9OUVvQ==
 
+"@cspotcode/source-map-support@^0.8.0":
+  version "0.8.1"
+  resolved "https://registry.yarnpkg.com/@cspotcode/source-map-support/-/source-map-support-0.8.1.tgz#00629c35a688e05a88b1cda684fb9d5e73f000a1"
+  integrity sha512-IchNf6dN4tHoMFIn/7OE8LWZ19Y6q/67Bmf6vnGREv8RSbBVb9LPJxEcnwrcwX6ixSvaiGoomAUvu4YSxXrVgw==
+  dependencies:
+    "@jridgewell/trace-mapping" "0.3.9"
+
 "@eslint/eslintrc@^0.4.2":
   version "0.4.2"
   resolved "https://registry.yarnpkg.com/@eslint/eslintrc/-/eslintrc-0.4.2.tgz#f63d0ef06f5c0c57d76c4ab5f63d3835c51b0179"
@@ -258,6 +265,24 @@
   resolved "https://registry.yarnpkg.com/@istanbuljs/schema/-/schema-0.1.3.tgz#e45e384e4b8ec16bce2fd903af78450f6bf7ec98"
   integrity sha512-ZXRY4jNvVgSVQ8DL3LTcakaAtXwTVUxE81hslsyD2AtoXW/wVob10HkOJ1X/pAlcI7D+2YoZKg5do8G/w6RYgA==
 
+"@jridgewell/resolve-uri@^3.0.3":
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/@jridgewell/resolve-uri/-/resolve-uri-3.1.0.tgz#2203b118c157721addfe69d47b70465463066d78"
+  integrity sha512-F2msla3tad+Mfht5cJq7LSXcdudKTWCVYUgw6pLFOOHSTtZlj6SWNYAp+AhuqLmWdBO2X5hPrLcu8cVP8fy28w==
+
+"@jridgewell/sourcemap-codec@^1.4.10":
+  version "1.4.14"
+  resolved "https://registry.yarnpkg.com/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.4.14.tgz#add4c98d341472a289190b424efbdb096991bb24"
+  integrity sha512-XPSJHWmi394fuUuzDnGz1wiKqWfo1yXecHQMRf2l6hztTO+nPru658AyDngaBe7isIxEkRsPR3FZh+s7iVa4Uw==
+
+"@jridgewell/trace-mapping@0.3.9":
+  version "0.3.9"
+  resolved "https://registry.yarnpkg.com/@jridgewell/trace-mapping/-/trace-mapping-0.3.9.tgz#6534fd5933a53ba7cbf3a17615e273a0d1273ff9"
+  integrity sha512-3Belt6tdc8bPgAtbcmdtNJlirVoTmEb5e2gC94PnkwEW9jI6CAHUeoG85tjWP5WquqfavoMtMwiG4P926ZKKuQ==
+  dependencies:
+    "@jridgewell/resolve-uri" "^3.0.3"
+    "@jridgewell/sourcemap-codec" "^1.4.10"
+
 "@mapbox/node-pre-gyp@^1.0.0":
   version "1.0.9"
   resolved "https://registry.yarnpkg.com/@mapbox/node-pre-gyp/-/node-pre-gyp-1.0.9.tgz#09a8781a3a036151cdebbe8719d6f8b25d4058bc"
@@ -508,10 +533,10 @@
   resolved "https://registry.yarnpkg.com/@tsconfig/node14/-/node14-1.0.0.tgz#5bd046e508b1ee90bc091766758838741fdefd6e"
   integrity sha512-RKkL8eTdPv6t5EHgFKIVQgsDapugbuOptNd9OOunN/HAkzmmTnZELx1kNCK0rSdUYGmiFMM3rRQMAWiyp023LQ==
 
-"@tsconfig/node16@^1.0.1":
-  version "1.0.1"
-  resolved "https://registry.yarnpkg.com/@tsconfig/node16/-/node16-1.0.1.tgz#a6ca6a9a0ff366af433f42f5f0e124794ff6b8f1"
-  integrity sha512-FTgBI767POY/lKNDNbIzgAX6miIDBs6NTCbdlDb8TrWovHsSvaVIZDlTqym29C6UqhzwcJx4CYr+AlrMywA0cA==
+"@tsconfig/node16@^1.0.2":
+  version "1.0.3"
+  resolved "https://registry.yarnpkg.com/@tsconfig/node16/-/node16-1.0.3.tgz#472eaab5f15c1ffdd7f8628bd4c4f753995ec79e"
+  integrity sha512-yOlFc+7UtL/89t2ZhjPvvB/DeAr3r+Dq58IgzsFkOAvVC6NMJXmCGjbptdXdR9qsX7pKcTL+s87FtYREi2dEEQ==
 
 "@types/body-parser@*":
   version "1.19.0"
@@ -941,11 +966,6 @@
     "@typescript-eslint/types" "4.26.0"
     eslint-visitor-keys "^2.0.0"
 
-"@ungap/promise-all-settled@1.1.2":
-  version "1.1.2"
-  resolved "https://registry.yarnpkg.com/@ungap/promise-all-settled/-/promise-all-settled-1.1.2.tgz#aa58042711d6e3275dd37dc597e5d31e8c290a44"
-  integrity sha512-sL/cEvJWAnClXw0wHk85/2L0G6Sj8UB0Ctc1TEMbKSsmpRosqhwj9gWgFRZSrBr2f9tiXISwNhCPmlfqUqyb9Q==
-
 abbrev@1, abbrev@~1.1.1:
   version "1.1.1"
   resolved "https://registry.yarnpkg.com/abbrev/-/abbrev-1.1.1.tgz#f8f2c887ad10bf67f634f005b6987fed3179aac8"
@@ -964,7 +984,12 @@ acorn-jsx@^5.3.1:
   resolved "https://registry.yarnpkg.com/acorn-jsx/-/acorn-jsx-5.3.1.tgz#fc8661e11b7ac1539c47dbfea2e72b3af34d267b"
   integrity sha512-K0Ptm/47OKfQRpNQ2J/oIN/3QYiK6FwW+eJbILhsdxh2WTLdl+30o8aGdTbm5JbffpFFAg/g+zi1E+jvJha5ng==
 
-acorn@>=8.3.0, acorn@^7.4.0:
+acorn-walk@^8.1.1:
+  version "8.2.0"
+  resolved "https://registry.yarnpkg.com/acorn-walk/-/acorn-walk-8.2.0.tgz#741210f2e2426454508853a2f44d0ab83b7f69c1"
+  integrity sha512-k+iyHEuPgSw6SbuDpGQM+06HQUa04DZ3o+F6CSzXMvvI5KMvnaEqXe+YVe555R9nn6GPt404fos4wcgpw12SDA==
+
+acorn@>=8.3.0, acorn@^7.4.0, acorn@^8.4.1:
   version "8.3.0"
   resolved "https://registry.yarnpkg.com/acorn/-/acorn-8.3.0.tgz#1193f9b96c4e8232f00b11a9edff81b2c8b98b88"
   integrity sha512-tqPKHZ5CaBJw0Xmy0ZZvLs1qTV+BNFSyvn77ASXkpBNfIRk8ev26fKrD9iLGwGA9zedPao52GSHzq8lyZG0NUw==
@@ -1321,11 +1346,6 @@ browserslist@^4.16.6:
     escalade "^3.1.1"
     node-releases "^1.1.71"
 
-buffer-from@^1.0.0:
-  version "1.1.1"
-  resolved "https://registry.yarnpkg.com/buffer-from/-/buffer-from-1.1.1.tgz#32713bc028f75c02fdb710d7c7bcec1f2c6070ef"
-  integrity sha512-MQcXEUbCKtEo7bhqEs6560Hyd4XaovZlO/k9V3hjVUF/zwW7KBVdSK4gIt/bzwS9MbR5qob+F5jusZsb0YQK2A==
-
 buffer@^5.5.0:
   version "5.7.1"
   resolved "https://registry.yarnpkg.com/buffer/-/buffer-5.7.1.tgz#ba62e7c13133053582197160851a8f648e99eed0"
@@ -3661,12 +3681,11 @@ mkdirp@^1.0.3, mkdirp@^1.0.4:
   resolved "https://registry.yarnpkg.com/mkdirp/-/mkdirp-1.0.4.tgz#3eb5ed62622756d79a5f0e2a221dfebad75c2f7e"
   integrity sha512-vVqVZQyf3WLx2Shd0qJ9xuvqgAyKPLAiqITEtqW0oIUjzo3PePDd6fW9iFz30ef7Ysp/oiWqbhszeGWW2T6Gzw==
 
-mocha@^10.0.0:
-  version "10.0.0"
-  resolved "https://registry.yarnpkg.com/mocha/-/mocha-10.0.0.tgz#205447d8993ec755335c4b13deba3d3a13c4def9"
-  integrity sha512-0Wl+elVUD43Y0BqPZBzZt8Tnkw9CMUdNYnUsTfOM1vuhJVZL+kiesFYsqwBkEEuEixaiPe5ZQdqDgX2jddhmoA==
+mocha@^10.1.0:
+  version "10.1.0"
+  resolved "https://registry.yarnpkg.com/mocha/-/mocha-10.1.0.tgz#dbf1114b7c3f9d0ca5de3133906aea3dfc89ef7a"
+  integrity sha512-vUF7IYxEoN7XhQpFLxQAEMtE4W91acW4B6En9l97MwE9stL1A9gusXfoHZCLVHDUJ/7V5+lbCM6yMqzo5vNymg==
   dependencies:
-    "@ungap/promise-all-settled" "1.1.2"
     ansi-colors "4.1.1"
     browser-stdout "1.3.1"
     chokidar "3.5.3"
@@ -4911,20 +4930,12 @@ sort-keys@^2.0.0:
   dependencies:
     is-plain-obj "^1.0.0"
 
-source-map-support@^0.5.17:
-  version "0.5.19"
-  resolved "https://registry.yarnpkg.com/source-map-support/-/source-map-support-0.5.19.tgz#a98b62f86dcaf4f67399648c085291ab9e8fed61"
-  integrity sha512-Wonm7zOCIJzBGQdB+thsPar0kYuCIzYvxZwlBa87yi/Mdjv7Tip2cyVbLj5o0cFPN4EVkuTwb3GDDyUx2DGnGw==
-  dependencies:
-    buffer-from "^1.0.0"
-    source-map "^0.6.0"
-
 source-map@^0.5.0:
   version "0.5.7"
   resolved "https://registry.yarnpkg.com/source-map/-/source-map-0.5.7.tgz#8a039d2d1021d22d1ea14c80d8ea468ba2ef3fcc"
   integrity sha1-igOdLRAh0i0eoUyA2OpGi6LvP8w=
 
-source-map@^0.6.0, source-map@^0.6.1:
+source-map@^0.6.1:
   version "0.6.1"
   resolved "https://registry.yarnpkg.com/source-map/-/source-map-0.6.1.tgz#74722af32e9614e9c287a8d0bbde48b5e2f1a263"
   integrity sha512-UjgapumWlbMhkBgzT7Ykc5YXUT46F0iKu8SGXq0bcwP5dz/h0Plj6enJqjz1Zbq2l5WaqYnrVbwWOWMyF3F47g==
@@ -5226,20 +5237,23 @@ ts-deferred@^1.0.4:
   resolved "https://registry.yarnpkg.com/ts-deferred/-/ts-deferred-1.0.4.tgz#58145ebaeef5b8f2a290b8cec3d060839f9489c7"
   integrity sha1-WBReuu71uPKikLjOw9Bgg5+Uicc=
 
-ts-node@^10.0.0:
-  version "10.0.0"
-  resolved "https://registry.yarnpkg.com/ts-node/-/ts-node-10.0.0.tgz#05f10b9a716b0b624129ad44f0ea05dac84ba3be"
-  integrity sha512-ROWeOIUvfFbPZkoDis0L/55Fk+6gFQNZwwKPLinacRl6tsxstTF1DbAcLKkovwnpKMVvOMHP1TIbnwXwtLg1gg==
+ts-node@^10.9.1:
+  version "10.9.1"
+  resolved "https://registry.yarnpkg.com/ts-node/-/ts-node-10.9.1.tgz#e73de9102958af9e1f0b168a6ff320e25adcff4b"
+  integrity sha512-NtVysVPkxxrwFGUUxGYhfux8k78pQB3JqYBXlLRZgdGUqTO5wU/UyHop5p70iEbGhB7q5KmiZiU0Y3KlJrScEw==
   dependencies:
+    "@cspotcode/source-map-support" "^0.8.0"
     "@tsconfig/node10" "^1.0.7"
     "@tsconfig/node12" "^1.0.7"
     "@tsconfig/node14" "^1.0.0"
-    "@tsconfig/node16" "^1.0.1"
+    "@tsconfig/node16" "^1.0.2"
+    acorn "^8.4.1"
+    acorn-walk "^8.1.1"
     arg "^4.1.0"
     create-require "^1.1.0"
     diff "^4.0.1"
     make-error "^1.1.1"
-    source-map-support "^0.5.17"
+    v8-compile-cache-lib "^3.0.1"
     yn "3.1.1"
 
 tslib@^1.8.1:
@@ -5384,6 +5398,11 @@ uuid@^3.0.0, uuid@^3.3.2, uuid@^3.3.3:
   resolved "https://registry.yarnpkg.com/uuid/-/uuid-3.4.0.tgz#b23e4358afa8a202fe7a100af1f5f883f02007ee"
   integrity sha512-HjSDRw6gZE5JMggctHBcjVak08+KEVhSIiDzFnT9S9aegmp85S/bReBVTb4QTFaRNptJ9kuYaNhnbNEOkbKb/A==
 
+v8-compile-cache-lib@^3.0.1:
+  version "3.0.1"
+  resolved "https://registry.yarnpkg.com/v8-compile-cache-lib/-/v8-compile-cache-lib-3.0.1.tgz#6336e8d71965cb3d35a1bbb7868445a7c05264bf"
+  integrity sha512-wa7YjyUGfNZngI/vtK0UHAN+lgDCxBPCylVXGp0zu59Fz5aiGtNXaq3DhIov063MorB+VfufLh3JlF2KdTK3xg==
+
 v8-compile-cache@^2.0.3:
   version "2.3.0"
   resolved "https://registry.yarnpkg.com/v8-compile-cache/-/v8-compile-cache-2.3.0.tgz#2de19618c66dc247dcfb6f99338035d8245a2cee"