utils.ts 12.3 KB
Newer Older
liuzhe-lz's avatar
liuzhe-lz committed
1
2
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
Deshui Yu's avatar
Deshui Yu committed
3

4
import assert from 'assert';
Deshui Yu's avatar
Deshui Yu committed
5
import { randomBytes } from 'crypto';
6
7
import cpp from 'child-process-promise';
import cp from 'child_process';
8
import { ChildProcess, spawn, StdioOptions } from 'child_process';
9
10
11
12
13
import dgram from 'dgram';
import fs from 'fs';
import net from 'net';
import os from 'os';
import path from 'path';
liuzhe-lz's avatar
liuzhe-lz committed
14
import * as timersPromises from 'timers/promises';
15
import lockfile from 'lockfile';
Deshui Yu's avatar
Deshui Yu committed
16
17
import { Deferred } from 'ts-deferred';
import { Container } from 'typescript-ioc';
18
import glob from 'glob';
Deshui Yu's avatar
Deshui Yu committed
19
20

import { Database, DataStore } from './datastore';
21
22
import globals from './globals';
import { resetGlobals } from './globals/unittest';  // TODO: this file should not contain unittest helpers
23
import { ExperimentConfig, Manager } from './manager';
24
import { ExperimentManager } from './experimentManager';
QuanluZhang's avatar
QuanluZhang committed
25
import { HyperParameters, TrainingService, TrialJobStatus } from './trainingService';
Deshui Yu's avatar
Deshui Yu committed
26

27
function getExperimentRootDir(): string {
28
    return globals.paths.experimentRoot;
Deshui Yu's avatar
Deshui Yu committed
29
30
}

31
function getLogDir(): string {
32
    return globals.paths.logDirectory;
Deshui Yu's avatar
Deshui Yu committed
33
34
}

35
function getLogLevel(): string {
36
    return globals.args.logLevel;
37
38
}

Deshui Yu's avatar
Deshui Yu committed
39
40
41
42
function getDefaultDatabaseDir(): string {
    return path.join(getExperimentRootDir(), 'db');
}

QuanluZhang's avatar
QuanluZhang committed
43
44
45
46
function getCheckpointDir(): string {
    return path.join(getExperimentRootDir(), 'checkpoint');
}

47
48
49
50
function getExperimentsInfoPath(): string {
    return path.join(os.homedir(), 'nni-experiments', '.experiment');
}

liuzhe-lz's avatar
liuzhe-lz committed
51
52
async function mkDirP(dirPath: string): Promise<void> {
    await fs.promises.mkdir(dirPath, { recursive: true });
Deshui Yu's avatar
Deshui Yu committed
53
54
55
}

function mkDirPSync(dirPath: string): void {
liuzhe-lz's avatar
liuzhe-lz committed
56
    fs.mkdirSync(dirPath, { recursive: true });
Deshui Yu's avatar
Deshui Yu committed
57
58
}

liuzhe-lz's avatar
liuzhe-lz committed
59
const delay = timersPromises.setTimeout;
Deshui Yu's avatar
Deshui Yu committed
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

/**
 * Convert index to character
 * @param index index
 * @returns a mapping character
 */
function charMap(index: number): number {
    if (index < 26) {
        return index + 97;
    } else if (index < 52) {
        return index - 26 + 65;
    } else {
        return index - 52 + 48;
    }
}

/**
 * Generate a unique string by length
 * @param len length of string
 * @returns a unique string
 */
function uniqueString(len: number): string {
    if (len === 0) {
        return '';
    }
    const byteLength: number = Math.ceil((Math.log2(52) + Math.log2(62) * (len - 1)) / 8);
    let num: number = randomBytes(byteLength).reduce((a: number, b: number) => a * 256 + b, 0);
    const codes: number[] = [];
    codes.push(charMap(num % 52));
    num = Math.floor(num / 52);
    for (let i: number = 1; i < len; i++) {
        codes.push(charMap(num % 62));
        num = Math.floor(num / 62);
    }

    return String.fromCharCode(...codes);
}

98
99
100
101
function randomInt(max: number): number {
    return Math.floor(Math.random() * max);
}

102
103
104
105
106
function randomSelect<T>(a: T[]): T {
    assert(a !== undefined);

    return a[Math.floor(Math.random() * a.length)];
}
107

108
function getCmdPy(): string {
109
    let cmd = 'python3';
110
    if (process.platform === 'win32') {
111
112
113
114
115
        cmd = 'python';
    }
    return cmd;
}

116
/**
117
 * Generate command line to start automl algorithm(s),
QuanluZhang's avatar
QuanluZhang committed
118
 * either start advisor or start a process which runs tuner and assessor
119
 *
chicm-ms's avatar
chicm-ms committed
120
 * @param expParams: experiment startup parameters
121
122
 *
 */
123
function getMsgDispatcherCommand(expParams: ExperimentConfig): string {
chicm-ms's avatar
chicm-ms committed
124
125
126
    const clonedParams = Object.assign({}, expParams);
    delete clonedParams.searchSpace;
    return `${getCmdPy()} -m nni --exp_params ${Buffer.from(JSON.stringify(clonedParams)).toString('base64')}`;
127
128
}

129
130
131
132
/**
 * Generate parameter file name based on HyperParameters object
 * @param hyperParameters HyperParameters instance
 */
chicm-ms's avatar
chicm-ms committed
133
function generateParamFileName(hyperParameters: HyperParameters): string {
134
135
136
    assert(hyperParameters !== undefined);
    assert(hyperParameters.index >= 0);

chicm-ms's avatar
chicm-ms committed
137
    let paramFileName: string;
138
    if (hyperParameters.index == 0) {
139
140
141
142
143
144
145
        paramFileName = 'parameter.cfg';
    } else {
        paramFileName = `parameter_${hyperParameters.index}.cfg`
    }
    return paramFileName;
}

Deshui Yu's avatar
Deshui Yu committed
146
147
148
149
150
151
152
153
154
/**
 * Initialize a pseudo experiment environment for unit test.
 * Must be paired with `cleanupUnitTest()`.
 */
function prepareUnitTest(): void {
    Container.snapshot(Database);
    Container.snapshot(DataStore);
    Container.snapshot(TrainingService);
    Container.snapshot(Manager);
155
    Container.snapshot(ExperimentManager);
Deshui Yu's avatar
Deshui Yu committed
156

157
    resetGlobals();
Deshui Yu's avatar
Deshui Yu committed
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175

    const sqliteFile: string = path.join(getDefaultDatabaseDir(), 'nni.sqlite');
    try {
        fs.unlinkSync(sqliteFile);
    } catch (err) {
        // file not exists, good
    }
}

/**
 * Clean up unit test pseudo experiment.
 * Must be paired with `prepareUnitTest()`.
 */
function cleanupUnitTest(): void {
    Container.restore(Manager);
    Container.restore(TrainingService);
    Container.restore(DataStore);
    Container.restore(Database);
176
    Container.restore(ExperimentManager);
Deshui Yu's avatar
Deshui Yu committed
177
178
}

179
180
let cachedIpv4Address: string | null = null;

181
/**
182
 * Get IPv4 address of current machine.
183
 */
liuzhe-lz's avatar
liuzhe-lz committed
184
async function getIPV4Address(): Promise<string> {
185
186
    if (cachedIpv4Address !== null) {
        return cachedIpv4Address;
187
    }
188

189
190
191
192
    // creates "udp connection" to a non-exist target, and get local address of the connection.
    // since udp is connectionless, this does not send actual packets.
    const socket = dgram.createSocket('udp4');
    socket.connect(1, '192.0.2.0');
liuzhe-lz's avatar
liuzhe-lz committed
193
    for (let i = 0; i < 10; i++) {  // wait the system to initialize "connection"
liuzhe-lz's avatar
liuzhe-lz committed
194
195
196
197
198
199
200
201
        await timersPromises.setTimeout(1);
        try {
            cachedIpv4Address = socket.address().address;
            socket.close();
            return cachedIpv4Address;
        } catch (error) {
            /* retry */
        }
liuzhe-lz's avatar
liuzhe-lz committed
202
    }
liuzhe-lz's avatar
liuzhe-lz committed
203

liuzhe-lz's avatar
liuzhe-lz committed
204
    cachedIpv4Address = socket.address().address;  // if it still fails, throw the error
205
206
    socket.close();
    return cachedIpv4Address;
207
208
}

QuanluZhang's avatar
QuanluZhang committed
209
210
211
212
213
214
215
/**
 * Get the status of canceled jobs according to the hint isEarlyStopped
 */
function getJobCancelStatus(isEarlyStopped: boolean): TrialJobStatus {
    return isEarlyStopped ? 'EARLY_STOPPED' : 'USER_CANCELED';
}

216
217
218
219
/**
 * Utility method to calculate file numbers under a directory, recursively
 * @param directory directory name
 */
chicm-ms's avatar
chicm-ms committed
220
function countFilesRecursively(directory: string): Promise<number> {
221
    if (!fs.existsSync(directory)) {
222
223
224
225
226
        throw Error(`Direcotory ${directory} doesn't exist`);
    }

    const deferred: Deferred<number> = new Deferred<number>();

chicm-ms's avatar
chicm-ms committed
227
    let timeoutId: NodeJS.Timer
228
    const delayTimeout: Promise<number> = new Promise((_resolve: Function, reject: Function): void => {
229
230
231
232
233
234
235
        // Set timeout and reject the promise once reach timeout (5 seconds)
        timeoutId = setTimeout(() => {
            reject(new Error(`Timeout: path ${directory} has too many files`));
        }, 5000);
    });

    let fileCount: number = -1;
236
    let cmd: string;
237
    if (process.platform === "win32") {
238
239
        cmd = `powershell "Get-ChildItem -Path ${directory} -Recurse -File | Measure-Object | %{$_.Count}"`
    } else {
240
        cmd = `find ${directory} -type f | wc -l`;
241
242
    }
    cpp.exec(cmd).then((result) => {
243
        if (result.stdout && parseInt(result.stdout)) {
244
            fileCount = parseInt(result.stdout);
245
246
247
248
249
250
251
252
        }
        deferred.resolve(fileCount);
    });
    return Promise.race([deferred.promise, delayTimeout]).finally(() => {
        clearTimeout(timeoutId);
    });
}

253
254
255
256
/**
 * get the version of current package
 */
async function getVersion(): Promise<string> {
chicm-ms's avatar
chicm-ms committed
257
    const deferred: Deferred<string> = new Deferred<string>();
258
    import(path.join(__dirname, '..', 'package.json')).then((pkg) => {
259
        deferred.resolve(pkg.version);
260
261
    }).catch(() => {
        deferred.resolve('999.0.0-developing');
262
263
    });
    return deferred.promise;
264
}
265

266
267
268
/**
 * run command as ChildProcess
 */
J-shang's avatar
J-shang committed
269
function getTunerProc(command: string, stdio: StdioOptions, newCwd: string, newEnv: any, newShell: boolean = true, isDetached: boolean = false): ChildProcess {
270
271
    let cmd: string = command;
    let arg: string[] = [];
272
    if (process.platform === "win32") {
273
        cmd = command.split(" ", 1)[0];
274
        arg = command.substr(cmd.length + 1).split(" ");
275
        newShell = false;
276
        isDetached = true;
277
278
279
280
281
    }
    const tunerProc: ChildProcess = spawn(cmd, arg, {
        stdio,
        cwd: newCwd,
        env: newEnv,
282
283
        shell: newShell,
        detached: isDetached
284
285
286
287
288
289
290
    });
    return tunerProc;
}

/**
 * judge whether the process is alive
 */
Yuge Zhang's avatar
Yuge Zhang committed
291
async function isAlive(pid: any): Promise<boolean> {
chicm-ms's avatar
chicm-ms committed
292
    const deferred: Deferred<boolean> = new Deferred<boolean>();
293
    let alive: boolean = false;
Yuge Zhang's avatar
Yuge Zhang committed
294
    if (process.platform === 'win32') {
295
296
297
298
299
300
301
        try {
            const str = cp.execSync(`powershell.exe Get-Process -Id ${pid} -ErrorAction SilentlyContinue`).toString();
            if (str) {
                alive = true;
            }
        }
        catch (error) {
chicm-ms's avatar
chicm-ms committed
302
            //ignore
303
304
        }
    }
Yuge Zhang's avatar
Yuge Zhang committed
305
    else {
306
307
308
309
310
311
312
313
314
315
316
317
        try {
            await cpp.exec(`kill -0 ${pid}`);
            alive = true;
        } catch (error) {
            //ignore
        }
    }
    deferred.resolve(alive);
    return deferred.promise;
}

/**
318
 * kill process
319
 */
Yuge Zhang's avatar
Yuge Zhang committed
320
async function killPid(pid: any): Promise<void> {
chicm-ms's avatar
chicm-ms committed
321
    const deferred: Deferred<void> = new Deferred<void>();
322
323
    try {
        if (process.platform === "win32") {
Yuge Zhang's avatar
Yuge Zhang committed
324
            await cpp.exec(`cmd.exe /c taskkill /PID ${pid} /F`);
325
        }
326
        else {
327
328
329
330
331
332
333
334
335
            await cpp.exec(`kill -9 ${pid}`);
        }
    } catch (error) {
        // pid does not exist, do nothing here
    }
    deferred.resolve();
    return deferred.promise;
}

336
function getNewLine(): string {
337
338
339
    if (process.platform === "win32") {
        return "\r\n";
    }
340
    else {
341
342
343
344
        return "\n";
    }
}

345
346
/**
 * Use '/' to join path instead of '\' for all kinds of platform
347
 * @param path
348
349
350
351
352
353
354
 */
function unixPathJoin(...paths: any[]): string {
    const dir: string = paths.filter((path: any) => path !== '').join('/');
    if (dir === '') return '.';
    return dir;
}

355
356
357
358
359
/**
 * lock a file sync
 */
function withLockSync(func: Function, filePath: string, lockOpts: {[key: string]: any}, ...args: any): any {
    const lockName = path.join(path.dirname(filePath), path.basename(filePath) + `.lock.${process.pid}`);
360
    if (typeof lockOpts['stale'] === 'number'){
361
362
363
        const lockPath = path.join(path.dirname(filePath), path.basename(filePath) + '.lock.*');
        const lockFileNames: string[] = glob.sync(lockPath);
        const canLock: boolean = lockFileNames.map((fileName) => {
364
            return fs.existsSync(fileName) && Date.now() - fs.statSync(fileName).mtimeMs < lockOpts['stale'];
365
        }).filter(unexpired=>unexpired === true).length === 0;
366
367
368
369
370
371
372
373
374
375
        if (!canLock) {
            throw new Error('File has been locked.');
        }
    }
    lockfile.lockSync(lockName, lockOpts);
    const result = func(...args);
    lockfile.unlockSync(lockName);
    return result;
}

J-shang's avatar
J-shang committed
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
async function isPortOpen(host: string, port: number): Promise<boolean> {
    return new Promise<boolean>((resolve, reject) => {
        try{
            const stream = net.createConnection(port, host);
            const id = setTimeout(() => {
                stream.destroy();
                resolve(false);
            }, 1000);

            stream.on('connect', () => {
                clearTimeout(id);
                stream.destroy();
                resolve(true);
            });

            stream.on('error', () => {
                clearTimeout(id);
                stream.destroy();
                resolve(false);
            });
        } catch (error) {
            reject(error);
        }
    });
}

async function getFreePort(host: string, start: number, end: number): Promise<number> {
    if (start > end) {
        throw new Error(`no more free port`);
    }
    if (await isPortOpen(host, start)) {
        return await getFreePort(host, start + 1, end);
    } else {
        return start;
    }
}

413
414
415
416
417
export function importModule(modulePath: string): any {
    module.paths.unshift(path.dirname(modulePath));
    return require(path.basename(modulePath));
}

418
export {
419
    countFilesRecursively, generateParamFileName, getMsgDispatcherCommand, getCheckpointDir, getExperimentsInfoPath,
J-shang's avatar
J-shang committed
420
    getLogDir, getExperimentRootDir, getJobCancelStatus, getDefaultDatabaseDir, getIPV4Address, unixPathJoin, withLockSync, getFreePort, isPortOpen,
liuzhe-lz's avatar
liuzhe-lz committed
421
    mkDirP, mkDirPSync, delay, prepareUnitTest, cleanupUnitTest, uniqueString, randomInt, randomSelect, getLogLevel, getVersion, getCmdPy, getTunerProc, isAlive, killPid, getNewLine
422
};