localTrainingService.ts 25.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';
29
import { 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;
129
    private gpuScheduler!: GPUScheduler;
130
    private readonly occupiedGpuIndexNumMap: Map<number, number>;
131
    private designatedGpuIndices!: Set<number>;
132
    private readonly log: Logger;
133
134
    private localTrailConfig?: TrialConfig;
    private localConfig?: LocalConfig;
135
    private isMultiPhase: boolean;
136
    private readonly jobStreamMap: Map<string, ts.Stream>;
137
138
    private maxTrialNumPerGpu: number;
    private useActiveGpu: boolean;
Deshui Yu's avatar
Deshui Yu committed
139
140
141
142
143
144
145
146

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

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

    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') {
187
            const alive: boolean = await isAlive(trialJob.pid);
Deshui Yu's avatar
Deshui Yu committed
188
            if (!alive) {
189
                trialJob.endTime = Date.now();
Deshui Yu's avatar
Deshui Yu committed
190
191
192
                this.setTrialJobStatus(trialJob, 'FAILED');
                try {
                    const state: string = await fs.promises.readFile(path.join(trialJob.workingDirectory, '.nni', 'state'), 'utf8');
193
194
                    const match: RegExpMatchArray | null = state.trim()
                        .match(/^(\d+)\s+(\d+)/);
Deshui Yu's avatar
Deshui Yu committed
195
196
197
198
199
                    if (match !== null) {
                        const { 1: code, 2: timestamp } = match;
                        if (parseInt(code, 10) === 0) {
                            this.setTrialJobStatus(trialJob, 'SUCCEEDED');
                        }
200
                        trialJob.endTime = parseInt(timestamp, 10);
Deshui Yu's avatar
Deshui Yu committed
201
202
203
204
                    }
                } catch (error) {
                    //ignore
                }
chicm-ms's avatar
chicm-ms committed
205
                this.log.debug(`trailJob status update: ${trialJobId}, ${trialJob.status}`);
Deshui Yu's avatar
Deshui Yu committed
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
            }
        }

        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',
228
                Date.now(),
Deshui Yu's avatar
Deshui Yu committed
229
                path.join(this.rootDir, 'trials', trialJobId),
230
231
232
                form,
                this.generateSequenceId()
            );
Deshui Yu's avatar
Deshui Yu committed
233
234
235
236
237
238
239
240
241
242
243
            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)}`));
        }
    }

244
245
246
247
248
    /**
     * Update trial job for multi-phase
     * @param trialJobId trial job id
     * @param form job application form
     */
chicm-ms's avatar
chicm-ms committed
249
250
251
252
253
254
255
256
257
258
259
260
    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;
261
262
263
264
265
266
    }

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

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

SparkSnail's avatar
SparkSnail committed
278
            return Promise.resolve();
279
        }
Deshui Yu's avatar
Deshui Yu committed
280
        if (trialJob.form.jobType === 'TRIAL') {
281
            tkill(trialJob.pid, 'SIGKILL');
Deshui Yu's avatar
Deshui Yu committed
282
283
284
285
286
        } 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
287
        this.setTrialJobStatus(trialJob, getJobCancelStatus(isEarlyStopped));
288

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

    public async setClusterMetadata(key: string, value: string): Promise<void> {
        if (!this.initialized) {
            this.rootDir = getExperimentRootDir();
295
296
            // tslint:disable-next-line:non-literal-fs-path
            if (!fs.existsSync(this.rootDir)) {
297
298
                await cpp.exec(`powershell.exe mkdir ${this.rootDir}`);
            }
Deshui Yu's avatar
Deshui Yu committed
299
300
301
            this.initialized = true;
        }
        switch (key) {
302
303
304
            case TrialConfigMetadataKey.TRIAL_CONFIG:
                this.localTrailConfig = <TrialConfig>JSON.parse(value);
                // Parse trial config failed, throw Error
305
                if (this.localTrailConfig === undefined) {
306
307
                    throw new Error('trial config parsed failed');
                }
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
                this.log.info(`required GPU number is ${this.localTrailConfig.gpuNum}`);
                if (this.gpuScheduler === undefined && this.localTrailConfig.gpuNum > 0) {
                    this.gpuScheduler = new GPUScheduler();
                }
                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.');
                    }
                }
323
324
325
326
327
328
329
                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
330
                break;
331
332
333
            case TrialConfigMetadataKey.MULTI_PHASE:
                this.isMultiPhase = (value === 'true' || value === 'True');
                break;
Deshui Yu's avatar
Deshui Yu committed
334
335
336
337
338
339
            default:
        }
    }

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

348
                return getResult;
Deshui Yu's avatar
Deshui Yu committed
349
350
351
352
353
            default:
                return Promise.reject(new NNIError(NNIErrorNames.NOT_FOUND, 'Key not found'));
        }
    }

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

Deshui Yu's avatar
Deshui Yu committed
364
365
366
        return Promise.resolve();
    }

367
    private onTrialJobStatusChanged(trialJob: LocalTrialJobDetail, oldStatus: TrialJobStatus): void {
368
        //if job is not running, destory job stream
369
370
371
        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);
372
                if (stream === undefined) {
373
374
375
376
377
378
                    throw new Error(`Could not find stream in trial ${trialJob.id}`);
                }
                stream.destroy();
                this.jobStreamMap.delete(trialJob.id);
            }
        }
379
380
381
        if (trialJob.gpuIndices !== undefined && trialJob.gpuIndices.length > 0 && this.gpuScheduler !== undefined) {
            if (oldStatus === 'RUNNING' && trialJob.status !== 'RUNNING') {
                for (const index of trialJob.gpuIndices) {
382
383
                    const num: number | undefined = this.occupiedGpuIndexNumMap.get(index);
                    if (num === undefined) {
384
                        throw new Error(`gpu resource schedule error`);
385
                    } else if (num === 1) {
386
387
                        this.occupiedGpuIndexNumMap.delete(index);
                    } else {
388
                        this.occupiedGpuIndexNumMap.set(index, num - 1);
389
                    }
390
391
392
                }
            }
        }
Deshui Yu's avatar
Deshui Yu committed
393
394
    }

395
396
    private getEnvironmentVariables(
        trialJobDetail: TrialJobDetail,
chicm-ms's avatar
chicm-ms committed
397
        resource: { gpuIndices: number[] }): { key: string; value: string }[] {
398
        const envVariables: { key: string; value: string }[] = [
Deshui Yu's avatar
Deshui Yu committed
399
400
401
            { key: 'NNI_PLATFORM', value: 'local' },
            { key: 'NNI_SYS_DIR', value: trialJobDetail.workingDirectory },
            { key: 'NNI_TRIAL_JOB_ID', value: trialJobDetail.id },
402
            { key: 'NNI_OUTPUT_DIR', value: trialJobDetail.workingDirectory },
403
            { key: 'NNI_TRIAL_SEQ_ID', value: trialJobDetail.sequenceId.toString() },
404
            { key: 'MULTI_PHASE', value: this.isMultiPhase.toString() }
Deshui Yu's avatar
Deshui Yu committed
405
        ];
406

chicm-ms's avatar
chicm-ms committed
407
408
        envVariables.push({
            key: 'CUDA_VISIBLE_DEVICES',
409
            value: this.gpuScheduler === undefined ? '-1' : resource.gpuIndices.join(',')
chicm-ms's avatar
chicm-ms committed
410
        });
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428

        return envVariables;
    }

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

    private tryGetAvailableResource(): [boolean, { gpuIndices: number[]}] {
        if (this.localTrailConfig === undefined) {
            throw new Error('localTrailConfig is not initialized!');
        }

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

429
        let selectedGPUIndices: number[] = [];
430
431
432
433
        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) {
434
435
436
                selectedGPUIndices.push(index);
            }
        }
437
438
439
440
441
442
443
444
445
446
447
448
449
450

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

        if (selectedGPUIndices.length < this.localTrailConfig.gpuNum) {
            return [false, resource];
        }

        selectedGPUIndices.splice(this.localTrailConfig.gpuNum);
        Object.assign(resource, { gpuIndices: selectedGPUIndices });

        return [true, resource];
Deshui Yu's avatar
Deshui Yu committed
451
452
    }

453
454
455
456
457
458
459
460
461
    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
462
463
    }

464
465
466
    private occupyResource(resource: {gpuIndices: number[]}): void {
        if (this.gpuScheduler !== undefined) {
            for (const index of resource.gpuIndices) {
467
468
469
                const num: number | undefined = this.occupiedGpuIndexNumMap.get(index);
                if (num === undefined) {
                    this.occupiedGpuIndexNumMap.set(index, 1);
470
                } else {
471
                    this.occupiedGpuIndexNumMap.set(index, num + 1);
472
                }
473
474
            }
        }
Deshui Yu's avatar
Deshui Yu committed
475
476
    }

477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
    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;
                    }
                    this.occupyResource(resource);
                    await this.runTrialJob(trialJobId, resource);
                }
                this.jobQueue.shift();
            }
            await delay(5000);
        }
Deshui Yu's avatar
Deshui Yu committed
494
495
496
497
498
499
500
501
502
503
    }

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

504
505
506
    private getScript(localTrailConfig: TrialConfig, workingDirectory: string): string[] {
        const script: string[] = [];
        if (process.platform === 'win32') {
507
508
509
            script.push(
                `cmd /c ${localTrailConfig.command} 2>${path.join(workingDirectory, 'stderr')}`,
                `$NOW_DATE = [int64](([datetime]::UtcNow)-(get-date "1/1/1970")).TotalSeconds`,
510
                `$NOW_DATE = "$NOW_DATE" + (Get-Date -Format fff).ToString()`,
511
                `Write $LASTEXITCODE " " $NOW_DATE  | Out-File ${path.join(workingDirectory, '.nni', 'state')} -NoNewline -encoding utf8`);
512
        } else {
513
514
            script.push(
                `eval ${localTrailConfig.command} 2>${path.join(workingDirectory, 'stderr')}`,
515
                `echo $? \`date +%s%3N\` >${path.join(workingDirectory, '.nni', 'state')}`);
516
        }
517

518
519
520
        return script;
    }

521
    private async runTrialJob(trialJobId: string, resource: {gpuIndices: number[]}): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
522
523
524
        const trialJobDetail: LocalTrialJobDetail = <LocalTrialJobDetail>this.jobMap.get(trialJobId);
        const variables: { key: string; value: string }[] = this.getEnvironmentVariables(trialJobDetail, resource);

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

571
        this.jobStreamMap.set(trialJobDetail.id, stream);
Deshui Yu's avatar
Deshui Yu committed
572
573
574
575
576
577
578
579
580
581
582
583
    }

    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',
584
            submitTime: Date.now(),
Deshui Yu's avatar
Deshui Yu committed
585
586
            workingDirectory: workDir,
            form: form,
587
            sequenceId: this.generateSequenceId(),
Deshui Yu's avatar
Deshui Yu committed
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
            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
615
616

    private async writeParameterFile(directory: string, hyperParameters: HyperParameters): Promise<void> {
617
        const filepath: string = path.join(directory, generateParamFileName(hyperParameters));
chicm-ms's avatar
chicm-ms committed
618
619
        await fs.promises.writeFile(filepath, hyperParameters.value, { encoding: 'utf8' });
    }
620
621

    private generateSequenceId(): number {
622
623
624
625
        if (this.trialSequenceId === -1) {
            this.trialSequenceId = getInitTrialSequenceId();
        }

626
627
        return this.trialSequenceId++;
    }
Deshui Yu's avatar
Deshui Yu committed
628
629
630
}

export { LocalTrainingService };