localTrainingService.ts 23.7 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
5
6
7
8
9
10

'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';
11
import * as tkill from 'tree-kill';
12
import { NNIError, NNIErrorNames } from '../../common/errors';
13
import { getExperimentId } from '../../common/experimentStartupInfo';
14
import { getLogger, Logger } from '../../common/log';
Deshui Yu's avatar
Deshui Yu committed
15
import {
16
    HyperParameters, TrainingService, TrialJobApplicationForm,
17
    TrialJobDetail, TrialJobMetric, TrialJobStatus, LogType
Deshui Yu's avatar
Deshui Yu committed
18
} from '../../common/trainingService';
19
20
21
import {
    delay, generateParamFileName, getExperimentRootDir, getJobCancelStatus, getNewLine, isAlive, uniqueString
} from '../../common/utils';
22
23
import { TrialConfig } from '../common/trialConfig';
import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
24
import { execMkdir, execNewFile, getScriptName, runScript, setEnvironmentVariable } from '../common/util';
25
import { GPUScheduler } from './gpuScheduler';
Deshui Yu's avatar
Deshui Yu committed
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

/**
 * Decode a command
 * @param Buffer binary incoming data
 * @returns a tuple of (success, commandType, content, remain)
 *          success: true if the buffer contains at least one complete command; otherwise false
 *          remain: remaining data after the first command
 */
function decodeCommand(data: Buffer): [boolean, string, string, Buffer] {
    if (data.length < 8) {
        return [false, '', '', data];
    }
    const commandType: string = data.slice(0, 2).toString();
    const contentLength: number = parseInt(data.slice(2, 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];
}

/**
 * LocalTrialJobDetail
 */
class LocalTrialJobDetail implements TrialJobDetail {
    public id: string;
    public status: TrialJobStatus;
55
56
57
    public submitTime: number;
    public startTime?: number;
    public endTime?: number;
Deshui Yu's avatar
Deshui Yu committed
58
59
60
    public tags?: string[];
    public url?: string;
    public workingDirectory: string;
61
    public form: TrialJobApplicationForm;
Deshui Yu's avatar
Deshui Yu committed
62
    public pid?: number;
63
    public gpuIndices?: number[];
Deshui Yu's avatar
Deshui Yu committed
64

65
66
    constructor(
        id: string, status: TrialJobStatus, submitTime: number,
67
        workingDirectory: string, form: TrialJobApplicationForm) {
Deshui Yu's avatar
Deshui Yu committed
68
69
70
71
72
73
        this.id = id;
        this.status = status;
        this.submitTime = submitTime;
        this.workingDirectory = workingDirectory;
        this.form = form;
        this.url = `file://localhost:${workingDirectory}`;
74
75
76
77
78
79
80
81
        this.gpuIndices = [];
    }
}

/**
 * Local training service config
 */
class LocalConfig {
82
    public maxTrialNumPerGpu?: number;
83
    public gpuIndices?: string;
84
85
    public useActiveGpu?: boolean;
    constructor(gpuIndices?: string, maxTrialNumPerGpu?: number, useActiveGpu?: boolean) {
86
87
88
        if (gpuIndices !== undefined) {
            this.gpuIndices = gpuIndices;
        }
89
90
91
92
93
94
        if (maxTrialNumPerGpu !== undefined) {
            this.maxTrialNumPerGpu = maxTrialNumPerGpu;
        }
        if (useActiveGpu !== undefined) {
            this.useActiveGpu = useActiveGpu;
        }
Deshui Yu's avatar
Deshui Yu committed
95
96
97
98
    }
}

/**
chicm-ms's avatar
chicm-ms committed
99
 * Local machine training service
Deshui Yu's avatar
Deshui Yu committed
100
101
 */
class LocalTrainingService implements TrainingService {
102
103
104
    private readonly eventEmitter: EventEmitter;
    private readonly jobMap: Map<string, LocalTrialJobDetail>;
    private readonly jobQueue: string[];
Deshui Yu's avatar
Deshui Yu committed
105
106
107
    private initialized: boolean;
    private stopping: boolean;
    private rootDir!: string;
chicm-ms's avatar
chicm-ms committed
108
    private readonly experimentId!: string;
109
    private gpuScheduler!: GPUScheduler;
110
    private readonly occupiedGpuIndexNumMap: Map<number, number>;
111
    private designatedGpuIndices!: Set<number>;
112
    private readonly log: Logger;
113
    private localTrialConfig?: TrialConfig;
114
    private localConfig?: LocalConfig;
115
    private isMultiPhase: boolean;
116
    private readonly jobStreamMap: Map<string, ts.Stream>;
117
118
    private maxTrialNumPerGpu: number;
    private useActiveGpu: boolean;
Deshui Yu's avatar
Deshui Yu committed
119
120
121
122
123
124
125
126

    constructor() {
        this.eventEmitter = new EventEmitter();
        this.jobMap = new Map<string, LocalTrialJobDetail>();
        this.jobQueue = [];
        this.initialized = false;
        this.stopping = false;
        this.log = getLogger();
suiguoxin's avatar
suiguoxin committed
127
        this.experimentId = getExperimentId();
128
        this.jobStreamMap = new Map<string, ts.Stream>();
chicm-ms's avatar
chicm-ms committed
129
        this.log.info('Construct local machine training service.');
130
131
132
133
        this.occupiedGpuIndexNumMap = new Map<number, number>();
        this.maxTrialNumPerGpu = 1;
        this.useActiveGpu = false;
        this.isMultiPhase = false;
Deshui Yu's avatar
Deshui Yu committed
134
135
136
    }

    public async run(): Promise<void> {
chicm-ms's avatar
chicm-ms committed
137
        this.log.info('Run local machine training service.');
138
139
140
        const longRunningTasks: Promise<void>[] = [this.runJobLoop()];
        if (this.gpuScheduler !== undefined) {
            longRunningTasks.push(this.gpuScheduler.run());
Deshui Yu's avatar
Deshui Yu committed
141
        }
142
        await Promise.all(longRunningTasks);
chicm-ms's avatar
chicm-ms committed
143
        this.log.info('Local machine training service exit.');
Deshui Yu's avatar
Deshui Yu committed
144
145
146
147
148
149
    }

    public async listTrialJobs(): Promise<TrialJobDetail[]> {
        const jobs: TrialJobDetail[] = [];
        for (const key of this.jobMap.keys()) {
            const trialJob: TrialJobDetail = await this.getTrialJob(key);
150
            jobs.push(trialJob);
Deshui Yu's avatar
Deshui Yu committed
151
152
153
154
155
156
157
158
159
160
161
        }

        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.status === 'RUNNING') {
162
            const alive: boolean = await isAlive(trialJob.pid);
Deshui Yu's avatar
Deshui Yu committed
163
            if (!alive) {
164
                trialJob.endTime = Date.now();
Deshui Yu's avatar
Deshui Yu committed
165
166
167
                this.setTrialJobStatus(trialJob, 'FAILED');
                try {
                    const state: string = await fs.promises.readFile(path.join(trialJob.workingDirectory, '.nni', 'state'), 'utf8');
168
169
                    const match: RegExpMatchArray | null = state.trim()
                        .match(/^(\d+)\s+(\d+)/);
Deshui Yu's avatar
Deshui Yu committed
170
171
172
173
174
                    if (match !== null) {
                        const { 1: code, 2: timestamp } = match;
                        if (parseInt(code, 10) === 0) {
                            this.setTrialJobStatus(trialJob, 'SUCCEEDED');
                        }
175
                        trialJob.endTime = parseInt(timestamp, 10);
Deshui Yu's avatar
Deshui Yu committed
176
177
178
179
                    }
                } catch (error) {
                    //ignore
                }
180
                this.log.debug(`trialJob status update: ${trialJobId}, ${trialJob.status}`);
Deshui Yu's avatar
Deshui Yu committed
181
182
183
184
185
186
            }
        }

        return trialJob;
    }

187
188
189
190
191
192
193
194
195
196
197
198
    public async getTrialLog(trialJobId: string, logType: LogType): Promise<string> {
        let logPath: string;
        if (logType === 'TRIAL_LOG') {
            logPath = path.join(this.rootDir, 'trials', trialJobId, 'trial.log');
        } else if (logType === 'TRIAL_ERROR') {
            logPath = path.join(this.rootDir, 'trials', trialJobId, 'stderr');
        } else {
            throw new Error('unexpected log type');
        }
        return fs.promises.readFile(logPath, 'utf8');
    }

Deshui Yu's avatar
Deshui Yu committed
199
200
201
202
203
204
205
206
    public addTrialJobMetricListener(listener: (metric: TrialJobMetric) => void): void {
        this.eventEmitter.on('metric', listener);
    }

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

207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
    public submitTrialJob(form: TrialJobApplicationForm): Promise<TrialJobDetail> {
        const trialJobId: string = uniqueString(5);
        const trialJobDetail: LocalTrialJobDetail = new LocalTrialJobDetail(
            trialJobId,
            'WAITING',
            Date.now(),
            path.join(this.rootDir, 'trials', trialJobId),
            form
        );
        this.jobQueue.push(trialJobId);
        this.jobMap.set(trialJobId, trialJobDetail);

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

        return Promise.resolve(trialJobDetail);
Deshui Yu's avatar
Deshui Yu committed
222
223
    }

224
225
226
227
228
    /**
     * Update trial job for multi-phase
     * @param trialJobId trial job id
     * @param form job application form
     */
229
    public async updateTrialJob(trialJobId: string, form: TrialJobApplicationForm): Promise<TrialJobDetail> {
chicm-ms's avatar
chicm-ms committed
230
231
232
233
        const trialJobDetail: undefined | TrialJobDetail = this.jobMap.get(trialJobId);
        if (trialJobDetail === undefined) {
            throw new Error(`updateTrialJob failed: ${trialJobId} not found`);
        }
234
        await this.writeParameterFile(trialJobDetail.workingDirectory, form.hyperParameters);
chicm-ms's avatar
chicm-ms committed
235
236

        return trialJobDetail;
237
238
239
240
241
242
    }

    /**
     * Is multiphase job supported in current training service
     */
    public get isMultiPhaseJobSupported(): boolean {
243
        return true;
244
245
    }

QuanluZhang's avatar
QuanluZhang committed
246
    public async cancelTrialJob(trialJobId: string, isEarlyStopped: boolean = false): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
247
248
249
250
        const trialJob: LocalTrialJobDetail | undefined = this.jobMap.get(trialJobId);
        if (trialJob === undefined) {
            throw new NNIError(NNIErrorNames.NOT_FOUND, 'Trial job not found');
        }
251
        if (trialJob.pid === undefined) {
252
            this.setTrialJobStatus(trialJob, 'USER_CANCELED');
253

SparkSnail's avatar
SparkSnail committed
254
            return Promise.resolve();
255
        }
256
        tkill(trialJob.pid, 'SIGKILL');
QuanluZhang's avatar
QuanluZhang committed
257
        this.setTrialJobStatus(trialJob, getJobCancelStatus(isEarlyStopped));
258

SparkSnail's avatar
SparkSnail committed
259
        return Promise.resolve();
Deshui Yu's avatar
Deshui Yu committed
260
261
262
263
264
    }

    public async setClusterMetadata(key: string, value: string): Promise<void> {
        if (!this.initialized) {
            this.rootDir = getExperimentRootDir();
265
            if (!fs.existsSync(this.rootDir)) {
266
267
                await cpp.exec(`powershell.exe mkdir ${this.rootDir}`);
            }
Deshui Yu's avatar
Deshui Yu committed
268
269
270
            this.initialized = true;
        }
        switch (key) {
271
            case TrialConfigMetadataKey.TRIAL_CONFIG:
272
                this.localTrialConfig = <TrialConfig>JSON.parse(value);
273
                // Parse trial config failed, throw Error
274
                if (this.localTrialConfig === undefined) {
275
276
                    throw new Error('trial config parsed failed');
                }
277
278
279
                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
280
281
                        this.gpuScheduler = new GPUScheduler();
                    }
282
283
284
285
286
287
288
289
290
291
292
293
                }
                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.');
                    }
                }
294
295
296
297
298
299
300
                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
301
                break;
302
303
304
            case TrialConfigMetadataKey.MULTI_PHASE:
                this.isMultiPhase = (value === 'true' || value === 'True');
                break;
Deshui Yu's avatar
Deshui Yu committed
305
306
307
308
309
310
            default:
        }
    }

    public getClusterMetadata(key: string): Promise<string> {
        switch (key) {
chicm-ms's avatar
chicm-ms committed
311
            case TrialConfigMetadataKey.TRIAL_CONFIG: {
312
                let getResult: Promise<string>;
313
                if (this.localTrialConfig === undefined) {
314
315
                    getResult = Promise.reject(new NNIError(NNIErrorNames.NOT_FOUND, `${key} is never set yet`));
                } else {
316
                    getResult = Promise.resolve(JSON.stringify(this.localTrialConfig));
317
                }
318

319
                return getResult;
chicm-ms's avatar
chicm-ms committed
320
            }
Deshui Yu's avatar
Deshui Yu committed
321
322
323
324
325
            default:
                return Promise.reject(new NNIError(NNIErrorNames.NOT_FOUND, 'Key not found'));
        }
    }

326
    public async cleanUp(): Promise<void> {
chicm-ms's avatar
chicm-ms committed
327
        this.log.info('Stopping local machine training service...');
Deshui Yu's avatar
Deshui Yu committed
328
        this.stopping = true;
329
        for (const stream of this.jobStreamMap.values()) {
330
331
            stream.end(0);
            stream.emit('end');
332
        }
333
334
335
336
        if (this.gpuScheduler !== undefined) {
            await this.gpuScheduler.stop();
        }

Deshui Yu's avatar
Deshui Yu committed
337
338
339
        return Promise.resolve();
    }

340
    private onTrialJobStatusChanged(trialJob: LocalTrialJobDetail, oldStatus: TrialJobStatus): void {
341
        //if job is not running, destory job stream
342
343
344
        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);
345
                if (stream === undefined) {
346
347
                    throw new Error(`Could not find stream in trial ${trialJob.id}`);
                }
348
                //Refer https://github.com/Juul/tail-stream/issues/20
Yan Ni's avatar
Yan Ni committed
349
350
351
352
353
                setTimeout(() => {
                    stream.end(0);
                    stream.emit('end');
                    this.jobStreamMap.delete(trialJob.id);
                }, 5000);
354
355
            }
        }
356
357
358
        if (trialJob.gpuIndices !== undefined && trialJob.gpuIndices.length > 0 && this.gpuScheduler !== undefined) {
            if (oldStatus === 'RUNNING' && trialJob.status !== 'RUNNING') {
                for (const index of trialJob.gpuIndices) {
359
360
                    const num: number | undefined = this.occupiedGpuIndexNumMap.get(index);
                    if (num === undefined) {
361
                        throw new Error(`gpu resource schedule error`);
362
                    } else if (num === 1) {
363
364
                        this.occupiedGpuIndexNumMap.delete(index);
                    } else {
365
                        this.occupiedGpuIndexNumMap.set(index, num - 1);
366
                    }
367
368
369
                }
            }
        }
Deshui Yu's avatar
Deshui Yu committed
370
371
    }

372
373
    private getEnvironmentVariables(
        trialJobDetail: TrialJobDetail,
SparkSnail's avatar
SparkSnail committed
374
375
        resource: { gpuIndices: number[] },
        gpuNum: number | undefined): { key: string; value: string }[] {
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
            if (this.localTrialConfig === undefined) {
                throw new Error('localTrialConfig is not initialized!');
            }
            const envVariables: { key: string; value: string }[] = [
                { key: 'NNI_PLATFORM', value: 'local' },
                { key: 'NNI_EXP_ID', value: this.experimentId },
                { key: 'NNI_SYS_DIR', value: trialJobDetail.workingDirectory },
                { key: 'NNI_TRIAL_JOB_ID', value: trialJobDetail.id },
                { key: 'NNI_OUTPUT_DIR', value: trialJobDetail.workingDirectory },
                { key: 'NNI_TRIAL_SEQ_ID', value: trialJobDetail.form.sequenceId.toString() },
                { key: 'MULTI_PHASE', value: this.isMultiPhase.toString() },
                { key: 'NNI_CODE_DIR', value: this.localTrialConfig.codeDir}
            ];
            if (gpuNum !== undefined) {
                envVariables.push({
                    key: 'CUDA_VISIBLE_DEVICES',
                    value: this.gpuScheduler === undefined ? '-1' : resource.gpuIndices.join(',')
                });
            }
395
396
397
398
399
400
401
402
403

        return envVariables;
    }

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

    private tryGetAvailableResource(): [boolean, { gpuIndices: number[]}] {
404
405
        if (this.localTrialConfig === undefined) {
            throw new Error('localTrialConfig is not initialized!');
406
407
408
409
410
411
412
        }

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

413
        let selectedGPUIndices: number[] = [];
414
415
416
417
        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) {
418
419
420
                selectedGPUIndices.push(index);
            }
        }
421
422
423
424
425
426

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

427
        if (selectedGPUIndices.length < this.localTrialConfig.gpuNum) {
428
429
430
            return [false, resource];
        }

431
        selectedGPUIndices.splice(this.localTrialConfig.gpuNum);
432
433
434
        Object.assign(resource, { gpuIndices: selectedGPUIndices });

        return [true, resource];
Deshui Yu's avatar
Deshui Yu committed
435
436
    }

437
438
439
440
441
442
443
444
445
    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
446
447
    }

448
449
450
    private occupyResource(resource: {gpuIndices: number[]}): void {
        if (this.gpuScheduler !== undefined) {
            for (const index of resource.gpuIndices) {
451
452
453
                const num: number | undefined = this.occupiedGpuIndexNumMap.get(index);
                if (num === undefined) {
                    this.occupiedGpuIndexNumMap.set(index, 1);
454
                } else {
455
                    this.occupiedGpuIndexNumMap.set(index, num + 1);
456
                }
457
458
            }
        }
Deshui Yu's avatar
Deshui Yu committed
459
460
    }

461
462
463
464
    private async runJobLoop(): Promise<void> {
        while (!this.stopping) {
            while (!this.stopping && this.jobQueue.length !== 0) {
                const trialJobId: string = this.jobQueue[0];
465
466
                const trialJobDetail: LocalTrialJobDetail | undefined = this.jobMap.get(trialJobId);
                if (trialJobDetail !== undefined && trialJobDetail.status === 'WAITING') {
467
468
469
470
                    const [success, resource] = this.tryGetAvailableResource();
                    if (!success) {
                        break;
                    }
471

472
473
474
475
476
477
478
                    this.occupyResource(resource);
                    await this.runTrialJob(trialJobId, resource);
                }
                this.jobQueue.shift();
            }
            await delay(5000);
        }
Deshui Yu's avatar
Deshui Yu committed
479
480
481
482
483
484
485
486
487
488
    }

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

489
    private getScript(localTrialConfig: TrialConfig, workingDirectory: string): string[] {
490
491
        const script: string[] = [];
        if (process.platform === 'win32') {
SparkSnail's avatar
SparkSnail committed
492
            script.push(`cd $env:NNI_CODE_DIR`);
493
            script.push(
494
                `cmd.exe /c ${localTrialConfig.command} 2>&1 | Out-File "${path.join(workingDirectory, 'stderr')}" -encoding utf8`,
495
                `$NOW_DATE = [int64](([datetime]::UtcNow)-(get-date "1/1/1970")).TotalSeconds`,
496
                `$NOW_DATE = "$NOW_DATE" + (Get-Date -Format fff).ToString()`,
SparkSnail's avatar
SparkSnail committed
497
                `Write $LASTEXITCODE " " $NOW_DATE  | Out-File "${path.join(workingDirectory, '.nni', 'state')}" -NoNewline -encoding utf8`);
498
        } else {
SparkSnail's avatar
SparkSnail committed
499
            script.push(`cd $NNI_CODE_DIR`);
SparkSnail's avatar
SparkSnail committed
500
            script.push(`eval ${localTrialConfig.command} 2>"${path.join(workingDirectory, 'stderr')}"`);
501
502
503
            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
SparkSnail's avatar
SparkSnail committed
504
                script.push(`echo $? \`date +%s999\` >'${path.join(workingDirectory, '.nni', 'state')}'`);
505
            } else {
SparkSnail's avatar
SparkSnail committed
506
                script.push(`echo $? \`date +%s%3N\` >'${path.join(workingDirectory, '.nni', 'state')}'`);
507
            }
508
        }
509

510
511
512
        return script;
    }

513
    private async runTrialJob(trialJobId: string, resource: {gpuIndices: number[]}): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
514
        const trialJobDetail: LocalTrialJobDetail = <LocalTrialJobDetail>this.jobMap.get(trialJobId);
515
        if (this.localTrialConfig === undefined) {
SparkSnail's avatar
SparkSnail committed
516
517
            throw new Error(`localTrialConfig not initialized!`);
        }
518
        const variables: { key: string; value: string }[] = this.getEnvironmentVariables(trialJobDetail, resource, this.localTrialConfig.gpuNum);
Deshui Yu's avatar
Deshui Yu committed
519

520
        if (this.localTrialConfig === undefined) {
521
522
            throw new Error('trial config is not initialized');
        }
523
524
525
        const runScriptContent: string[] = [];
        if (process.platform !== 'win32') {
            runScriptContent.push('#!/bin/bash');
526
527
        } else {
            runScriptContent.push(`$env:PATH="${process.env.path}"`)
528
        }
Deshui Yu's avatar
Deshui Yu committed
529
        for (const variable of variables) {
530
            runScriptContent.push(setEnvironmentVariable(variable));
Deshui Yu's avatar
Deshui Yu committed
531
        }
532
        const scripts: string[] = this.getScript(this.localTrialConfig, trialJobDetail.workingDirectory);
533
534
        scripts.forEach((script: string) => {
            runScriptContent.push(script);
535
536
537
538
539
        });
        await execMkdir(trialJobDetail.workingDirectory);
        await execMkdir(path.join(trialJobDetail.workingDirectory, '.nni'));
        await execNewFile(path.join(trialJobDetail.workingDirectory, '.nni', 'metrics'));
        const scriptName: string = getScriptName('run');
540
541
        await fs.promises.writeFile(path.join(trialJobDetail.workingDirectory, scriptName),
                                    runScriptContent.join(getNewLine()), { encoding: 'utf8', mode: 0o777 });
542
        await this.writeParameterFile(trialJobDetail.workingDirectory, trialJobDetail.form.hyperParameters);
543
        const trialJobProcess: cp.ChildProcess = runScript(path.join(trialJobDetail.workingDirectory, scriptName));
Deshui Yu's avatar
Deshui Yu committed
544
        this.setTrialJobStatus(trialJobDetail, 'RUNNING');
545
546
        trialJobDetail.startTime = Date.now(); // eslint-disable-line require-atomic-updates
        trialJobDetail.pid = trialJobProcess.pid; // eslint-disable-line require-atomic-updates
Deshui Yu's avatar
Deshui Yu committed
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
        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;
            }
        });
566
        this.jobStreamMap.set(trialJobDetail.id, stream);
Deshui Yu's avatar
Deshui Yu committed
567
568
    }

chicm-ms's avatar
chicm-ms committed
569
    private async writeParameterFile(directory: string, hyperParameters: HyperParameters): Promise<void> {
570
        const filepath: string = path.join(directory, generateParamFileName(hyperParameters));
chicm-ms's avatar
chicm-ms committed
571
572
        await fs.promises.writeFile(filepath, hyperParameters.value, { encoding: 'utf8' });
    }
Deshui Yu's avatar
Deshui Yu committed
573
574
575
}

export { LocalTrainingService };