localTrainingService.ts 26.9 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
24
25
26
/**
 * 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 cpp from 'child-process-promise';
import * as cp from 'child_process';
import { EventEmitter } from 'events';
import * as fs from 'fs';
import * as path from 'path';
import * as ts from 'tail-stream';
27
import * as tkill from 'tree-kill';
28
import { NNIError, NNIErrorNames } from '../../common/errors';
suiguoxin's avatar
suiguoxin committed
29
import { getExperimentId, getInitTrialSequenceId } from '../../common/experimentStartupInfo';
30
import { getLogger, Logger } from '../../common/log';
Deshui Yu's avatar
Deshui Yu committed
31
import {
32
    HostJobApplicationForm, HyperParameters, JobApplicationForm, TrainingService, TrialJobApplicationForm,
Deshui Yu's avatar
Deshui Yu committed
33
34
    TrialJobDetail, TrialJobMetric, TrialJobStatus
} from '../../common/trainingService';
35
36
37
import {
    delay, generateParamFileName, getExperimentRootDir, getJobCancelStatus, getNewLine, isAlive, uniqueString
} from '../../common/utils';
38
39
import { TrialConfig } from '../common/trialConfig';
import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
40
import { execMkdir, execNewFile, getScriptName, runScript, setEnvironmentVariable } from '../common/util';
41
import { GPUScheduler } from './gpuScheduler';
Deshui Yu's avatar
Deshui Yu committed
42
43
44
45
46
47
48
49

/**
 * 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
 */
50
// tslint:disable:newline-per-chained-call informative-docs
Deshui Yu's avatar
Deshui Yu committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
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, 8).toString(), 10);
    if (data.length < contentLength + 8) {
        return [false, '', '', data];
    }
    const content: string = data.slice(8, contentLength + 8).toString();
    const remain: Buffer = data.slice(contentLength + 8);

    return [true, commandType, content, remain];
}
65
// tslint:enable:newline-per-chained-call informative-docs
Deshui Yu's avatar
Deshui Yu committed
66
67
68
69
70
71
72

/**
 * LocalTrialJobDetail
 */
class LocalTrialJobDetail implements TrialJobDetail {
    public id: string;
    public status: TrialJobStatus;
73
74
75
    public submitTime: number;
    public startTime?: number;
    public endTime?: number;
Deshui Yu's avatar
Deshui Yu committed
76
77
78
79
    public tags?: string[];
    public url?: string;
    public workingDirectory: string;
    public form: JobApplicationForm;
80
    public sequenceId: number;
Deshui Yu's avatar
Deshui Yu committed
81
    public pid?: number;
82
    public gpuIndices?: number[];
Deshui Yu's avatar
Deshui Yu committed
83

84
85
    constructor(
        id: string, status: TrialJobStatus, submitTime: number,
86
        workingDirectory: string, form: JobApplicationForm, sequenceId: number) {
Deshui Yu's avatar
Deshui Yu committed
87
88
89
90
91
92
        this.id = id;
        this.status = status;
        this.submitTime = submitTime;
        this.workingDirectory = workingDirectory;
        this.form = form;
        this.url = `file://localhost:${workingDirectory}`;
93
        this.sequenceId = sequenceId;
94
95
96
97
98
99
100
101
        this.gpuIndices = [];
    }
}

/**
 * Local training service config
 */
class LocalConfig {
102
    public maxTrialNumPerGpu?: number;
103
    public gpuIndices?: string;
104
105
    public useActiveGpu?: boolean;
    constructor(gpuIndices?: string, maxTrialNumPerGpu?: number, useActiveGpu?: boolean) {
106
107
108
        if (gpuIndices !== undefined) {
            this.gpuIndices = gpuIndices;
        }
109
110
111
112
113
114
        if (maxTrialNumPerGpu !== undefined) {
            this.maxTrialNumPerGpu = maxTrialNumPerGpu;
        }
        if (useActiveGpu !== undefined) {
            this.useActiveGpu = useActiveGpu;
        }
Deshui Yu's avatar
Deshui Yu committed
115
116
117
118
    }
}

/**
chicm-ms's avatar
chicm-ms committed
119
 * Local machine training service
Deshui Yu's avatar
Deshui Yu committed
120
121
 */
class LocalTrainingService implements TrainingService {
122
123
124
    private readonly eventEmitter: EventEmitter;
    private readonly jobMap: Map<string, LocalTrialJobDetail>;
    private readonly jobQueue: string[];
Deshui Yu's avatar
Deshui Yu committed
125
126
127
    private initialized: boolean;
    private stopping: boolean;
    private rootDir!: string;
128
    private trialSequenceId: number;
suiguoxin's avatar
suiguoxin committed
129
    private readonly experimentId! : string;
130
    private gpuScheduler!: GPUScheduler;
131
    private readonly occupiedGpuIndexNumMap: Map<number, number>;
132
    private designatedGpuIndices!: Set<number>;
133
    private readonly log: Logger;
134
    private localTrialConfig?: TrialConfig;
135
    private localConfig?: LocalConfig;
136
    private isMultiPhase: boolean;
137
    private readonly jobStreamMap: Map<string, ts.Stream>;
138
139
    private maxTrialNumPerGpu: number;
    private useActiveGpu: boolean;
Deshui Yu's avatar
Deshui Yu committed
140
141
142
143
144
145
146
147

    constructor() {
        this.eventEmitter = new EventEmitter();
        this.jobMap = new Map<string, LocalTrialJobDetail>();
        this.jobQueue = [];
        this.initialized = false;
        this.stopping = false;
        this.log = getLogger();
148
        this.trialSequenceId = -1;
suiguoxin's avatar
suiguoxin committed
149
        this.experimentId = getExperimentId();
150
        this.jobStreamMap = new Map<string, ts.Stream>();
chicm-ms's avatar
chicm-ms committed
151
        this.log.info('Construct local machine training service.');
152
153
154
155
        this.occupiedGpuIndexNumMap = new Map<number, number>();
        this.maxTrialNumPerGpu = 1;
        this.useActiveGpu = false;
        this.isMultiPhase = false;
Deshui Yu's avatar
Deshui Yu committed
156
157
158
    }

    public async run(): Promise<void> {
chicm-ms's avatar
chicm-ms committed
159
        this.log.info('Run local machine training service.');
160
161
162
        const longRunningTasks: Promise<void>[] = [this.runJobLoop()];
        if (this.gpuScheduler !== undefined) {
            longRunningTasks.push(this.gpuScheduler.run());
Deshui Yu's avatar
Deshui Yu committed
163
        }
164
        await Promise.all(longRunningTasks);
chicm-ms's avatar
chicm-ms committed
165
        this.log.info('Local machine training service exit.');
Deshui Yu's avatar
Deshui Yu committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
    }

    public async listTrialJobs(): Promise<TrialJobDetail[]> {
        const jobs: TrialJobDetail[] = [];
        for (const key of this.jobMap.keys()) {
            const trialJob: TrialJobDetail = await this.getTrialJob(key);
            if (trialJob.form.jobType === 'TRIAL') {
                jobs.push(trialJob);
            }
        }

        return jobs;
    }

    public async getTrialJob(trialJobId: string): Promise<TrialJobDetail> {
        const trialJob: LocalTrialJobDetail | undefined = this.jobMap.get(trialJobId);
        if (trialJob === undefined) {
            throw new NNIError(NNIErrorNames.NOT_FOUND, 'Trial job not found');
        }
        if (trialJob.form.jobType === 'HOST') {
            return this.getHostJob(trialJobId);
        }
        if (trialJob.status === 'RUNNING') {
189
            const alive: boolean = await isAlive(trialJob.pid);
Deshui Yu's avatar
Deshui Yu committed
190
            if (!alive) {
191
                trialJob.endTime = Date.now();
Deshui Yu's avatar
Deshui Yu committed
192
193
194
                this.setTrialJobStatus(trialJob, 'FAILED');
                try {
                    const state: string = await fs.promises.readFile(path.join(trialJob.workingDirectory, '.nni', 'state'), 'utf8');
195
196
                    const match: RegExpMatchArray | null = state.trim()
                        .match(/^(\d+)\s+(\d+)/);
Deshui Yu's avatar
Deshui Yu committed
197
198
199
200
201
                    if (match !== null) {
                        const { 1: code, 2: timestamp } = match;
                        if (parseInt(code, 10) === 0) {
                            this.setTrialJobStatus(trialJob, 'SUCCEEDED');
                        }
202
                        trialJob.endTime = parseInt(timestamp, 10);
Deshui Yu's avatar
Deshui Yu committed
203
204
205
206
                    }
                } catch (error) {
                    //ignore
                }
207
                this.log.debug(`trialJob status update: ${trialJobId}, ${trialJob.status}`);
Deshui Yu's avatar
Deshui Yu committed
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
            }
        }

        return trialJob;
    }

    public addTrialJobMetricListener(listener: (metric: TrialJobMetric) => void): void {
        this.eventEmitter.on('metric', listener);
    }

    public removeTrialJobMetricListener(listener: (metric: TrialJobMetric) => void): void {
        this.eventEmitter.off('metric', listener);
    }

    public submitTrialJob(form: JobApplicationForm): Promise<TrialJobDetail> {
        if (form.jobType === 'HOST') {
            return this.runHostJob(<HostJobApplicationForm>form);
        } else if (form.jobType === 'TRIAL') {
            const trialJobId: string = uniqueString(5);
            const trialJobDetail: LocalTrialJobDetail = new LocalTrialJobDetail(
                trialJobId,
                'WAITING',
230
                Date.now(),
Deshui Yu's avatar
Deshui Yu committed
231
                path.join(this.rootDir, 'trials', trialJobId),
232
233
234
                form,
                this.generateSequenceId()
            );
Deshui Yu's avatar
Deshui Yu committed
235
236
237
238
239
240
241
242
243
244
245
            this.jobQueue.push(trialJobId);
            this.jobMap.set(trialJobId, trialJobDetail);

            this.log.debug(`submitTrialJob: return: ${JSON.stringify(trialJobDetail)} `);

            return Promise.resolve(trialJobDetail);
        } else {
            return Promise.reject(new Error(`Job form not supported: ${JSON.stringify(form)}`));
        }
    }

246
247
248
249
250
    /**
     * Update trial job for multi-phase
     * @param trialJobId trial job id
     * @param form job application form
     */
chicm-ms's avatar
chicm-ms committed
251
252
253
254
255
256
257
258
259
260
261
262
    public async updateTrialJob(trialJobId: string, form: JobApplicationForm): Promise<TrialJobDetail> {
        const trialJobDetail: undefined | TrialJobDetail = this.jobMap.get(trialJobId);
        if (trialJobDetail === undefined) {
            throw new Error(`updateTrialJob failed: ${trialJobId} not found`);
        }
        if (form.jobType === 'TRIAL') {
            await this.writeParameterFile(trialJobDetail.workingDirectory, (<TrialJobApplicationForm>form).hyperParameters);
        } else {
            throw new Error(`updateTrialJob failed: jobType ${form.jobType} not supported.`);
        }

        return trialJobDetail;
263
264
265
266
267
268
    }

    /**
     * Is multiphase job supported in current training service
     */
    public get isMultiPhaseJobSupported(): boolean {
269
        return true;
270
271
    }

QuanluZhang's avatar
QuanluZhang committed
272
    public async cancelTrialJob(trialJobId: string, isEarlyStopped: boolean = false): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
273
274
275
276
        const trialJob: LocalTrialJobDetail | undefined = this.jobMap.get(trialJobId);
        if (trialJob === undefined) {
            throw new NNIError(NNIErrorNames.NOT_FOUND, 'Trial job not found');
        }
277
        if (trialJob.pid === undefined) {
278
            this.setTrialJobStatus(trialJob, 'USER_CANCELED');
279

SparkSnail's avatar
SparkSnail committed
280
            return Promise.resolve();
281
        }
Deshui Yu's avatar
Deshui Yu committed
282
        if (trialJob.form.jobType === 'TRIAL') {
283
            tkill(trialJob.pid, 'SIGKILL');
Deshui Yu's avatar
Deshui Yu committed
284
285
286
287
288
        } else if (trialJob.form.jobType === 'HOST') {
            await cpp.exec(`pkill -9 -P ${trialJob.pid}`);
        } else {
            throw new Error(`Job type not supported: ${trialJob.form.jobType}`);
        }
QuanluZhang's avatar
QuanluZhang committed
289
        this.setTrialJobStatus(trialJob, getJobCancelStatus(isEarlyStopped));
290

SparkSnail's avatar
SparkSnail committed
291
        return Promise.resolve();
Deshui Yu's avatar
Deshui Yu committed
292
293
294
295
296
    }

    public async setClusterMetadata(key: string, value: string): Promise<void> {
        if (!this.initialized) {
            this.rootDir = getExperimentRootDir();
297
298
            // tslint:disable-next-line:non-literal-fs-path
            if (!fs.existsSync(this.rootDir)) {
299
300
                await cpp.exec(`powershell.exe mkdir ${this.rootDir}`);
            }
Deshui Yu's avatar
Deshui Yu committed
301
302
303
            this.initialized = true;
        }
        switch (key) {
304
            case TrialConfigMetadataKey.TRIAL_CONFIG:
305
                this.localTrialConfig = <TrialConfig>JSON.parse(value);
306
                // Parse trial config failed, throw Error
307
                if (this.localTrialConfig === undefined) {
308
309
                    throw new Error('trial config parsed failed');
                }
310
311
312
                if (this.localTrialConfig.gpuNum !== undefined) {
                    this.log.info(`required GPU number is ${this.localTrialConfig.gpuNum}`);
                    if (this.gpuScheduler === undefined && this.localTrialConfig.gpuNum > 0) {
SparkSnail's avatar
SparkSnail committed
313
314
                        this.gpuScheduler = new GPUScheduler();
                    }
315
316
317
318
319
320
321
322
323
324
325
326
                }
                break;
            case TrialConfigMetadataKey.LOCAL_CONFIG:
                this.localConfig = <LocalConfig>JSON.parse(value);
                this.log.info(`Specified GPU indices: ${this.localConfig.gpuIndices}`);
                if (this.localConfig.gpuIndices !== undefined) {
                    this.designatedGpuIndices = new Set(this.localConfig.gpuIndices.split(',')
                            .map((x: string) => parseInt(x, 10)));
                    if (this.designatedGpuIndices.size === 0) {
                        throw new Error('gpuIndices can not be empty if specified.');
                    }
                }
327
328
329
330
331
332
333
                if (this.localConfig.maxTrialNumPerGpu !== undefined) {
                    this.maxTrialNumPerGpu = this.localConfig.maxTrialNumPerGpu;
                }

                if (this.localConfig.useActiveGpu !== undefined) {
                    this.useActiveGpu = this.localConfig.useActiveGpu;
                }
Deshui Yu's avatar
Deshui Yu committed
334
                break;
335
336
337
            case TrialConfigMetadataKey.MULTI_PHASE:
                this.isMultiPhase = (value === 'true' || value === 'True');
                break;
Deshui Yu's avatar
Deshui Yu committed
338
339
340
341
342
343
            default:
        }
    }

    public getClusterMetadata(key: string): Promise<string> {
        switch (key) {
344
            case TrialConfigMetadataKey.TRIAL_CONFIG:
345
                let getResult: Promise<string>;
346
                if (this.localTrialConfig === undefined) {
347
348
                    getResult = Promise.reject(new NNIError(NNIErrorNames.NOT_FOUND, `${key} is never set yet`));
                } else {
349
                    getResult = Promise.resolve(JSON.stringify(this.localTrialConfig));
350
                }
351

352
                return getResult;
Deshui Yu's avatar
Deshui Yu committed
353
354
355
356
357
            default:
                return Promise.reject(new NNIError(NNIErrorNames.NOT_FOUND, 'Key not found'));
        }
    }

358
    public async cleanUp(): Promise<void> {
chicm-ms's avatar
chicm-ms committed
359
        this.log.info('Stopping local machine training service...');
Deshui Yu's avatar
Deshui Yu committed
360
        this.stopping = true;
361
        for (const stream of this.jobStreamMap.values()) {
362
363
            stream.end(0);
            stream.emit('end');
364
        }
365
366
367
368
        if (this.gpuScheduler !== undefined) {
            await this.gpuScheduler.stop();
        }

Deshui Yu's avatar
Deshui Yu committed
369
370
371
        return Promise.resolve();
    }

372
    private onTrialJobStatusChanged(trialJob: LocalTrialJobDetail, oldStatus: TrialJobStatus): void {
373
        //if job is not running, destory job stream
374
375
376
        if (['SUCCEEDED', 'FAILED', 'USER_CANCELED', 'SYS_CANCELED', 'EARLY_STOPPED'].includes(trialJob.status)) {
            if (this.jobStreamMap.has(trialJob.id)) {
                const stream: ts.Stream | undefined = this.jobStreamMap.get(trialJob.id);
377
                if (stream === undefined) {
378
379
                    throw new Error(`Could not find stream in trial ${trialJob.id}`);
                }
380
                //Refer https://github.com/Juul/tail-stream/issues/20
381
382
                stream.end(0);
                stream.emit('end');
383
384
385
                this.jobStreamMap.delete(trialJob.id);
            }
        }
386
387
388
        if (trialJob.gpuIndices !== undefined && trialJob.gpuIndices.length > 0 && this.gpuScheduler !== undefined) {
            if (oldStatus === 'RUNNING' && trialJob.status !== 'RUNNING') {
                for (const index of trialJob.gpuIndices) {
389
390
                    const num: number | undefined = this.occupiedGpuIndexNumMap.get(index);
                    if (num === undefined) {
391
                        throw new Error(`gpu resource schedule error`);
392
                    } else if (num === 1) {
393
394
                        this.occupiedGpuIndexNumMap.delete(index);
                    } else {
395
                        this.occupiedGpuIndexNumMap.set(index, num - 1);
396
                    }
397
398
399
                }
            }
        }
Deshui Yu's avatar
Deshui Yu committed
400
401
    }

402
403
    private getEnvironmentVariables(
        trialJobDetail: TrialJobDetail,
SparkSnail's avatar
SparkSnail committed
404
405
        resource: { gpuIndices: number[] },
        gpuNum: number | undefined): { key: string; value: string }[] {
406
        const envVariables: { key: string; value: string }[] = [
Deshui Yu's avatar
Deshui Yu committed
407
            { key: 'NNI_PLATFORM', value: 'local' },
suiguoxin's avatar
suiguoxin committed
408
            { key: 'NNI_EXP_ID', value: this.experimentId },
Deshui Yu's avatar
Deshui Yu committed
409
410
            { key: 'NNI_SYS_DIR', value: trialJobDetail.workingDirectory },
            { key: 'NNI_TRIAL_JOB_ID', value: trialJobDetail.id },
411
            { key: 'NNI_OUTPUT_DIR', value: trialJobDetail.workingDirectory },
412
            { key: 'NNI_TRIAL_SEQ_ID', value: trialJobDetail.sequenceId.toString() },
413
            { key: 'MULTI_PHASE', value: this.isMultiPhase.toString() }
Deshui Yu's avatar
Deshui Yu committed
414
        ];
SparkSnail's avatar
SparkSnail committed
415
416
417
418
419
420
        if (gpuNum !== undefined) {
            envVariables.push({
                key: 'CUDA_VISIBLE_DEVICES',
                value: this.gpuScheduler === undefined ? '-1' : resource.gpuIndices.join(',')
            });
        }
421
422
423
424
425
426
427
428
429

        return envVariables;
    }

    private setExtraProperties(trialJobDetail: LocalTrialJobDetail, resource: { gpuIndices: number[] }): void {
        trialJobDetail.gpuIndices = resource.gpuIndices;
    }

    private tryGetAvailableResource(): [boolean, { gpuIndices: number[]}] {
430
431
        if (this.localTrialConfig === undefined) {
            throw new Error('localTrialConfig is not initialized!');
432
433
434
435
436
437
438
        }

        const resource: { gpuIndices: number[] } = { gpuIndices: [] };
        if (this.gpuScheduler === undefined) {
            return [true, resource];
        }

439
        let selectedGPUIndices: number[] = [];
440
441
442
443
        const availableGpuIndices: number[] = this.gpuScheduler.getAvailableGPUIndices(this.useActiveGpu, this.occupiedGpuIndexNumMap);
        for (const index of availableGpuIndices) {
            const num: number | undefined = this.occupiedGpuIndexNumMap.get(index);
            if (num === undefined || num < this.maxTrialNumPerGpu) {
444
445
446
                selectedGPUIndices.push(index);
            }
        }
447
448
449
450
451
452

        if (this.designatedGpuIndices !== undefined) {
            this.checkSpecifiedGpuIndices();
            selectedGPUIndices = selectedGPUIndices.filter((index: number) => this.designatedGpuIndices.has(index));
        }

453
        if (selectedGPUIndices.length < this.localTrialConfig.gpuNum) {
454
455
456
            return [false, resource];
        }

457
        selectedGPUIndices.splice(this.localTrialConfig.gpuNum);
458
459
460
        Object.assign(resource, { gpuIndices: selectedGPUIndices });

        return [true, resource];
Deshui Yu's avatar
Deshui Yu committed
461
462
    }

463
464
465
466
467
468
469
470
471
    private checkSpecifiedGpuIndices(): void {
        const gpuCount: number = this.gpuScheduler.getSystemGpuCount();
        if (this.designatedGpuIndices !== undefined) {
            for (const index of this.designatedGpuIndices) {
                if (index >= gpuCount) {
                    throw new Error(`Specified GPU index not found: ${index}`);
                }
            }
        }
Deshui Yu's avatar
Deshui Yu committed
472
473
    }

474
475
476
    private occupyResource(resource: {gpuIndices: number[]}): void {
        if (this.gpuScheduler !== undefined) {
            for (const index of resource.gpuIndices) {
477
478
479
                const num: number | undefined = this.occupiedGpuIndexNumMap.get(index);
                if (num === undefined) {
                    this.occupiedGpuIndexNumMap.set(index, 1);
480
                } else {
481
                    this.occupiedGpuIndexNumMap.set(index, num + 1);
482
                }
483
484
            }
        }
Deshui Yu's avatar
Deshui Yu committed
485
486
    }

487
488
489
490
491
492
493
494
495
496
    private async runJobLoop(): Promise<void> {
        while (!this.stopping) {
            while (!this.stopping && this.jobQueue.length !== 0) {
                const trialJobId: string = this.jobQueue[0];
                const trialJobDeatil: LocalTrialJobDetail | undefined = this.jobMap.get(trialJobId);
                if (trialJobDeatil !== undefined && trialJobDeatil.status === 'WAITING') {
                    const [success, resource] = this.tryGetAvailableResource();
                    if (!success) {
                        break;
                    }
497

498
499
500
501
502
503
504
                    this.occupyResource(resource);
                    await this.runTrialJob(trialJobId, resource);
                }
                this.jobQueue.shift();
            }
            await delay(5000);
        }
Deshui Yu's avatar
Deshui Yu committed
505
506
507
508
509
510
511
512
513
514
    }

    private setTrialJobStatus(trialJob: LocalTrialJobDetail, newStatus: TrialJobStatus): void {
        if (trialJob.status !== newStatus) {
            const oldStatus: TrialJobStatus = trialJob.status;
            trialJob.status = newStatus;
            this.onTrialJobStatusChanged(trialJob, oldStatus);
        }
    }

515
    private getScript(localTrialConfig: TrialConfig, workingDirectory: string): string[] {
516
517
        const script: string[] = [];
        if (process.platform === 'win32') {
518
            script.push(
519
                `cmd /c ${localTrialConfig.command} 2>${path.join(workingDirectory, 'stderr')}`,
520
                `$NOW_DATE = [int64](([datetime]::UtcNow)-(get-date "1/1/1970")).TotalSeconds`,
521
                `$NOW_DATE = "$NOW_DATE" + (Get-Date -Format fff).ToString()`,
522
                `Write $LASTEXITCODE " " $NOW_DATE  | Out-File ${path.join(workingDirectory, '.nni', 'state')} -NoNewline -encoding utf8`);
523
        } else {
524
525
526
527
528
529
530
531
            script.push(`eval ${localTrialConfig.command} 2>${path.join(workingDirectory, 'stderr')}`);
            if (process.platform === 'darwin') {
                // https://superuser.com/questions/599072/how-to-get-bash-execution-time-in-milliseconds-under-mac-os-x
                // Considering the worst case, write 999 to avoid negative duration
                script.push(`echo $? \`date +%s999\` >${path.join(workingDirectory, '.nni', 'state')}`);
            } else {
                script.push(`echo $? \`date +%s%3N\` >${path.join(workingDirectory, '.nni', 'state')}`);
            }
532
        }
533

534
535
536
        return script;
    }

537
    private async runTrialJob(trialJobId: string, resource: {gpuIndices: number[]}): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
538
        const trialJobDetail: LocalTrialJobDetail = <LocalTrialJobDetail>this.jobMap.get(trialJobId);
539
        if (this.localTrialConfig === undefined) {
SparkSnail's avatar
SparkSnail committed
540
541
            throw new Error(`localTrialConfig not initialized!`);
        }
542
        const variables: { key: string; value: string }[] = this.getEnvironmentVariables(trialJobDetail, resource, this.localTrialConfig.gpuNum);
Deshui Yu's avatar
Deshui Yu committed
543

544
        if (this.localTrialConfig === undefined) {
545
546
            throw new Error('trial config is not initialized');
        }
547
548
549
        const runScriptContent: string[] = [];
        if (process.platform !== 'win32') {
            runScriptContent.push('#!/bin/bash');
550
        }
551
        runScriptContent.push(`cd ${this.localTrialConfig.codeDir}`);
Deshui Yu's avatar
Deshui Yu committed
552
        for (const variable of variables) {
553
            runScriptContent.push(setEnvironmentVariable(variable));
Deshui Yu's avatar
Deshui Yu committed
554
        }
555
        const scripts: string[] = this.getScript(this.localTrialConfig, trialJobDetail.workingDirectory);
556
557
        scripts.forEach((script: string) => {
            runScriptContent.push(script);
558
559
560
561
562
        });
        await execMkdir(trialJobDetail.workingDirectory);
        await execMkdir(path.join(trialJobDetail.workingDirectory, '.nni'));
        await execNewFile(path.join(trialJobDetail.workingDirectory, '.nni', 'metrics'));
        const scriptName: string = getScriptName('run');
563
564
        await fs.promises.writeFile(path.join(trialJobDetail.workingDirectory, scriptName),
                                    runScriptContent.join(getNewLine()), { encoding: 'utf8', mode: 0o777 });
chicm-ms's avatar
chicm-ms committed
565
        await this.writeParameterFile(trialJobDetail.workingDirectory, (<TrialJobApplicationForm>trialJobDetail.form).hyperParameters);
566
        const trialJobProcess: cp.ChildProcess = runScript(path.join(trialJobDetail.workingDirectory, scriptName));
Deshui Yu's avatar
Deshui Yu committed
567
        this.setTrialJobStatus(trialJobDetail, 'RUNNING');
568
        trialJobDetail.startTime = Date.now();
569
        trialJobDetail.pid = trialJobProcess.pid;
Deshui Yu's avatar
Deshui Yu committed
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
        this.setExtraProperties(trialJobDetail, resource);

        let buffer: Buffer = Buffer.alloc(0);
        const stream: ts.Stream = ts.createReadStream(path.join(trialJobDetail.workingDirectory, '.nni', 'metrics'));
        stream.on('data', (data: Buffer) => {
            buffer = Buffer.concat([buffer, data]);
            while (buffer.length > 0) {
                const [success, , content, remain] = decodeCommand(buffer);
                if (!success) {
                    break;
                }
                this.eventEmitter.emit('metric', {
                    id: trialJobDetail.id,
                    data: content
                });
                this.log.debug(`Sending metrics, job id: ${trialJobDetail.id}, metrics: ${content}`);
                buffer = remain;
            }
        });
589
        this.jobStreamMap.set(trialJobDetail.id, stream);
Deshui Yu's avatar
Deshui Yu committed
590
591
592
593
594
595
596
597
598
599
600
601
    }

    private async runHostJob(form: HostJobApplicationForm): Promise<TrialJobDetail> {
        const jobId: string = uniqueString(5);
        const workDir: string = path.join(this.rootDir, 'hostjobs', jobId);
        await cpp.exec(`mkdir -p ${workDir}`);
        const wrappedCmd: string = `cd ${workDir} && ${form.cmd}>stdout 2>stderr`;
        this.log.debug(`runHostJob: command: ${wrappedCmd}`);
        const process: cp.ChildProcess = cp.exec(wrappedCmd);
        const jobDetail: LocalTrialJobDetail = {
            id: jobId,
            status: 'RUNNING',
602
            submitTime: Date.now(),
Deshui Yu's avatar
Deshui Yu committed
603
604
            workingDirectory: workDir,
            form: form,
605
            sequenceId: this.generateSequenceId(),
Deshui Yu's avatar
Deshui Yu committed
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
            pid: process.pid
        };
        this.jobMap.set(jobId, jobDetail);
        this.log.debug(`runHostJob: return: ${JSON.stringify(jobDetail)} `);

        return jobDetail;
    }

    private async getHostJob(jobId: string): Promise<TrialJobDetail> {
        const jobDetail: LocalTrialJobDetail | undefined = this.jobMap.get(jobId);
        if (jobDetail === undefined) {
            throw new NNIError(NNIErrorNames.NOT_FOUND, `Host Job not found: ${jobId}`);
        }
        try {
            await cpp.exec(`kill -0 ${jobDetail.pid}`);

            return jobDetail;
        } catch (error) {
            if (error instanceof Error) {
                this.log.debug(`getHostJob: error: ${error.message}`);
                this.jobMap.delete(jobId);
                throw new NNIError(NNIErrorNames.NOT_FOUND, `Host Job not found: ${error.message}`);
            } else {
                throw error;
            }
        }
    }
chicm-ms's avatar
chicm-ms committed
633
634

    private async writeParameterFile(directory: string, hyperParameters: HyperParameters): Promise<void> {
635
        const filepath: string = path.join(directory, generateParamFileName(hyperParameters));
chicm-ms's avatar
chicm-ms committed
636
637
        await fs.promises.writeFile(filepath, hyperParameters.value, { encoding: 'utf8' });
    }
638
639

    private generateSequenceId(): number {
640
641
642
643
        if (this.trialSequenceId === -1) {
            this.trialSequenceId = getInitTrialSequenceId();
        }

644
645
        return this.trialSequenceId++;
    }
Deshui Yu's avatar
Deshui Yu committed
646
647
648
}

export { LocalTrainingService };