nnimanager.ts 34.6 KB
Newer Older
Deshui Yu's avatar
Deshui Yu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
/**
 * Copyright (c) Microsoft Corporation
 * All rights reserved.
 *
 * MIT License
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
 * documentation files (the "Software"), to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
 * to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
 * BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

'use strict';

import * as assert from 'assert';
import * as cpp from 'child-process-promise';
goooxu's avatar
goooxu committed
24
import { ChildProcess, spawn, StdioOptions } from 'child_process';
Deshui Yu's avatar
Deshui Yu committed
25
26
27
import { Deferred } from 'ts-deferred';
import * as component from '../common/component';
import { DataStore, MetricDataRecord, MetricType, TrialJobInfo } from '../common/datastore';
28
import { NNIError } from '../common/errors';
29
import { getExperimentId, setInitTrialSequenceId } from '../common/experimentStartupInfo';
Deshui Yu's avatar
Deshui Yu committed
30
31
import { getLogger, Logger } from '../common/log';
import {
chicm-ms's avatar
chicm-ms committed
32
    ExperimentParams, ExperimentProfile, Manager, ExperimentStatus,
33
    NNIManagerStatus, ProfileUpdateType, TrialJobStatistics
Deshui Yu's avatar
Deshui Yu committed
34
35
36
37
} from '../common/manager';
import {
    TrainingService, TrialJobApplicationForm, TrialJobDetail, TrialJobMetric, TrialJobStatus
} from '../common/trainingService';
38
import { delay, getCheckpointDir, getExperimentRootDir, getLogDir, getMsgDispatcherCommand, mkDirP, getTunerProc, getLogLevel, isAlive, killPid } from '../common/utils';
Deshui Yu's avatar
Deshui Yu committed
39
import {
chicm-ms's avatar
chicm-ms committed
40
    ADD_CUSTOMIZED_TRIAL_JOB, INITIALIZE, INITIALIZED, KILL_TRIAL_JOB, NEW_TRIAL_JOB, NO_MORE_TRIAL_JOBS, PING,
41
    REPORT_METRIC_DATA, REQUEST_TRIAL_JOBS, SEND_TRIAL_JOB_PARAMETER, TERMINATE, TRIAL_END, UPDATE_SEARCH_SPACE, IMPORT_DATA
Deshui Yu's avatar
Deshui Yu committed
42
} from './commands';
43
import { createDispatcherInterface, IpcInterface } from './ipcInterface';
Deshui Yu's avatar
Deshui Yu committed
44
45

/**
chicm-ms's avatar
chicm-ms committed
46
 * NNIManager which implements Manager interface
Deshui Yu's avatar
Deshui Yu committed
47
48
49
 */
class NNIManager implements Manager {
    private trainingService: TrainingService;
50
    private dispatcher: IpcInterface | undefined;
51
    private currSubmittedTrialNum: number;  // need to be recovered
QuanluZhang's avatar
QuanluZhang committed
52
    private trialConcurrencyChange: number; // >0: increase, <0: decrease
Deshui Yu's avatar
Deshui Yu committed
53
54
55
56
    private customizedTrials: string[]; // need to be recovered
    private log: Logger;
    private dataStore: DataStore;
    private experimentProfile: ExperimentProfile;
57
    private dispatcherPid: number;
58
    private status: NNIManagerStatus;
QuanluZhang's avatar
QuanluZhang committed
59
60
    private waitingTrials: string[];
    private trialJobs: Map<string, TrialJobDetail>;
61
    private trialDataForTuner: string;
SparkSnail's avatar
SparkSnail committed
62
    private readonly: boolean;
63

64
65
    private trialJobMetricListener: (metric: TrialJobMetric) => void;
    
Deshui Yu's avatar
Deshui Yu committed
66
67
    constructor() {
        this.currSubmittedTrialNum = 0;
QuanluZhang's avatar
QuanluZhang committed
68
        this.trialConcurrencyChange = 0;
Deshui Yu's avatar
Deshui Yu committed
69
70
71
        this.customizedTrials = [];
        this.trainingService = component.get(TrainingService);
        assert(this.trainingService);
72
        this.dispatcherPid = 0;
QuanluZhang's avatar
QuanluZhang committed
73
74
        this.waitingTrials = [];
        this.trialJobs = new Map<string, TrialJobDetail>();
75
        this.trialDataForTuner = '';
SparkSnail's avatar
SparkSnail committed
76
        this.readonly = false;
Deshui Yu's avatar
Deshui Yu committed
77
78
79

        this.log = getLogger();
        this.dataStore = component.get(DataStore);
80
81
82
83
        this.experimentProfile = this.createEmptyExperimentProfile();
        this.status = {
            status: 'INITIALIZED',
            errors: []
Deshui Yu's avatar
Deshui Yu committed
84
        };
85
86
87
88
89
        this.trialJobMetricListener = (metric: TrialJobMetric) => {
            this.onTrialJobMetrics(metric).catch((err: Error) => {
                this.criticalError(NNIError.FromError(err, 'Job metrics error: '));
            });
        };
Deshui Yu's avatar
Deshui Yu committed
90
91
92
    }

    public updateExperimentProfile(experimentProfile: ExperimentProfile, updateType: ProfileUpdateType): Promise<void> {
SparkSnail's avatar
SparkSnail committed
93
94
95
        if (this.readonly) {
            return Promise.reject(new Error('Error: can not update experiment profile in readonly mode!'));
        }
Deshui Yu's avatar
Deshui Yu committed
96
97
98
99
100
101
102
103
104
105
        switch (updateType) {
            case 'TRIAL_CONCURRENCY':
                this.updateTrialConcurrency(experimentProfile.params.trialConcurrency);
                break;
            case 'MAX_EXEC_DURATION':
                this.updateMaxExecDuration(experimentProfile.params.maxExecDuration);
                break;
            case 'SEARCH_SPACE':
                this.updateSearchSpace(experimentProfile.params.searchSpace);
                break;
QuanluZhang's avatar
QuanluZhang committed
106
107
108
            case 'MAX_TRIAL_NUM':
                this.updateMaxTrialNum(experimentProfile.params.maxTrialNum);
                break;
Deshui Yu's avatar
Deshui Yu committed
109
110
111
112
113
114
115
            default:
                throw new Error('Error: unrecognized updateType');
        }

        return this.storeExperimentProfile();
    }

116
    public importData(data: string): Promise<void> {
SparkSnail's avatar
SparkSnail committed
117
118
119
        if (this.readonly) {
            return Promise.reject(new Error('Error: can not import data in readonly mode!'));
        }
120
121
122
123
124
125
126
127
128
129
        if (this.dispatcher === undefined) {
            return Promise.reject(
                new Error('tuner has not been setup')
            );
        }
        this.dispatcher.sendCommand(IMPORT_DATA, data);

        return this.dataStore.storeTrialJobEvent('IMPORT_DATA', '', data);
    }

130
131
132
133
    public async exportData(): Promise<string> {
        return this.dataStore.exportTrialHpConfigs();
    }

Deshui Yu's avatar
Deshui Yu committed
134
    public addCustomizedTrialJob(hyperParams: string): Promise<void> {
SparkSnail's avatar
SparkSnail committed
135
136
137
        if (this.readonly) {
            return Promise.reject(new Error('Error: can not add customized trial job in readonly mode!'));
        }
Deshui Yu's avatar
Deshui Yu committed
138
139
140
141
142
143
144
145
146
147
148
149
        if (this.currSubmittedTrialNum >= this.experimentProfile.params.maxTrialNum) {
            return Promise.reject(
                new Error('reach maxTrialNum')
            );
        }
        this.customizedTrials.push(hyperParams);

        // trial id has not been generated yet, thus use '' instead
        return this.dataStore.storeTrialJobEvent('ADD_CUSTOMIZED', '', hyperParams);
    }

    public async cancelTrialJobByUser(trialJobId: string): Promise<void> {
SparkSnail's avatar
SparkSnail committed
150
151
152
        if (this.readonly) {
            return Promise.reject(new Error('Error: can not cancel trial job in readonly mode!'));
        }
chicm-ms's avatar
chicm-ms committed
153
        this.log.info(`User cancelTrialJob: ${trialJobId}`);
Deshui Yu's avatar
Deshui Yu committed
154
155
156
157
158
        await this.trainingService.cancelTrialJob(trialJobId);
        await this.dataStore.storeTrialJobEvent('USER_TO_CANCEL', trialJobId, '');
    }

    public async startExperiment(expParams: ExperimentParams): Promise<string> {
chicm-ms's avatar
chicm-ms committed
159
        this.log.info(`Starting experiment: ${this.experimentProfile.id}`);
Deshui Yu's avatar
Deshui Yu committed
160
161
162
        this.experimentProfile.params = expParams;
        await this.storeExperimentProfile();
        this.log.debug('Setup tuner...');
163

164
        // Set up multiphase config
165
        if (expParams.multiPhase && this.trainingService.isMultiPhaseJobSupported) {
166
167
            this.trainingService.setClusterMetadata('multiPhase', expParams.multiPhase.toString());
        }
168
169
170
171
        // Set up versionCheck config
        if (expParams.versionCheck !== undefined) {
            this.trainingService.setClusterMetadata('version_check', expParams.versionCheck.toString());
        }
SparkSnail's avatar
SparkSnail committed
172
173
174
175
        // Set up logCollection config
        if (expParams.logCollection !== undefined) {
            this.trainingService.setClusterMetadata('log_collection', expParams.logCollection.toString());
        }
176

QuanluZhang's avatar
QuanluZhang committed
177
        const dispatcherCommand: string = getMsgDispatcherCommand(expParams.tuner, expParams.assessor, expParams.advisor,
goooxu's avatar
goooxu committed
178
            expParams.multiPhase, expParams.multiThread);
179
        this.log.debug(`dispatcher command: ${dispatcherCommand}`);
QuanluZhang's avatar
QuanluZhang committed
180
        const checkpointDir: string = await this.createCheckpointDir();
Deshui Yu's avatar
Deshui Yu committed
181
        this.setupTuner(
182
183
            dispatcherCommand,
            undefined,
Deshui Yu's avatar
Deshui Yu committed
184
            'start',
QuanluZhang's avatar
QuanluZhang committed
185
            checkpointDir);
Deshui Yu's avatar
Deshui Yu committed
186

187
        this.experimentProfile.startTime = Date.now();
chicm-ms's avatar
chicm-ms committed
188
        this.setStatus('RUNNING');
Deshui Yu's avatar
Deshui Yu committed
189
        await this.storeExperimentProfile();
190
191
        this.run().catch((err: Error) => {
            this.criticalError(err);
Deshui Yu's avatar
Deshui Yu committed
192
        });
193

Deshui Yu's avatar
Deshui Yu committed
194
195
196
        return this.experimentProfile.id;
    }

SparkSnail's avatar
SparkSnail committed
197
    public async resumeExperiment(readonly: boolean): Promise<void> {
chicm-ms's avatar
chicm-ms committed
198
        this.log.info(`Resuming experiment: ${this.experimentProfile.id}`);
Deshui Yu's avatar
Deshui Yu committed
199
200
201
        //Fetch back the experiment profile
        const experimentId: string = getExperimentId();
        this.experimentProfile = await this.dataStore.getExperimentProfile(experimentId);
SparkSnail's avatar
SparkSnail committed
202
203
204
205
        this.readonly = readonly;
        if (readonly) {
            return Promise.resolve();
        }
Deshui Yu's avatar
Deshui Yu committed
206
        const expParams: ExperimentParams = this.experimentProfile.params;
207
208
        setInitTrialSequenceId(this.experimentProfile.maxSequenceId + 1);

209
        // Set up multiphase config
210
        if (expParams.multiPhase && this.trainingService.isMultiPhaseJobSupported) {
211
212
213
            this.trainingService.setClusterMetadata('multiPhase', expParams.multiPhase.toString());
        }

214
215
        // Set up versionCheck config
        if (expParams.versionCheck !== undefined) {
SparkSnail's avatar
SparkSnail committed
216
            this.trainingService.setClusterMetadata('version_check', expParams.versionCheck.toString());
217
218
        }

QuanluZhang's avatar
QuanluZhang committed
219
        const dispatcherCommand: string = getMsgDispatcherCommand(expParams.tuner, expParams.assessor, expParams.advisor,
goooxu's avatar
goooxu committed
220
            expParams.multiPhase, expParams.multiThread);
221
        this.log.debug(`dispatcher command: ${dispatcherCommand}`);
QuanluZhang's avatar
QuanluZhang committed
222
        const checkpointDir: string = await this.createCheckpointDir();
Deshui Yu's avatar
Deshui Yu committed
223
        this.setupTuner(
224
225
            dispatcherCommand,
            undefined,
Deshui Yu's avatar
Deshui Yu committed
226
            'resume',
QuanluZhang's avatar
QuanluZhang committed
227
            checkpointDir);
Deshui Yu's avatar
Deshui Yu committed
228
229
230
231
232
233
234
235
236
237
238

        const allTrialJobs: TrialJobInfo[] = await this.dataStore.listTrialJobs();

        // Resume currSubmittedTrialNum
        this.currSubmittedTrialNum = allTrialJobs.length;

        // Check the final status for WAITING and RUNNING jobs
        await Promise.all(allTrialJobs
            .filter((job: TrialJobInfo) => job.status === 'WAITING' || job.status === 'RUNNING')
            .map((job: TrialJobInfo) => this.dataStore.storeTrialJobEvent('FAILED', job.id)));

239
240
241
242
243
244
245
246
247
248
        // Collect generated trials and imported trials
        const finishedTrialData: string = await this.exportData();
        const importedData: string[] = await this.dataStore.getImportedData();
        let trialData: Object[] = JSON.parse(finishedTrialData);
        for (const oneImportedData of importedData) {
            // do not deduplicate
            trialData = trialData.concat(<Object[]>JSON.parse(oneImportedData));
        }
        this.trialDataForTuner = JSON.stringify(trialData);

chicm-ms's avatar
chicm-ms committed
249
250
251
252
253
        if (this.experimentProfile.execDuration < this.experimentProfile.params.maxExecDuration &&
            this.currSubmittedTrialNum < this.experimentProfile.params.maxTrialNum &&
            this.experimentProfile.endTime) {
            delete this.experimentProfile.endTime;
        }
chicm-ms's avatar
chicm-ms committed
254
        this.setStatus('RUNNING');
255

Deshui Yu's avatar
Deshui Yu committed
256
        // TO DO: update database record for resume event
257
258
259
        this.run().catch((err: Error) => {
            this.criticalError(err);
        });
Deshui Yu's avatar
Deshui Yu committed
260
261
    }

262
263
    public getTrialJob(trialJobId: string): Promise<TrialJobInfo> {
        return this.dataStore.getTrialJob(trialJobId);
Deshui Yu's avatar
Deshui Yu committed
264
265
266
    }

    public async setClusterMetadata(key: string, value: string): Promise<void> {
SparkSnail's avatar
SparkSnail committed
267
268
269
        if (this.readonly) {
            return Promise.reject(new Error('Error: can not set cluster metadata in readonly mode!'));
        }
chicm-ms's avatar
chicm-ms committed
270
        this.log.info(`NNIManager setClusterMetadata, key: ${key}, value: ${value}`);
Deshui Yu's avatar
Deshui Yu committed
271
272
273
274
        let timeoutId: NodeJS.Timer;
        // TO DO: move timeout value to constants file
        const delay1: Promise<{}> = new Promise((resolve: Function, reject: Function): void => {
            timeoutId = setTimeout(
275
                () => { reject(new Error('TrainingService setClusterMetadata timeout. Please check your config file.')); },
Deshui Yu's avatar
Deshui Yu committed
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
                10000);
        });
        await Promise.race([delay1, this.trainingService.setClusterMetadata(key, value)]).finally(() => {
            clearTimeout(timeoutId);
        });
    }

    public getClusterMetadata(key: string): Promise<string> {
        return Promise.resolve(
            this.trainingService.getClusterMetadata(key)
        );
    }

    public async getTrialJobStatistics(): Promise<TrialJobStatistics[]> {
        return this.dataStore.getTrialJobStatistics();
    }

293
    public async stopExperiment(): Promise<void> {
chicm-ms's avatar
chicm-ms committed
294
295
        this.setStatus('STOPPING');
        this.log.info('Stopping experiment, cleaning up ...');
296
        await this.experimentDoneCleanUp();
chicm-ms's avatar
chicm-ms committed
297
        this.log.info('Experiment stopped.');
Deshui Yu's avatar
Deshui Yu committed
298
299
    }

300
    public async getMetricData(trialJobId?: string, metricType?: MetricType): Promise<MetricDataRecord[]> {
Deshui Yu's avatar
Deshui Yu committed
301
302
303
304
305
306
307
308
309
310
311
        return this.dataStore.getMetricData(trialJobId, metricType);
    }

    public getExperimentProfile(): Promise<ExperimentProfile> {
        // TO DO: using Promise.resolve()
        const deferred: Deferred<ExperimentProfile> = new Deferred<ExperimentProfile>();
        deferred.resolve(this.experimentProfile);

        return deferred.promise;
    }

312
313
314
315
    public getStatus(): NNIManagerStatus {
        return this.status;
    }

Deshui Yu's avatar
Deshui Yu committed
316
317
318
319
    public async listTrialJobs(status?: TrialJobStatus): Promise<TrialJobInfo[]> {
        return this.dataStore.listTrialJobs(status);
    }

320
321
    private setupTuner(command: string, cwd: string | undefined, mode: 'start' | 'resume', dataDirectory: string): void {
        if (this.dispatcher !== undefined) {
Deshui Yu's avatar
Deshui Yu committed
322
323
            return;
        }
goooxu's avatar
goooxu committed
324
        const stdio: StdioOptions = ['ignore', process.stdout, process.stderr, 'pipe', 'pipe'];
Deshui Yu's avatar
Deshui Yu committed
325
326
327
328
329
330
331
        let newCwd: string;
        if (cwd === undefined || cwd === '') {
            newCwd = getLogDir();
        } else {
            newCwd = cwd;
        }
        // TO DO: add CUDA_VISIBLE_DEVICES
332
333
334
335
336
        let includeIntermediateResultsEnv: boolean | undefined = false;
        if (this.experimentProfile.params.tuner !== undefined) {
            includeIntermediateResultsEnv = this.experimentProfile.params.tuner.includeIntermediateResults;
        }

Zejun Lin's avatar
Zejun Lin committed
337
338
339
        let nniEnv = {
            NNI_MODE: mode,
            NNI_CHECKPOINT_DIRECTORY: dataDirectory,
340
            NNI_LOG_DIRECTORY: getLogDir(),
341
342
            NNI_LOG_LEVEL: getLogLevel(),
            NNI_INCLUDE_INTERMEDIATE_RESULTS: includeIntermediateResultsEnv
Zejun Lin's avatar
Zejun Lin committed
343
344
        };
        let newEnv = Object.assign({}, process.env, nniEnv);
345
        const tunerProc: ChildProcess = getTunerProc(command,stdio,newCwd,newEnv);
346
347
        this.dispatcherPid = tunerProc.pid;
        this.dispatcher = createDispatcherInterface(tunerProc);
Deshui Yu's avatar
Deshui Yu committed
348
349
350
351
352

        return;
    }

    private updateTrialConcurrency(trialConcurrency: number): void {
QuanluZhang's avatar
QuanluZhang committed
353
354
        // we assume trialConcurrency >= 0, which is checked by restserver
        this.trialConcurrencyChange += (trialConcurrency - this.experimentProfile.params.trialConcurrency);
Deshui Yu's avatar
Deshui Yu committed
355
356
357
358
359
360
361
362
363
364
365
366
        this.experimentProfile.params.trialConcurrency = trialConcurrency;

        return;
    }

    private updateMaxExecDuration(duration: number): void {
        this.experimentProfile.params.maxExecDuration = duration;

        return;
    }

    private updateSearchSpace(searchSpace: string): void {
367
        if (this.dispatcher === undefined) {
Deshui Yu's avatar
Deshui Yu committed
368
369
            throw new Error('Error: tuner has not been setup');
        }
370
        this.dispatcher.sendCommand(UPDATE_SEARCH_SPACE, searchSpace);
Deshui Yu's avatar
Deshui Yu committed
371
372
373
374
375
        this.experimentProfile.params.searchSpace = searchSpace;

        return;
    }

QuanluZhang's avatar
QuanluZhang committed
376
377
378
379
380
381
    private updateMaxTrialNum(maxTrialNum: number): void {
        this.experimentProfile.params.maxTrialNum = maxTrialNum;

        return;
    }

Deshui Yu's avatar
Deshui Yu committed
382
    private async experimentDoneCleanUp(): Promise<void> {
383
        if (this.dispatcher === undefined) {
Deshui Yu's avatar
Deshui Yu committed
384
385
            throw new Error('Error: tuner has not been setup');
        }
386
        this.trainingService.removeTrialJobMetricListener(this.trialJobMetricListener); 
387
        this.dispatcher.sendCommand(TERMINATE);
Deshui Yu's avatar
Deshui Yu committed
388
389
390
        let tunerAlive: boolean = true;
        // gracefully terminate tuner and assessor here, wait at most 30 seconds.
        for (let i: number = 0; i < 30; i++) {
391
            if (!tunerAlive) { break; }
392
            tunerAlive = await isAlive(this.dispatcherPid);
Deshui Yu's avatar
Deshui Yu committed
393
394
            await delay(1000);
        }
395
        await killPid(this.dispatcherPid);
Deshui Yu's avatar
Deshui Yu committed
396
397
398
399
400
401
        const trialJobList: TrialJobDetail[] = await this.trainingService.listTrialJobs();
        // TO DO: to promise all
        for (const trialJob of trialJobList) {
            if (trialJob.status === 'RUNNING' ||
                trialJob.status === 'WAITING') {
                try {
chicm-ms's avatar
chicm-ms committed
402
                    this.log.info(`cancelTrialJob: ${trialJob.id}`);
Deshui Yu's avatar
Deshui Yu committed
403
404
405
406
407
408
409
                    await this.trainingService.cancelTrialJob(trialJob.id);
                } catch (error) {
                    // pid does not exist, do nothing here
                }
            }
        }
        await this.trainingService.cleanUp();
410
        this.experimentProfile.endTime = Date.now();
Deshui Yu's avatar
Deshui Yu committed
411
        await this.storeExperimentProfile();
chicm-ms's avatar
chicm-ms committed
412
        this.setStatus('STOPPED');
Deshui Yu's avatar
Deshui Yu committed
413
414
415
    }

    private async periodicallyUpdateExecDuration(): Promise<void> {
416
        let count: number = 1;
417
        while (!['ERROR', 'STOPPING', 'STOPPED'].includes(this.status.status)) {
418
            await delay(1000 * 1); // 1 seconds
419
            if (this.status.status === 'RUNNING') {
420
421
422
423
424
425
                this.experimentProfile.execDuration += 1;
                if (count % 10 === 0) {
                    await this.storeExperimentProfile();
                }
            }
            count += 1;
Deshui Yu's avatar
Deshui Yu committed
426
427
428
        }
    }

chicm-ms's avatar
chicm-ms committed
429
430
431
432
433
434
    private async pingDispatcher(): Promise<void> {
        if (this.dispatcher === undefined) {
            throw new Error('Error: tuner has not been setup');
        }
        while (!['ERROR', 'STOPPING', 'STOPPED'].includes(this.status.status)) {
            this.dispatcher.sendCommand(PING);
chicm-ms's avatar
chicm-ms committed
435
            await delay(1000 * 5);
chicm-ms's avatar
chicm-ms committed
436
437
438
        }
    }

QuanluZhang's avatar
QuanluZhang committed
439
440
    private async requestTrialJobsStatus(): Promise<number> {
        let finishedTrialJobNum: number = 0;
QuanluZhang's avatar
QuanluZhang committed
441
442
443
        if (this.dispatcher === undefined) {
            throw new Error('Error: tuner has not been setup');
        }
QuanluZhang's avatar
QuanluZhang committed
444
445
446
447
        for (const trialJobId of Array.from(this.trialJobs.keys())) {
            const trialJobDetail: TrialJobDetail = await this.trainingService.getTrialJob(trialJobId);
            const oldTrialJobDetail: TrialJobDetail | undefined = this.trialJobs.get(trialJobId);
            if (oldTrialJobDetail !== undefined && oldTrialJobDetail.status !== trialJobDetail.status) {
chicm-ms's avatar
chicm-ms committed
448
                this.log.info(`Trial job ${trialJobDetail.id} status changed from ${oldTrialJobDetail.status} to ${trialJobDetail.status}`);
QuanluZhang's avatar
QuanluZhang committed
449
                this.trialJobs.set(trialJobId, Object.assign({}, trialJobDetail));
450
                await this.dataStore.storeTrialJobEvent(trialJobDetail.status, trialJobDetail.id, undefined, trialJobDetail);
QuanluZhang's avatar
QuanluZhang committed
451
            }
QuanluZhang's avatar
QuanluZhang committed
452
            let hyperParams: string | undefined = undefined;
QuanluZhang's avatar
QuanluZhang committed
453
454
455
            switch (trialJobDetail.status) {
                case 'SUCCEEDED':
                case 'USER_CANCELED':
QuanluZhang's avatar
QuanluZhang committed
456
                case 'EARLY_STOPPED':
QuanluZhang's avatar
QuanluZhang committed
457
458
                    this.trialJobs.delete(trialJobId);
                    finishedTrialJobNum++;
QuanluZhang's avatar
QuanluZhang committed
459
460
461
462
463
464
465
466
                    if (trialJobDetail.form.jobType === 'TRIAL') {
                        hyperParams = (<TrialJobApplicationForm>trialJobDetail.form).hyperParameters.value;
                    } else {
                        throw new Error('Error: jobType error, not TRIAL');
                    }
                    this.dispatcher.sendCommand(TRIAL_END, JSON.stringify({
                        trial_job_id: trialJobDetail.id,
                        event: trialJobDetail.status,
goooxu's avatar
goooxu committed
467
468
                        hyper_params: hyperParams
                    }));
QuanluZhang's avatar
QuanluZhang committed
469
470
471
472
473
474
475
                    break;
                case 'FAILED':
                case 'SYS_CANCELED':
                    // In the current version, we do not retry
                    // TO DO: push this job to queue for retry
                    this.trialJobs.delete(trialJobId);
                    finishedTrialJobNum++;
QuanluZhang's avatar
QuanluZhang committed
476
477
478
479
480
481
482
483
                    if (trialJobDetail.form.jobType === 'TRIAL') {
                        hyperParams = (<TrialJobApplicationForm>trialJobDetail.form).hyperParameters.value;
                    } else {
                        throw new Error('Error: jobType error, not TRIAL');
                    }
                    this.dispatcher.sendCommand(TRIAL_END, JSON.stringify({
                        trial_job_id: trialJobDetail.id,
                        event: trialJobDetail.status,
goooxu's avatar
goooxu committed
484
485
                        hyper_params: hyperParams
                    }));
QuanluZhang's avatar
QuanluZhang committed
486
487
488
489
490
491
492
493
494
495
                    break;
                case 'WAITING':
                case 'RUNNING':
                case 'UNKNOWN':
                    // Do nothing
                    break;
                default:
                // TO DO: add warning in log
            }
        }
goooxu's avatar
goooxu committed
496

Gems Guo's avatar
Gems Guo committed
497
        return finishedTrialJobNum;
QuanluZhang's avatar
QuanluZhang committed
498
499
500
501
502
503
    }

    private async manageTrials(): Promise<void> {
        if (this.dispatcher === undefined) {
            throw new Error('Error: tuner has not been setup');
        }
QuanluZhang's avatar
QuanluZhang committed
504
        let allFinishedTrialJobNum: number = this.currSubmittedTrialNum;
QuanluZhang's avatar
QuanluZhang committed
505
        let waitSubmittedToFinish: number;
506
        while (!['ERROR', 'STOPPING', 'STOPPED'].includes(this.status.status)) {
QuanluZhang's avatar
QuanluZhang committed
507
508
509
510
511
512
513
            const finishedTrialJobNum: number = await this.requestTrialJobsStatus();
            allFinishedTrialJobNum += finishedTrialJobNum;

            // requestTrialNum is the number of trials that will be requested from tuner.
            // If trialConcurrency does not change, requestTrialNum equals finishedTrialJobNum.
            // If trialConcurrency changes, for example, trialConcurrency increases by 2 (trialConcurrencyChange=2), then
            // requestTrialNum equals 2 + finishedTrialJobNum and trialConcurrencyChange becomes 0.
514
            // If trialConcurrency changes, for example, trialConcurrency decreases by 4 (trialConcurrencyChange=-4) and
QuanluZhang's avatar
QuanluZhang committed
515
516
517
518
519
520
521
522
            // finishedTrialJobNum is 2, then requestTrialNum becomes -2. No trial will be requested from tuner,
            // and trialConcurrencyChange becomes -2.
            const requestTrialNum: number = this.trialConcurrencyChange + finishedTrialJobNum;
            if (requestTrialNum >= 0) {
                this.trialConcurrencyChange = 0;
            } else {
                this.trialConcurrencyChange = requestTrialNum;
            }
chicm-ms's avatar
chicm-ms committed
523
524
525

            const requestCustomTrialNum: number = Math.min(requestTrialNum, this.customizedTrials.length);
            for (let i: number = 0; i < requestCustomTrialNum; i++) {
QuanluZhang's avatar
QuanluZhang committed
526
527
528
529
530
531
532
                // ask tuner for more trials
                if (this.customizedTrials.length > 0) {
                    const hyperParams: string | undefined = this.customizedTrials.shift();
                    this.dispatcher.sendCommand(ADD_CUSTOMIZED_TRIAL_JOB, hyperParams);
                }
            }

chicm-ms's avatar
chicm-ms committed
533
534
535
536
            if (requestTrialNum - requestCustomTrialNum > 0) {
                this.requestTrialJobs(requestTrialNum - requestCustomTrialNum);
            }

QuanluZhang's avatar
QuanluZhang committed
537
            // check maxtrialnum and maxduration here
538
            // NO_MORE_TRIAL is more like a subset of RUNNING, because during RUNNING tuner
539
            // might tell nnimanager that this is no more trials. In NO_MORE_TRIAL state, the experiment is viewed
540
541
            // as still running. DONE could be transfered from RUNNING or NO_MORE_TRIAL.
            assert(this.status.status === 'RUNNING' ||
542
                this.status.status === 'DONE' ||
QuanluZhang's avatar
QuanluZhang committed
543
544
                this.status.status === 'NO_MORE_TRIAL' ||
                this.status.status === 'TUNER_NO_MORE_TRIAL');
545
            if (this.experimentProfile.execDuration > this.experimentProfile.params.maxExecDuration ||
QuanluZhang's avatar
QuanluZhang committed
546
                this.currSubmittedTrialNum >= this.experimentProfile.params.maxTrialNum) {
QuanluZhang's avatar
QuanluZhang committed
547
                if (this.status.status !== 'DONE') {
chicm-ms's avatar
chicm-ms committed
548
                    this.setStatus('NO_MORE_TRIAL');
QuanluZhang's avatar
QuanluZhang committed
549
550
551
552
                    waitSubmittedToFinish = this.currSubmittedTrialNum;

                    assert(allFinishedTrialJobNum <= waitSubmittedToFinish);
                    if (allFinishedTrialJobNum >= waitSubmittedToFinish) {
chicm-ms's avatar
chicm-ms committed
553
                        this.setStatus('DONE');
QuanluZhang's avatar
QuanluZhang committed
554
555
556
557
558
                        this.experimentProfile.endTime = Date.now();
                        await this.storeExperimentProfile();
                        // write this log for travis CI
                        this.log.info('Experiment done.');
                    }
QuanluZhang's avatar
QuanluZhang committed
559
560
561
                }
            } else {
                if (this.status.status === 'DONE') {
562
563
                    delete this.experimentProfile.endTime;
                    await this.storeExperimentProfile();
QuanluZhang's avatar
QuanluZhang committed
564
                }
QuanluZhang's avatar
QuanluZhang committed
565
                if (this.status.status !== 'TUNER_NO_MORE_TRIAL') {
chicm-ms's avatar
chicm-ms committed
566
                    this.setStatus('RUNNING');
567
                }
QuanluZhang's avatar
QuanluZhang committed
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
                for (let i: number = this.trialJobs.size; i < this.experimentProfile.params.trialConcurrency; i++) {
                    if (this.waitingTrials.length === 0 ||
                        this.currSubmittedTrialNum >= this.experimentProfile.params.maxTrialNum) {
                        break;
                    }
                    const hyperParams: string | undefined = this.waitingTrials.shift();
                    if (hyperParams === undefined) {
                        throw new Error(`Error: invalid hyper-parameters for job submission: ${hyperParams}`);
                    }
                    this.currSubmittedTrialNum++;
                    const trialJobAppForm: TrialJobApplicationForm = {
                        jobType: 'TRIAL',
                        hyperParameters: {
                            value: hyperParams,
                            index: 0
                        }
                    };
chicm-ms's avatar
chicm-ms committed
585
                    this.log.info(`submitTrialJob: form: ${JSON.stringify(trialJobAppForm)}`);
QuanluZhang's avatar
QuanluZhang committed
586
                    const trialJobDetail: TrialJobDetail = await this.trainingService.submitTrialJob(trialJobAppForm);
587
                    await this.storeMaxSequenceId(trialJobDetail.sequenceId);
QuanluZhang's avatar
QuanluZhang committed
588
589
590
591
                    this.trialJobs.set(trialJobDetail.id, Object.assign({}, trialJobDetail));
                    const trialJobDetailSnapshot: TrialJobDetail | undefined = this.trialJobs.get(trialJobDetail.id);
                    if (trialJobDetailSnapshot != undefined) {
                        await this.dataStore.storeTrialJobEvent(
592
                            trialJobDetailSnapshot.status, trialJobDetailSnapshot.id, hyperParams, trialJobDetailSnapshot);
QuanluZhang's avatar
QuanluZhang committed
593
594
595
596
597
598
599
600
601
                    } else {
                        assert(false, `undefined trialJobDetail in trialJobs: ${trialJobDetail.id}`);
                    }
                }
            }
            await delay(1000 * 5); // 5 seconds
        }
    }

Deshui Yu's avatar
Deshui Yu committed
602
603
604
605
606
607
    private storeExperimentProfile(): Promise<void> {
        this.experimentProfile.revision += 1;

        return this.dataStore.storeExperimentProfile(this.experimentProfile);
    }

608
    private async run(): Promise<void> {
QuanluZhang's avatar
QuanluZhang committed
609
        assert(this.dispatcher !== undefined);
610
611
612
613
614
615
616

        this.addEventListeners();

        this.sendInitTunerCommands();

        await Promise.all([
            this.periodicallyUpdateExecDuration(),
chicm-ms's avatar
chicm-ms committed
617
            this.pingDispatcher().catch((err: Error) => {
chicm-ms's avatar
chicm-ms committed
618
                throw NNIError.FromError(err, 'Dispatcher error: ');
chicm-ms's avatar
chicm-ms committed
619
            }),
620
            this.trainingService.run().catch((err: Error) => {
chicm-ms's avatar
chicm-ms committed
621
                throw NNIError.FromError(err, 'Training service error: ');
622
            }),
QuanluZhang's avatar
QuanluZhang committed
623
            this.manageTrials().catch((err: Error) => {
chicm-ms's avatar
chicm-ms committed
624
                throw NNIError.FromError(err, 'Job management error: ');
625
            })]);
626
627
    }

QuanluZhang's avatar
QuanluZhang committed
628
    private addEventListeners(): void {
chicm-ms's avatar
chicm-ms committed
629
        this.log.info('Add event listeners');
630
        // TO DO: cannot run this method more than once in one NNIManager instance
QuanluZhang's avatar
QuanluZhang committed
631
        if (this.dispatcher === undefined) {
632
633
            throw new Error('Error: tuner or job maintainer have not been setup');
        }
634
        this.trainingService.addTrialJobMetricListener(this.trialJobMetricListener);
635
636
637

        this.dispatcher.onCommand((commandType: string, content: string) => {
            this.onTunerCommand(commandType, content).catch((err: Error) => {
chicm-ms's avatar
chicm-ms committed
638
                this.criticalError(NNIError.FromError(err, 'Tuner command event error: '));
639
640
641
642
643
644
            });
        });
    }

    private sendInitTunerCommands(): void {
        if (this.dispatcher === undefined) {
645
            throw new Error('Dispatcher error: tuner has not been setup');
646
        }
chicm-ms's avatar
chicm-ms committed
647
648
649
        this.log.debug(`Send tuner command: INITIALIZE: ${this.experimentProfile.params.searchSpace}`);
        // Tuner need to be initialized with search space before generating any hyper parameters
        this.dispatcher.sendCommand(INITIALIZE, this.experimentProfile.params.searchSpace);
650
651
652
    }

    private async onTrialJobMetrics(metric: TrialJobMetric): Promise<void> {
chicm-ms's avatar
chicm-ms committed
653
        this.log.debug(`NNIManager received trial job metrics: ${metric}`);
654
655
656
657
658
659
660
        await this.dataStore.storeMetricData(metric.id, metric.data);
        if (this.dispatcher === undefined) {
            throw new Error('Error: tuner has not been setup');
        }
        this.dispatcher.sendCommand(REPORT_METRIC_DATA, metric.data);
    }

chicm-ms's avatar
chicm-ms committed
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
    private requestTrialJobs(jobNum: number): void {
        if (jobNum < 1) {
            return;
        }
        if (this.dispatcher === undefined) {
            throw new Error('Dispatcher error: tuner has not been setup');
        }
        if (this.experimentProfile.params.multiThread) {
            // Send multiple requests to ensure multiple hyper parameters are generated in non-blocking way.
            // For a single REQUEST_TRIAL_JOBS request, hyper parameters are generated one by one
            // sequentially.
            for (let i: number = 0; i < jobNum; i++) {
                this.dispatcher.sendCommand(REQUEST_TRIAL_JOBS, '1');
            }
        } else {
            this.dispatcher.sendCommand(REQUEST_TRIAL_JOBS, String(jobNum));
        }
    }

680
    private async onTunerCommand(commandType: string, content: string): Promise<void> {
horizon365's avatar
horizon365 committed
681
        this.log.info(`NNIManager received command from dispatcher: ${commandType}, ${content}`);
682
        switch (commandType) {
chicm-ms's avatar
chicm-ms committed
683
684
            case INITIALIZED:
                // Tuner is intialized, search space is set, request tuner to generate hyper parameters
685
686
687
688
689
690
                if (this.trialDataForTuner.length > 0) {
                    if (this.dispatcher === undefined) {
                        throw new Error('Dispatcher error: tuner has not been setup');
                    }
                    this.dispatcher.sendCommand(IMPORT_DATA, this.trialDataForTuner);
                }
chicm-ms's avatar
chicm-ms committed
691
692
                this.requestTrialJobs(this.experimentProfile.params.trialConcurrency);
                break;
693
            case NEW_TRIAL_JOB:
QuanluZhang's avatar
QuanluZhang committed
694
                if (this.status.status === 'TUNER_NO_MORE_TRIAL') {
695
                    this.log.warning('It is not supposed to receive more trials after NO_MORE_TRIAL is set');
chicm-ms's avatar
chicm-ms committed
696
                    this.setStatus('RUNNING');
697
                }
QuanluZhang's avatar
QuanluZhang committed
698
                this.waitingTrials.push(content);
699
                break;
chicm-ms's avatar
chicm-ms committed
700
701
702
703
704
705
706
707
708
709
710
711
            case SEND_TRIAL_JOB_PARAMETER:
                const tunerCommand: any = JSON.parse(content);
                assert(tunerCommand.parameter_index >= 0);
                assert(tunerCommand.trial_job_id !== undefined);

                const trialJobForm: TrialJobApplicationForm = {
                    jobType: 'TRIAL',
                    hyperParameters: {
                        value: content,
                        index: tunerCommand.parameter_index
                    }
                };
chicm-ms's avatar
chicm-ms committed
712
                this.log.info(`updateTrialJob: job id: ${tunerCommand.trial_job_id}, form: ${JSON.stringify(trialJobForm)}`);
chicm-ms's avatar
chicm-ms committed
713
                await this.trainingService.updateTrialJob(tunerCommand.trial_job_id, trialJobForm);
714
715
716
717
718
                if (tunerCommand['parameters'] !== null) {
                    // parameters field is set as empty string if no more hyper parameter can be generated by tuner.
                    await this.dataStore.storeTrialJobEvent(
                        'ADD_HYPERPARAMETER', tunerCommand.trial_job_id, content, undefined);
                }
chicm-ms's avatar
chicm-ms committed
719
                break;
720
            case NO_MORE_TRIAL_JOBS:
721
722
723
                if (!['ERROR', 'STOPPING', 'STOPPED'].includes(this.status.status)) {
                    this.setStatus('TUNER_NO_MORE_TRIAL');
                }
724
725
                break;
            case KILL_TRIAL_JOB:
chicm-ms's avatar
chicm-ms committed
726
                this.log.info(`cancelTrialJob: ${JSON.parse(content)}`);
QuanluZhang's avatar
QuanluZhang committed
727
                await this.trainingService.cancelTrialJob(JSON.parse(content), true);
728
729
730
731
                break;
            default:
                throw new Error('Error: unsupported command type from tuner');
        }
Deshui Yu's avatar
Deshui Yu committed
732
733
    }

734
735
736
737
738
739
740
741
742
743
    private criticalError(err: Error): void {
        this.logError(err);
        console.error(err);
    }

    private logError(err: Error): void {
        if (err.stack !== undefined) {
            this.log.error(err.stack);
        }
        this.status.errors.push(err.message);
chicm-ms's avatar
chicm-ms committed
744
745
746
747
748
749
750
751
        this.setStatus('ERROR');
    }

    private setStatus(status: ExperimentStatus): void {
        if (status !== this.status.status) {
            this.log.info(`Change NNIManager status from: ${this.status.status} to: ${status}`);
            this.status.status = status;
        }
752
753
754
755
756
757
758
    }

    private createEmptyExperimentProfile(): ExperimentProfile {
        return {
            id: getExperimentId(),
            revision: 0,
            execDuration: 0,
759
            logDir: getExperimentRootDir(),
760
            maxSequenceId: 0,
761
762
763
764
765
766
            params: {
                authorName: '',
                experimentName: '',
                trialConcurrency: 0,
                maxExecDuration: 0, // unit: second
                maxTrialNum: 0, // maxTrialNum includes all the submitted trial jobs
767
                trainingServicePlatform: '',
QuanluZhang's avatar
QuanluZhang committed
768
                searchSpace: ''
769
770
            }
        };
Deshui Yu's avatar
Deshui Yu committed
771
    }
772

QuanluZhang's avatar
QuanluZhang committed
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
    private async createCheckpointDir(): Promise<string> {
        // TODO: test
        const chkpDir: string = getCheckpointDir();
        // create checkpoint directory
        await mkDirP(chkpDir);
        // assign this directory to exp profile's checkpointDir
        if (this.experimentProfile.params.advisor) {
            this.experimentProfile.params.advisor.checkpointDir = chkpDir;
        }
        if (this.experimentProfile.params.tuner) {
            this.experimentProfile.params.tuner.checkpointDir = chkpDir;
        }
        if (this.experimentProfile.params.assessor) {
            this.experimentProfile.params.assessor.checkpointDir = chkpDir;
        }

        return Promise.resolve(chkpDir);
    }
goooxu's avatar
goooxu committed
791

792
793
794
795
796
797
    private async storeMaxSequenceId(sequenceId: number): Promise<void> {
        if (sequenceId > this.experimentProfile.maxSequenceId) {
            this.experimentProfile.maxSequenceId = sequenceId;
            await this.storeExperimentProfile();
        }
    }
Deshui Yu's avatar
Deshui Yu committed
798
799
800
}

export { NNIManager };