remoteMachineTrainingService.ts 29.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

'use strict';

6
import * as assert from 'assert';
Deshui Yu's avatar
Deshui Yu committed
7
8
9
10
11
12
import { EventEmitter } from 'events';
import * as fs from 'fs';
import * as path from 'path';
import { Deferred } from 'ts-deferred';
import { String } from 'typescript-string-operations';
import * as component from '../../common/component';
SparkSnail's avatar
SparkSnail committed
13
import { NNIError, NNIErrorNames } from '../../common/errors';
14
import { getExperimentId } from '../../common/experimentStartupInfo';
Deshui Yu's avatar
Deshui Yu committed
15
16
17
import { getLogger, Logger } from '../../common/log';
import { ObservableTimer } from '../../common/observableTimer';
import {
18
    HyperParameters, NNIManagerIpConfig, TrainingService, TrialJobApplicationForm,
19
    TrialJobDetail, TrialJobMetric
Deshui Yu's avatar
Deshui Yu committed
20
} from '../../common/trainingService';
21
22
23
24
25
import {
    delay, generateParamFileName, getExperimentRootDir, getIPV4Address, getJobCancelStatus, getRemoteTmpDir,
    getVersion, uniqueString, unixPathJoin
} from '../../common/utils';
import { CONTAINER_INSTALL_NNI_SHELL_FORMAT } from '../common/containerJobData';
26
import { GPUSummary } from '../common/gpuData';
27
28
import { TrialConfig } from '../common/trialConfig';
import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
chicm-ms's avatar
chicm-ms committed
29
import { execCopydir, execMkdir, validateCodeDir, getGpuMetricsCollectorBashScriptContent } from '../common/util';
Deshui Yu's avatar
Deshui Yu committed
30
31
import { GPUScheduler } from './gpuScheduler';
import {
32
    REMOTEMACHINE_TRIAL_COMMAND_FORMAT, RemoteMachineMeta,
33
    RemoteMachineScheduleInfo, RemoteMachineScheduleResult, RemoteMachineTrialJobDetail,
34
    ScheduleResultType, ExecutorManager
Deshui Yu's avatar
Deshui Yu committed
35
} from './remoteMachineData';
SparkSnail's avatar
SparkSnail committed
36
import { RemoteMachineJobRestServer } from './remoteMachineJobRestServer';
37
import { ShellExecutor } from 'training_service/remote_machine/shellExecutor';
Deshui Yu's avatar
Deshui Yu committed
38
39
40
41

/**
 * Training Service implementation for Remote Machine (Linux)
 */
SparkSnail's avatar
SparkSnail committed
42
@component.Singleton
Deshui Yu's avatar
Deshui Yu committed
43
class RemoteMachineTrainingService implements TrainingService {
44
45
    private readonly machineExecutorManagerMap: Map<RemoteMachineMeta, ExecutorManager>; //machine excutor map
    private readonly trialExecutorMap: Map<string, ShellExecutor>; //trial excutor map
46
    private readonly trialJobsMap: Map<string, RemoteMachineTrialJobDetail>;
47
    private readonly MAX_TRIAL_NUMBER_PER_EXECUTOR: number = 5; // every excutor has a max trial concurrency number
48
49
    private readonly expRootDir: string;
    private readonly remoteExpRootDir: string;
50
    private trialConfig: TrialConfig | undefined;
51
    private gpuScheduler?: GPUScheduler;
52
53
    private readonly jobQueue: string[];
    private readonly timer: ObservableTimer;
Deshui Yu's avatar
Deshui Yu committed
54
    private stopping: boolean = false;
55
56
    private readonly metricsEmitter: EventEmitter;
    private readonly log: Logger;
57
    private isMultiPhase: boolean = false;
58
    private trialSequenceId: number;
SparkSnail's avatar
SparkSnail committed
59
    private remoteRestServerPort?: number;
60
    private readonly remoteOS: string;
SparkSnail's avatar
SparkSnail committed
61
    private nniManagerIpConfig?: NNIManagerIpConfig;
62
    private versionCheck: boolean = true;
SparkSnail's avatar
SparkSnail committed
63
    private logCollection: string;
Deshui Yu's avatar
Deshui Yu committed
64
65

    constructor(@component.Inject timer: ObservableTimer) {
66
        this.remoteOS = 'linux';
Deshui Yu's avatar
Deshui Yu committed
67
68
        this.metricsEmitter = new EventEmitter();
        this.trialJobsMap = new Map<string, RemoteMachineTrialJobDetail>();
69
70
        this.trialExecutorMap = new Map<string, ShellExecutor>();
        this.machineExecutorManagerMap = new Map<RemoteMachineMeta, ExecutorManager>();
Deshui Yu's avatar
Deshui Yu committed
71
72
        this.jobQueue = [];
        this.expRootDir = getExperimentRootDir();
73
        this.remoteExpRootDir = this.getRemoteExperimentRootDir();
Deshui Yu's avatar
Deshui Yu committed
74
75
        this.timer = timer;
        this.log = getLogger();
76
        this.trialSequenceId = -1;
SparkSnail's avatar
SparkSnail committed
77
        this.logCollection = 'none';
chicm-ms's avatar
chicm-ms committed
78
        this.log.info('Construct remote machine training service.');
Deshui Yu's avatar
Deshui Yu committed
79
80
81
82
83
84
    }

    /**
     * Loop to launch trial jobs and collect trial metrics
     */
    public async run(): Promise<void> {
SparkSnail's avatar
SparkSnail committed
85
86
        const restServer: RemoteMachineJobRestServer = component.get(RemoteMachineJobRestServer);
        await restServer.start();
87
        restServer.setEnableVersionCheck = this.versionCheck;
chicm-ms's avatar
chicm-ms committed
88
        this.log.info('Run remote machine training service.');
Deshui Yu's avatar
Deshui Yu committed
89
90
        while (!this.stopping) {
            while (this.jobQueue.length > 0) {
SparkSnail's avatar
SparkSnail committed
91
                this.updateGpuReservation();
Deshui Yu's avatar
Deshui Yu committed
92
                const trialJobId: string = this.jobQueue[0];
chicm-ms's avatar
chicm-ms committed
93
                const prepareResult: boolean = await this.prepareTrialJob(trialJobId);
Deshui Yu's avatar
Deshui Yu committed
94
95
96
97
                if (prepareResult) {
                    // Remove trial job with trialJobId from job queue
                    this.jobQueue.shift();
                } else {
98
                    // Break the while loop since no GPU resource is available right now,
Deshui Yu's avatar
Deshui Yu committed
99
100
101
                    // Wait to schedule job in next time iteration
                    break;
                }
102
            }
103
            if (restServer.getErrorMessage !== undefined) {
104
105
106
                throw new Error(restServer.getErrorMessage);
                this.stopping = true;
            }
Deshui Yu's avatar
Deshui Yu committed
107
108
            await delay(3000);
        }
chicm-ms's avatar
chicm-ms committed
109
        this.log.info('Remote machine training service exit.');
Deshui Yu's avatar
Deshui Yu committed
110
    }
111

SparkSnail's avatar
SparkSnail committed
112
    /**
113
     * give trial an executor
114
     * @param trial remote machine trial job detail
SparkSnail's avatar
SparkSnail committed
115
     */
116
    public async allocateExecutorForTrial(trial: RemoteMachineTrialJobDetail): Promise<void> {
117
        if (trial.rmMeta === undefined) {
SparkSnail's avatar
SparkSnail committed
118
119
            throw new Error(`rmMeta not set in trial ${trial.id}`);
        }
120
121
122
        const executorManager: ExecutorManager | undefined = this.machineExecutorManagerMap.get(trial.rmMeta);
        if (executorManager === undefined) {
            throw new Error(`executorManager not initialized`);
SparkSnail's avatar
SparkSnail committed
123
        }
124
125
        const shellExecutor: ShellExecutor = await executorManager.getAvailableExecutor();
        this.trialExecutorMap.set(trial.id, shellExecutor);
SparkSnail's avatar
SparkSnail committed
126
    }
127

SparkSnail's avatar
SparkSnail committed
128
129
    /**
     * If a trial is finished, release the connection resource
130
     * @param trial remote machine trial job detail
SparkSnail's avatar
SparkSnail committed
131
     */
132
    public releaseTrialExecutor(trial: RemoteMachineTrialJobDetail): void {
133
        if (trial.rmMeta === undefined) {
SparkSnail's avatar
SparkSnail committed
134
135
            throw new Error(`rmMeta not set in trial ${trial.id}`);
        }
136
137
138
        const executorManager: ExecutorManager | undefined = this.machineExecutorManagerMap.get(trial.rmMeta);
        if (executorManager === undefined) {
            throw new Error(`executorManager not initialized`);
SparkSnail's avatar
SparkSnail committed
139
        }
140
        executorManager.releaseConnection(this.trialExecutorMap.get(trial.id));
SparkSnail's avatar
SparkSnail committed
141
    }
Deshui Yu's avatar
Deshui Yu committed
142
143
144
145

    /**
     * List submitted trial jobs
     */
146
    public async listTrialJobs(): Promise<TrialJobDetail[]> {
Deshui Yu's avatar
Deshui Yu committed
147
148
149
        const jobs: TrialJobDetail[] = [];
        const deferred: Deferred<TrialJobDetail[]> = new Deferred<TrialJobDetail[]>();

150
        for (const [key,] of this.trialJobsMap) {
151
            jobs.push(await this.getTrialJob(key));
152
        }
Deshui Yu's avatar
Deshui Yu committed
153
154
155
156
157
158
159
160
161
162
163
        deferred.resolve(jobs);

        return deferred.promise;
    }

    /**
     * Get trial job detail information
     * @param trialJobId ID of trial job
     */
    public async getTrialJob(trialJobId: string): Promise<TrialJobDetail> {
        const trialJob: RemoteMachineTrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);
164
        if (trialJob === undefined) {
Deshui Yu's avatar
Deshui Yu committed
165
166
167
168
            throw new NNIError(NNIErrorNames.NOT_FOUND, `trial job id ${trialJobId} not found`);
        }
        //TO DO: add another job status, and design new job status change logic
        if (trialJob.status === 'RUNNING' || trialJob.status === 'UNKNOWN') {
169
            // Get executor where the job is running
Deshui Yu's avatar
Deshui Yu committed
170
171
172
            if (trialJob.rmMeta === undefined) {
                throw new Error(`rmMeta not set for submitted job ${trialJobId}`);
            }
173
174
175
            const executor: ShellExecutor | undefined = this.trialExecutorMap.get(trialJob.id);
            if (executor === undefined) {
                throw new Error(`Invalid job id: ${trialJobId}, cannot find executor`);
Deshui Yu's avatar
Deshui Yu committed
176
177
            }

178
            return this.updateTrialJobStatus(trialJob, executor);
Deshui Yu's avatar
Deshui Yu committed
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
        } else {
            return trialJob;
        }
    }

    /**
     * Add job metrics listener
     * @param listener callback listener
     */
    public addTrialJobMetricListener(listener: (metric: TrialJobMetric) => void): void {
        this.metricsEmitter.on('metric', listener);
    }

    /**
     * Remove job metrics listener
     * @param listener callback listener
     */
    public removeTrialJobMetricListener(listener: (metric: TrialJobMetric) => void): void {
        this.metricsEmitter.off('metric', listener);
    }

    /**
     * Submit trial job
     * @param form trial job description form
     */
204
    public async submitTrialJob(form: TrialJobApplicationForm): Promise<TrialJobDetail> {
205
        if (this.trialConfig === undefined) {
Deshui Yu's avatar
Deshui Yu committed
206
207
208
            throw new Error('trial config is not initialized');
        }

209
210
211
        // Generate trial job id(random)
        const trialJobId: string = uniqueString(5);
        const trialWorkingFolder: string = unixPathJoin(this.remoteExpRootDir, 'trials', trialJobId);
Deshui Yu's avatar
Deshui Yu committed
212

213
214
215
216
217
218
219
220
221
        const trialJobDetail: RemoteMachineTrialJobDetail = new RemoteMachineTrialJobDetail(
            trialJobId,
            'WAITING',
            Date.now(),
            trialWorkingFolder,
            form
        );
        this.jobQueue.push(trialJobId);
        this.trialJobsMap.set(trialJobId, trialJobDetail);
222

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

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

        return trialJobDetail;
239
    }
240

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

Deshui Yu's avatar
Deshui Yu committed
248
249
250
251
    /**
     * Cancel trial job
     * @param trialJobId ID of trial job
     */
QuanluZhang's avatar
QuanluZhang committed
252
    public async cancelTrialJob(trialJobId: string, isEarlyStopped: boolean = false): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
253
        const trialJob: RemoteMachineTrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);
254
        if (trialJob === undefined) {
Deshui Yu's avatar
Deshui Yu committed
255
256
257
258
            throw new Error(`trial job id ${trialJobId} not found`);
        }

        // Remove the job with trialJobId from job queue
chicm-ms's avatar
chicm-ms committed
259
        const index: number = this.jobQueue.indexOf(trialJobId);
260
        if (index >= 0) {
Deshui Yu's avatar
Deshui Yu committed
261
262
263
            this.jobQueue.splice(index, 1);
        }

264
        // Get executor where the job is running
Deshui Yu's avatar
Deshui Yu committed
265
266
        if (trialJob.rmMeta !== undefined) {
            // If the trial job is already scheduled, check its status and kill the trial process in remote machine
267
268
269
            const executor: ShellExecutor | undefined = this.trialExecutorMap.get(trialJob.id);
            if (executor === undefined) {
                throw new Error(`Invalid job id ${trialJobId}, cannot find executor`);
Deshui Yu's avatar
Deshui Yu committed
270
271
            }

chicm-ms's avatar
chicm-ms committed
272
            if (trialJob.status === 'UNKNOWN') {
273
                this.releaseTrialExecutor(trialJob);
chicm-ms's avatar
chicm-ms committed
274
275
276
277
                trialJob.status = 'USER_CANCELED';
                return
            }

Deshui Yu's avatar
Deshui Yu committed
278
279
            const jobpidPath: string = this.getJobPidPath(trialJob.id);
            try {
280
281
                // Mark the toEarlyStop tag here
                trialJob.isEarlyStopped = isEarlyStopped;
282
283
                await executor.killChildProcesses(jobpidPath);
                this.releaseTrialExecutor(trialJob);
Deshui Yu's avatar
Deshui Yu committed
284
285
286
287
288
289
            } catch (error) {
                // Not handle the error since pkill failed will not impact trial job's current status
                this.log.error(`remoteTrainingService.cancelTrialJob: ${error.message}`);
            }
        } else {
            // Job is not scheduled yet, set status to 'USER_CANCELLED' directly
QuanluZhang's avatar
QuanluZhang committed
290
291
            assert(isEarlyStopped === false, 'isEarlyStopped is not supposed to be true here.');
            trialJob.status = getJobCancelStatus(isEarlyStopped);
Deshui Yu's avatar
Deshui Yu committed
292
293
294
295
296
297
        }
    }

    /**
     * Set culster metadata
     * @param key metadata key
298
     * //1. MACHINE_LIST -- create executor of machine list
Deshui Yu's avatar
Deshui Yu committed
299
300
301
302
303
     * //2. TRIAL_CONFIG -- trial configuration
     * @param value metadata value
     */
    public async setClusterMetadata(key: string, value: string): Promise<void> {
        switch (key) {
SparkSnail's avatar
SparkSnail committed
304
305
306
            case TrialConfigMetadataKey.NNI_MANAGER_IP:
                this.nniManagerIpConfig = <NNIManagerIpConfig>JSON.parse(value);
                break;
307
            case TrialConfigMetadataKey.MACHINE_LIST:
Deshui Yu's avatar
Deshui Yu committed
308
                await this.setupConnections(value);
309
                this.gpuScheduler = new GPUScheduler(this.machineExecutorManagerMap);
Deshui Yu's avatar
Deshui Yu committed
310
                break;
chicm-ms's avatar
chicm-ms committed
311
            case TrialConfigMetadataKey.TRIAL_CONFIG: {
312
                const remoteMachineTrailConfig: TrialConfig = <TrialConfig>JSON.parse(value);
Deshui Yu's avatar
Deshui Yu committed
313
                // Parse trial config failed, throw Error
314
                if (remoteMachineTrailConfig === undefined) {
Deshui Yu's avatar
Deshui Yu committed
315
316
317
                    throw new Error('trial config parsed failed');
                }
                // codeDir is not a valid directory, throw Error
318
                if (!fs.lstatSync(remoteMachineTrailConfig.codeDir)
319
                    .isDirectory()) {
Deshui Yu's avatar
Deshui Yu committed
320
321
                    throw new Error(`codeDir ${remoteMachineTrailConfig.codeDir} is not a directory`);
                }
322
323
324
325

                // Validate to make sure codeDir doesn't have too many files
                try {
                    await validateCodeDir(remoteMachineTrailConfig.codeDir);
326
                } catch (error) {
327
                    this.log.error(error);
328

329
                    return Promise.reject(new Error(error));
330
331
                }

Deshui Yu's avatar
Deshui Yu committed
332
333
                this.trialConfig = remoteMachineTrailConfig;
                break;
chicm-ms's avatar
chicm-ms committed
334
            }
335
336
337
            case TrialConfigMetadataKey.MULTI_PHASE:
                this.isMultiPhase = (value === 'true' || value === 'True');
                break;
338
339
340
            case TrialConfigMetadataKey.VERSION_CHECK:
                this.versionCheck = (value === 'true' || value === 'True');
                break;
SparkSnail's avatar
SparkSnail committed
341
342
343
            case TrialConfigMetadataKey.LOG_COLLECTION:
                this.logCollection = value;
                break;
Deshui Yu's avatar
Deshui Yu committed
344
345
346
347
348
349
350
351
352
353
            default:
                //Reject for unknown keys
                throw new Error(`Uknown key: ${key}`);
        }
    }

    /**
     * Get culster metadata
     * @param key metadata key
     */
354
355
    public async getClusterMetadata(key: string): Promise<string> {
        return "";
Deshui Yu's avatar
Deshui Yu committed
356
    }
357

SparkSnail's avatar
SparkSnail committed
358
    /**
359
     * cleanup() has a time out of 10s to clean remote connections
SparkSnail's avatar
SparkSnail committed
360
361
     */
    public async cleanUp(): Promise<void> {
chicm-ms's avatar
chicm-ms committed
362
        this.log.info('Stopping remote machine training service...');
Deshui Yu's avatar
Deshui Yu committed
363
        this.stopping = true;
SparkSnail's avatar
SparkSnail committed
364
365
        await Promise.race([delay(10000), this.cleanupConnections()]);
    }
366

367
368
369
370
    /**
     * remove gpu reversion when job is not running
     */
    private updateGpuReservation(): void {
371
372
373
374
375
        if (this.gpuScheduler) {
            for (const [key, value] of this.trialJobsMap) {
                if (!['WAITING', 'RUNNING'].includes(value.status)) {
                    this.gpuScheduler.removeGpuReservation(key, this.trialJobsMap);
                }
376
377
378
379
            }
        }
    }

SparkSnail's avatar
SparkSnail committed
380
381
382
383
    /**
     * stop gpu_metric_collector process in remote machine and remove unused scripts
     */
    private async cleanupConnections(): Promise<void> {
384
        try {
385
            for (const [rmMeta, executorManager] of this.machineExecutorManagerMap.entries()) {
386
                const jobpidPath: string = unixPathJoin(this.getRemoteScriptsPath(rmMeta.username), 'pid');
387
388
389
390
                const executor: ShellExecutor | undefined = executorManager.getFirstExecutor();
                if (executor !== undefined) {
                    await executor.killChildProcesses(jobpidPath);
                    await executor.removeFolder(this.getRemoteScriptsPath(rmMeta.username));
SparkSnail's avatar
SparkSnail committed
391
                }
392
                executorManager.closeAllExecutor();
SparkSnail's avatar
SparkSnail committed
393
            }
394
        } catch (error) {
SparkSnail's avatar
SparkSnail committed
395
396
397
            //ignore error, this function is called to cleanup remote connections when experiment is stopping
            this.log.error(`Cleanup connection exception, error is ${error.message}`);
        }
Deshui Yu's avatar
Deshui Yu committed
398
399

        return Promise.resolve();
400
401
    }

Deshui Yu's avatar
Deshui Yu committed
402
    private async setupConnections(machineList: string): Promise<void> {
chicm-ms's avatar
chicm-ms committed
403
        this.log.debug(`Connecting to remote machines: ${machineList}`);
Deshui Yu's avatar
Deshui Yu committed
404
405
406
407
        const deferred: Deferred<void> = new Deferred<void>();
        //TO DO: verify if value's format is wrong, and json parse failed, how to handle error
        const rmMetaList: RemoteMachineMeta[] = <RemoteMachineMeta[]>JSON.parse(machineList);
        let connectedRMNum: number = 0;
SparkSnail's avatar
SparkSnail committed
408

SparkSnail's avatar
SparkSnail committed
409
        rmMetaList.forEach(async (rmMeta: RemoteMachineMeta) => {
410
            rmMeta.occupiedGpuIndexMap = new Map<number, number>();
411
412
413
414
            const executorManager: ExecutorManager = new ExecutorManager([], this.MAX_TRIAL_NUMBER_PER_EXECUTOR, rmMeta);
            const executor: ShellExecutor = await executorManager.getAvailableExecutor();
            this.machineExecutorManagerMap.set(rmMeta, executorManager);
            await this.initRemoteMachineOnConnected(rmMeta, executor);
SparkSnail's avatar
SparkSnail committed
415
416
            if (++connectedRMNum === rmMetaList.length) {
                deferred.resolve();
417
            }
Deshui Yu's avatar
Deshui Yu committed
418
        });
419

Deshui Yu's avatar
Deshui Yu committed
420
421
422
        return deferred.promise;
    }

423
424
    private async initRemoteMachineOnConnected(rmMeta: RemoteMachineMeta, executor: ShellExecutor): Promise<void> {
        // Create root working directory after executor is ready
425
        const nniRootDir: string = unixPathJoin(getRemoteTmpDir(this.remoteOS), 'nni');
426
        await executor.createFolder(this.remoteExpRootDir);
427

428
429
        // the directory to store temp scripts in remote machine
        const remoteGpuScriptCollectorDir: string = this.getRemoteScriptsPath(rmMeta.username);
430
431
        await executor.createFolder(remoteGpuScriptCollectorDir, true);
        await executor.allowPermission(false, nniRootDir, `${nniRootDir}/*`, `${nniRootDir}/scripts/*`);
432

Deshui Yu's avatar
Deshui Yu committed
433
        //Begin to execute gpu_metrics_collection scripts
434
        const script = getGpuMetricsCollectorBashScriptContent(remoteGpuScriptCollectorDir);
435
        executor.executeScript(script, false, true);
436

437
        const disposable: Rx.IDisposable = this.timer.subscribe(
438
439
440
441
            async () => {
                const cmdresult = await executor.readLastLines(unixPathJoin(remoteGpuScriptCollectorDir, 'gpu_metrics'));
                if (cmdresult !== "") {
                    rmMeta.gpuSummary = <GPUSummary>JSON.parse(cmdresult);
442
443
444
445
                    if (rmMeta.gpuSummary.gpuCount === 0) {
                        this.log.warning(`No GPU found on remote machine ${rmMeta.ip}`);
                        this.timer.unsubscribe(disposable);
                    }
Deshui Yu's avatar
Deshui Yu committed
446
447
448
449
450
451
                }
            }
        );
    }

    private async prepareTrialJob(trialJobId: string): Promise<boolean> {
chicm-ms's avatar
chicm-ms committed
452
        const deferred: Deferred<boolean> = new Deferred<boolean>();
Deshui Yu's avatar
Deshui Yu committed
453

454
        if (this.trialConfig === undefined) {
Deshui Yu's avatar
Deshui Yu committed
455
456
            throw new Error('trial config is not initialized');
        }
457
458
459
        if (this.gpuScheduler === undefined) {
            throw new Error('gpuScheduler is not initialized');
        }
Deshui Yu's avatar
Deshui Yu committed
460
461
462
463
        const trialJobDetail: RemoteMachineTrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);
        if (trialJobDetail === undefined) {
            throw new NNIError(NNIErrorNames.INVALID_JOB_DETAIL, `Invalid job detail information for trial job ${trialJobId}`);
        }
464
465
466
        // If job is not WATIING, Don't prepare and resolve true immediately
        if (trialJobDetail.status !== 'WAITING') {
            deferred.resolve(true);
467

468
469
            return deferred.promise;
        }
470
        // get an executor from scheduler
471
        const rmScheduleResult: RemoteMachineScheduleResult = this.gpuScheduler.scheduleMachine(this.trialConfig.gpuNum, trialJobDetail);
Deshui Yu's avatar
Deshui Yu committed
472
        if (rmScheduleResult.resultType === ScheduleResultType.REQUIRE_EXCEED_TOTAL) {
chicm-ms's avatar
chicm-ms committed
473
            const errorMessage: string = `Required GPU number ${this.trialConfig.gpuNum} is too large, no machine can meet`;
Deshui Yu's avatar
Deshui Yu committed
474
475
476
            this.log.error(errorMessage);
            deferred.reject();
            throw new NNIError(NNIErrorNames.RESOURCE_NOT_AVAILABLE, errorMessage);
477
        } else if (rmScheduleResult.resultType === ScheduleResultType.SUCCEED
Deshui Yu's avatar
Deshui Yu committed
478
            && rmScheduleResult.scheduleInfo !== undefined) {
chicm-ms's avatar
chicm-ms committed
479
            const rmScheduleInfo: RemoteMachineScheduleInfo = rmScheduleResult.scheduleInfo;
480
            const trialWorkingFolder: string = unixPathJoin(this.remoteExpRootDir, 'trials', trialJobId);
SparkSnail's avatar
SparkSnail committed
481
482
483

            trialJobDetail.rmMeta = rmScheduleInfo.rmMeta;

484
            await this.allocateExecutorForTrial(trialJobDetail);
Deshui Yu's avatar
Deshui Yu committed
485
            await this.launchTrialOnScheduledMachine(
486
                trialJobId, trialWorkingFolder, trialJobDetail.form, rmScheduleInfo);
Deshui Yu's avatar
Deshui Yu committed
487
488
489

            trialJobDetail.status = 'RUNNING';
            trialJobDetail.url = `file://${rmScheduleInfo.rmMeta.ip}:${trialWorkingFolder}`;
490
            trialJobDetail.startTime = Date.now();
Deshui Yu's avatar
Deshui Yu committed
491

492
            this.trialJobsMap.set(trialJobId, trialJobDetail);
Deshui Yu's avatar
Deshui Yu committed
493
            deferred.resolve(true);
494
        } else if (rmScheduleResult.resultType === ScheduleResultType.TMP_NO_AVAILABLE_GPU) {
Deshui Yu's avatar
Deshui Yu committed
495
496
497
            this.log.info(`Right now no available GPU can be allocated for trial ${trialJobId}, will try to schedule later`);
            deferred.resolve(false);
        } else {
498
            deferred.reject(`Invalid schedule resutl type: ${rmScheduleResult.resultType}`);
Deshui Yu's avatar
Deshui Yu committed
499
500
501
502
503
504
        }

        return deferred.promise;
    }

    private async launchTrialOnScheduledMachine(trialJobId: string, trialWorkingFolder: string, form: TrialJobApplicationForm,
505
        rmScheduleInfo: RemoteMachineScheduleInfo): Promise<void> {
506
        if (this.trialConfig === undefined) {
Deshui Yu's avatar
Deshui Yu committed
507
508
            throw new Error('trial config is not initialized');
        }
chicm-ms's avatar
chicm-ms committed
509
        const cudaVisibleDevice: string = rmScheduleInfo.cudaVisibleDevice;
510
511
512
        const executor: ShellExecutor | undefined = this.trialExecutorMap.get(trialJobId);
        if (executor === undefined) {
            assert(false, 'ShellExecutor is undefined.');
513
514
515
516

            // for lint
            return;
        }
517
518
519
520
521
        const trialJobDetail: RemoteMachineTrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);
        if (trialJobDetail === undefined) {
            throw new Error(`Can not get trial job detail for job: ${trialJobId}`);
        }

Deshui Yu's avatar
Deshui Yu committed
522
523
        const trialLocalTempFolder: string = path.join(this.expRootDir, 'trials-local', trialJobId);

524
525
        await executor.createFolder(trialWorkingFolder);
        await executor.createFolder(unixPathJoin(trialWorkingFolder, '.nni'));
Deshui Yu's avatar
Deshui Yu committed
526
527
528

        // RemoteMachineRunShellFormat is the run shell format string,
        // See definition in remoteMachineData.ts
SparkSnail's avatar
SparkSnail committed
529
530

        let command: string;
chicm-ms's avatar
chicm-ms committed
531
532
        // Set CUDA_VISIBLE_DEVICES environment variable based on cudaVisibleDevice
        // If no valid cudaVisibleDevice is defined, set CUDA_VISIBLE_DEVICES to empty string to hide GPU device
SparkSnail's avatar
SparkSnail committed
533
534
535
        // If gpuNum is undefined, will not set CUDA_VISIBLE_DEVICES in script
        if (this.trialConfig.gpuNum === undefined) {
            command = this.trialConfig.command;
SparkSnail's avatar
SparkSnail committed
536
        } else {
chicm-ms's avatar
chicm-ms committed
537
538
            if (typeof cudaVisibleDevice === 'string' && cudaVisibleDevice.length > 0) {
                command = `CUDA_VISIBLE_DEVICES=${cudaVisibleDevice} ${this.trialConfig.command}`;
SparkSnail's avatar
SparkSnail committed
539
540
541
            } else {
                command = `CUDA_VISIBLE_DEVICES=" " ${this.trialConfig.command}`;
            }
SparkSnail's avatar
SparkSnail committed
542
        }
543
544
        const nniManagerIp: string = this.nniManagerIpConfig ? this.nniManagerIpConfig.nniManagerIp : getIPV4Address();
        if (this.remoteRestServerPort === undefined) {
SparkSnail's avatar
SparkSnail committed
545
546
547
            const restServer: RemoteMachineJobRestServer = component.get(RemoteMachineJobRestServer);
            this.remoteRestServerPort = restServer.clusterRestServerPort;
        }
548
        const version: string = this.versionCheck ? await getVersion() : '';
SparkSnail's avatar
SparkSnail committed
549
550
551
        const runScriptTrialContent: string = String.Format(
            REMOTEMACHINE_TRIAL_COMMAND_FORMAT,
            trialWorkingFolder,
Deshui Yu's avatar
Deshui Yu committed
552
553
            trialWorkingFolder,
            trialJobId,
SparkSnail's avatar
SparkSnail committed
554
            getExperimentId(),
555
            trialJobDetail.form.sequenceId.toString(),
556
            this.isMultiPhase,
557
            unixPathJoin(trialWorkingFolder, '.nni', 'jobpid'),
SparkSnail's avatar
SparkSnail committed
558
559
560
            command,
            nniManagerIp,
            this.remoteRestServerPort,
561
            version,
SparkSnail's avatar
SparkSnail committed
562
            this.logCollection,
563
            unixPathJoin(trialWorkingFolder, '.nni', 'code')
564
        );
Deshui Yu's avatar
Deshui Yu committed
565
566

        //create tmp trial working folder locally.
567
        await execMkdir(path.join(trialLocalTempFolder, '.nni'));
Deshui Yu's avatar
Deshui Yu committed
568

SparkSnail's avatar
SparkSnail committed
569
        //create tmp trial working folder locally.
SparkSnail's avatar
SparkSnail committed
570
        await execCopydir(this.trialConfig.codeDir, trialLocalTempFolder);
chicm-ms's avatar
chicm-ms committed
571
        const installScriptContent: string = CONTAINER_INSTALL_NNI_SHELL_FORMAT;
SparkSnail's avatar
SparkSnail committed
572
573
        // Write NNI installation file to local tmp files
        await fs.promises.writeFile(path.join(trialLocalTempFolder, 'install_nni.sh'), installScriptContent, { encoding: 'utf8' });
574
        // Write file content ( run.sh and parameter.cfg ) to local tmp files
SparkSnail's avatar
SparkSnail committed
575
        await fs.promises.writeFile(path.join(trialLocalTempFolder, 'run.sh'), runScriptTrialContent, { encoding: 'utf8' });
chicm-ms's avatar
chicm-ms committed
576
        await this.writeParameterFile(trialJobId, form.hyperParameters);
Deshui Yu's avatar
Deshui Yu committed
577
        // Copy files in codeDir to remote working directory
578
        await executor.copyDirectoryToRemote(trialLocalTempFolder, trialWorkingFolder, this.remoteOS);
Deshui Yu's avatar
Deshui Yu committed
579
        // Execute command in remote machine
580
        executor.executeScript(unixPathJoin(trialWorkingFolder, 'run.sh'), true, true);
Deshui Yu's avatar
Deshui Yu committed
581
582
583
    }

    private getRmMetaByHost(host: string): RemoteMachineMeta {
584
        for (const rmMeta of this.machineExecutorManagerMap.keys()) {
Deshui Yu's avatar
Deshui Yu committed
585
586
587
588
589
590
591
            if (rmMeta.ip === host) {
                return rmMeta;
            }
        }
        throw new Error(`Host not found: ${host}`);
    }

592
    private async updateTrialJobStatus(trialJob: RemoteMachineTrialJobDetail, executor: ShellExecutor): Promise<TrialJobDetail> {
Deshui Yu's avatar
Deshui Yu committed
593
594
        const deferred: Deferred<TrialJobDetail> = new Deferred<TrialJobDetail>();
        const jobpidPath: string = this.getJobPidPath(trialJob.id);
595
        const trialReturnCodeFilePath: string = unixPathJoin(this.remoteExpRootDir, 'trials', trialJob.id, '.nni', 'code');
596
        /* eslint-disable require-atomic-updates */
Deshui Yu's avatar
Deshui Yu committed
597
        try {
598
            const isAlive = await executor.isProcessAlive(jobpidPath);
Deshui Yu's avatar
Deshui Yu committed
599
            // if the process of jobpid is not alive any more
600
601
602
603
            if (!isAlive) {
                const trialReturnCode: string = await executor.getRemoteFileContent(trialReturnCodeFilePath);
                this.log.debug(`trailjob ${trialJob.id} return code: ${trialReturnCode}`);
                const match: RegExpMatchArray | null = trialReturnCode.trim()
604
                    .match(/^(\d+)\s+(\d+)$/);
605
                if (match !== null) {
Deshui Yu's avatar
Deshui Yu committed
606
607
608
609
610
                    const { 1: code, 2: timestamp } = match;
                    // Update trial job's status based on result code
                    if (parseInt(code, 10) === 0) {
                        trialJob.status = 'SUCCEEDED';
                    } else {
611
612
613
614
615
616
                        // isEarlyStopped is never set, mean it's not cancelled by NNI, so if the process's exit code >0, mark it as FAILED
                        if (trialJob.isEarlyStopped === undefined) {
                            trialJob.status = 'FAILED';
                        } else {
                            trialJob.status = getJobCancelStatus(trialJob.isEarlyStopped);
                        }
Deshui Yu's avatar
Deshui Yu committed
617
                    }
618
                    trialJob.endTime = parseInt(timestamp, 10);
619
                    this.releaseTrialExecutor(trialJob);
Deshui Yu's avatar
Deshui Yu committed
620
                }
chicm-ms's avatar
chicm-ms committed
621
                this.log.debug(`trailJob status update: ${trialJob.id}, ${trialJob.status}`);
Deshui Yu's avatar
Deshui Yu committed
622
623
624
625
626
627
628
629
630
631
632
            }
            deferred.resolve(trialJob);
        } catch (error) {
            this.log.error(`Update job status exception, error is ${error.message}`);
            if (error instanceof NNIError && error.name === NNIErrorNames.NOT_FOUND) {
                deferred.resolve(trialJob);
            } else {
                trialJob.status = 'UNKNOWN';
                deferred.resolve(trialJob);
            }
        }
633
        /* eslint-enable require-atomic-updates */
Deshui Yu's avatar
Deshui Yu committed
634
635
636
        return deferred.promise;
    }

SparkSnail's avatar
SparkSnail committed
637
    private getRemoteScriptsPath(userName: string): string {
638
        return unixPathJoin(getRemoteTmpDir(this.remoteOS), userName, 'nni', 'scripts');
Deshui Yu's avatar
Deshui Yu committed
639
640
641
    }

    private getHostJobRemoteDir(jobId: string): string {
642
        return unixPathJoin(this.remoteExpRootDir, 'hostjobs', jobId);
Deshui Yu's avatar
Deshui Yu committed
643
644
    }

645
    private getRemoteExperimentRootDir(): string {
646
        return unixPathJoin(getRemoteTmpDir(this.remoteOS), 'nni', 'experiments', getExperimentId());
Deshui Yu's avatar
Deshui Yu committed
647
648
    }

chicm-ms's avatar
chicm-ms committed
649
    public get MetricsEmitter(): EventEmitter {
SparkSnail's avatar
SparkSnail committed
650
651
652
        return this.metricsEmitter;
    }

Deshui Yu's avatar
Deshui Yu committed
653
654
655
656
657
658
    private getJobPidPath(jobId: string): string {
        const trialJobDetail: RemoteMachineTrialJobDetail | undefined = this.trialJobsMap.get(jobId);
        if (trialJobDetail === undefined) {
            throw new NNIError(NNIErrorNames.INVALID_JOB_DETAIL, `Invalid job detail information for trial job ${jobId}`);
        }

chicm-ms's avatar
chicm-ms committed
659
        return unixPathJoin(trialJobDetail.workingDirectory, '.nni', 'jobpid');
Deshui Yu's avatar
Deshui Yu committed
660
    }
chicm-ms's avatar
chicm-ms committed
661

chicm-ms's avatar
chicm-ms committed
662
    private async writeParameterFile(trialJobId: string, hyperParameters: HyperParameters): Promise<void> {
663
664
665
        const executor: ShellExecutor | undefined = this.trialExecutorMap.get(trialJobId);
        if (executor === undefined) {
            throw new Error('ShellExecutor is undefined.');
chicm-ms's avatar
chicm-ms committed
666
667
        }

668
        const trialWorkingFolder: string = unixPathJoin(this.remoteExpRootDir, 'trials', trialJobId);
chicm-ms's avatar
chicm-ms committed
669
670
        const trialLocalTempFolder: string = path.join(this.expRootDir, 'trials-local', trialJobId);

671
        const fileName: string = generateParamFileName(hyperParameters);
chicm-ms's avatar
chicm-ms committed
672
673
674
        const localFilepath: string = path.join(trialLocalTempFolder, fileName);
        await fs.promises.writeFile(localFilepath, hyperParameters.value, { encoding: 'utf8' });

675
        await executor.copyFileToRemote(localFilepath, unixPathJoin(trialWorkingFolder, fileName));
chicm-ms's avatar
chicm-ms committed
676
    }
Deshui Yu's avatar
Deshui Yu committed
677
678
679
}

export { RemoteMachineTrainingService };