localTrainingService.ts 22.8 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,
Deshui Yu's avatar
Deshui Yu committed
17
18
    TrialJobDetail, TrialJobMetric, TrialJobStatus
} 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
187
188
189
190
191
192
193
194
            }
        }

        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);
    }

195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
    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
210
211
    }

212
213
214
215
216
    /**
     * Update trial job for multi-phase
     * @param trialJobId trial job id
     * @param form job application form
     */
217
    public async updateTrialJob(trialJobId: string, form: TrialJobApplicationForm): Promise<TrialJobDetail> {
chicm-ms's avatar
chicm-ms committed
218
219
220
221
        const trialJobDetail: undefined | TrialJobDetail = this.jobMap.get(trialJobId);
        if (trialJobDetail === undefined) {
            throw new Error(`updateTrialJob failed: ${trialJobId} not found`);
        }
222
        await this.writeParameterFile(trialJobDetail.workingDirectory, form.hyperParameters);
chicm-ms's avatar
chicm-ms committed
223
224

        return trialJobDetail;
225
226
227
228
229
230
    }

    /**
     * Is multiphase job supported in current training service
     */
    public get isMultiPhaseJobSupported(): boolean {
231
        return true;
232
233
    }

QuanluZhang's avatar
QuanluZhang committed
234
    public async cancelTrialJob(trialJobId: string, isEarlyStopped: boolean = false): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
235
236
237
238
        const trialJob: LocalTrialJobDetail | undefined = this.jobMap.get(trialJobId);
        if (trialJob === undefined) {
            throw new NNIError(NNIErrorNames.NOT_FOUND, 'Trial job not found');
        }
239
        if (trialJob.pid === undefined) {
240
            this.setTrialJobStatus(trialJob, 'USER_CANCELED');
241

SparkSnail's avatar
SparkSnail committed
242
            return Promise.resolve();
243
        }
244
        tkill(trialJob.pid, 'SIGKILL');
QuanluZhang's avatar
QuanluZhang committed
245
        this.setTrialJobStatus(trialJob, getJobCancelStatus(isEarlyStopped));
246

SparkSnail's avatar
SparkSnail committed
247
        return Promise.resolve();
Deshui Yu's avatar
Deshui Yu committed
248
249
250
251
252
    }

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

    public getClusterMetadata(key: string): Promise<string> {
        switch (key) {
chicm-ms's avatar
chicm-ms committed
299
            case TrialConfigMetadataKey.TRIAL_CONFIG: {
300
                let getResult: Promise<string>;
301
                if (this.localTrialConfig === undefined) {
302
303
                    getResult = Promise.reject(new NNIError(NNIErrorNames.NOT_FOUND, `${key} is never set yet`));
                } else {
304
                    getResult = Promise.resolve(JSON.stringify(this.localTrialConfig));
305
                }
306

307
                return getResult;
chicm-ms's avatar
chicm-ms committed
308
            }
Deshui Yu's avatar
Deshui Yu committed
309
310
311
312
313
            default:
                return Promise.reject(new NNIError(NNIErrorNames.NOT_FOUND, 'Key not found'));
        }
    }

314
    public async cleanUp(): Promise<void> {
chicm-ms's avatar
chicm-ms committed
315
        this.log.info('Stopping local machine training service...');
Deshui Yu's avatar
Deshui Yu committed
316
        this.stopping = true;
317
        for (const stream of this.jobStreamMap.values()) {
318
319
            stream.end(0);
            stream.emit('end');
320
        }
321
322
323
324
        if (this.gpuScheduler !== undefined) {
            await this.gpuScheduler.stop();
        }

Deshui Yu's avatar
Deshui Yu committed
325
326
327
        return Promise.resolve();
    }

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

360
361
    private getEnvironmentVariables(
        trialJobDetail: TrialJobDetail,
SparkSnail's avatar
SparkSnail committed
362
363
        resource: { gpuIndices: number[] },
        gpuNum: number | undefined): { key: string; value: string }[] {
364
        const envVariables: { key: string; value: string }[] = [
Deshui Yu's avatar
Deshui Yu committed
365
            { key: 'NNI_PLATFORM', value: 'local' },
suiguoxin's avatar
suiguoxin committed
366
            { key: 'NNI_EXP_ID', value: this.experimentId },
Deshui Yu's avatar
Deshui Yu committed
367
368
            { key: 'NNI_SYS_DIR', value: trialJobDetail.workingDirectory },
            { key: 'NNI_TRIAL_JOB_ID', value: trialJobDetail.id },
369
            { key: 'NNI_OUTPUT_DIR', value: trialJobDetail.workingDirectory },
370
            { key: 'NNI_TRIAL_SEQ_ID', value: trialJobDetail.form.sequenceId.toString() },
371
            { key: 'MULTI_PHASE', value: this.isMultiPhase.toString() }
Deshui Yu's avatar
Deshui Yu committed
372
        ];
SparkSnail's avatar
SparkSnail committed
373
374
375
376
377
378
        if (gpuNum !== undefined) {
            envVariables.push({
                key: 'CUDA_VISIBLE_DEVICES',
                value: this.gpuScheduler === undefined ? '-1' : resource.gpuIndices.join(',')
            });
        }
379
380
381
382
383
384
385
386
387

        return envVariables;
    }

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

    private tryGetAvailableResource(): [boolean, { gpuIndices: number[]}] {
388
389
        if (this.localTrialConfig === undefined) {
            throw new Error('localTrialConfig is not initialized!');
390
391
392
393
394
395
396
        }

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

397
        let selectedGPUIndices: number[] = [];
398
399
400
401
        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) {
402
403
404
                selectedGPUIndices.push(index);
            }
        }
405
406
407
408
409
410

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

411
        if (selectedGPUIndices.length < this.localTrialConfig.gpuNum) {
412
413
414
            return [false, resource];
        }

415
        selectedGPUIndices.splice(this.localTrialConfig.gpuNum);
416
417
418
        Object.assign(resource, { gpuIndices: selectedGPUIndices });

        return [true, resource];
Deshui Yu's avatar
Deshui Yu committed
419
420
    }

421
422
423
424
425
426
427
428
429
    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
430
431
    }

432
433
434
    private occupyResource(resource: {gpuIndices: number[]}): void {
        if (this.gpuScheduler !== undefined) {
            for (const index of resource.gpuIndices) {
435
436
437
                const num: number | undefined = this.occupiedGpuIndexNumMap.get(index);
                if (num === undefined) {
                    this.occupiedGpuIndexNumMap.set(index, 1);
438
                } else {
439
                    this.occupiedGpuIndexNumMap.set(index, num + 1);
440
                }
441
442
            }
        }
Deshui Yu's avatar
Deshui Yu committed
443
444
    }

445
446
447
448
449
450
451
452
453
454
    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;
                    }
455

456
457
458
459
460
461
462
                    this.occupyResource(resource);
                    await this.runTrialJob(trialJobId, resource);
                }
                this.jobQueue.shift();
            }
            await delay(5000);
        }
Deshui Yu's avatar
Deshui Yu committed
463
464
465
466
467
468
469
470
471
472
    }

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

473
    private getScript(localTrialConfig: TrialConfig, workingDirectory: string): string[] {
474
475
        const script: string[] = [];
        if (process.platform === 'win32') {
476
            script.push(
SparkSnail's avatar
SparkSnail committed
477
                `cmd.exe /c ${localTrialConfig.command} 2>"${path.join(workingDirectory, 'stderr')}"`,
478
                `$NOW_DATE = [int64](([datetime]::UtcNow)-(get-date "1/1/1970")).TotalSeconds`,
479
                `$NOW_DATE = "$NOW_DATE" + (Get-Date -Format fff).ToString()`,
SparkSnail's avatar
SparkSnail committed
480
                `Write $LASTEXITCODE " " $NOW_DATE  | Out-File "${path.join(workingDirectory, '.nni', 'state')}" -NoNewline -encoding utf8`);
481
        } else {
SparkSnail's avatar
SparkSnail committed
482
            script.push(`eval ${localTrialConfig.command} 2>"${path.join(workingDirectory, 'stderr')}"`);
483
484
485
            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
486
                script.push(`echo $? \`date +%s999\` >'${path.join(workingDirectory, '.nni', 'state')}'`);
487
            } else {
SparkSnail's avatar
SparkSnail committed
488
                script.push(`echo $? \`date +%s%3N\` >'${path.join(workingDirectory, '.nni', 'state')}'`);
489
            }
490
        }
491

492
493
494
        return script;
    }

495
    private async runTrialJob(trialJobId: string, resource: {gpuIndices: number[]}): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
496
        const trialJobDetail: LocalTrialJobDetail = <LocalTrialJobDetail>this.jobMap.get(trialJobId);
497
        if (this.localTrialConfig === undefined) {
SparkSnail's avatar
SparkSnail committed
498
499
            throw new Error(`localTrialConfig not initialized!`);
        }
500
        const variables: { key: string; value: string }[] = this.getEnvironmentVariables(trialJobDetail, resource, this.localTrialConfig.gpuNum);
Deshui Yu's avatar
Deshui Yu committed
501

502
        if (this.localTrialConfig === undefined) {
503
504
            throw new Error('trial config is not initialized');
        }
505
506
507
        const runScriptContent: string[] = [];
        if (process.platform !== 'win32') {
            runScriptContent.push('#!/bin/bash');
508
        }
SparkSnail's avatar
SparkSnail committed
509
        runScriptContent.push(`cd '${this.localTrialConfig.codeDir}'`);
Deshui Yu's avatar
Deshui Yu committed
510
        for (const variable of variables) {
511
            runScriptContent.push(setEnvironmentVariable(variable));
Deshui Yu's avatar
Deshui Yu committed
512
        }
513
        const scripts: string[] = this.getScript(this.localTrialConfig, trialJobDetail.workingDirectory);
514
515
        scripts.forEach((script: string) => {
            runScriptContent.push(script);
516
517
518
519
520
        });
        await execMkdir(trialJobDetail.workingDirectory);
        await execMkdir(path.join(trialJobDetail.workingDirectory, '.nni'));
        await execNewFile(path.join(trialJobDetail.workingDirectory, '.nni', 'metrics'));
        const scriptName: string = getScriptName('run');
521
522
        await fs.promises.writeFile(path.join(trialJobDetail.workingDirectory, scriptName),
                                    runScriptContent.join(getNewLine()), { encoding: 'utf8', mode: 0o777 });
523
        await this.writeParameterFile(trialJobDetail.workingDirectory, trialJobDetail.form.hyperParameters);
524
        const trialJobProcess: cp.ChildProcess = runScript(path.join(trialJobDetail.workingDirectory, scriptName));
Deshui Yu's avatar
Deshui Yu committed
525
        this.setTrialJobStatus(trialJobDetail, 'RUNNING');
526
527
        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
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
        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;
            }
        });
547
        this.jobStreamMap.set(trialJobDetail.id, stream);
Deshui Yu's avatar
Deshui Yu committed
548
549
    }

chicm-ms's avatar
chicm-ms committed
550
    private async writeParameterFile(directory: string, hyperParameters: HyperParameters): Promise<void> {
551
        const filepath: string = path.join(directory, generateParamFileName(hyperParameters));
chicm-ms's avatar
chicm-ms committed
552
553
        await fs.promises.writeFile(filepath, hyperParameters.value, { encoding: 'utf8' });
    }
Deshui Yu's avatar
Deshui Yu committed
554
555
556
}

export { LocalTrainingService };