remoteMachineTrainingService.ts 34.2 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
/**
 * 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';

22
import * as assert from 'assert';
Deshui Yu's avatar
Deshui Yu committed
23
24
25
26
27
import * as cpp from 'child-process-promise';
import { EventEmitter } from 'events';
import * as fs from 'fs';
import * as os from 'os';
import * as path from 'path';
28
import { Client, ConnectConfig } from 'ssh2';
Deshui Yu's avatar
Deshui Yu committed
29
30
31
import { Deferred } from 'ts-deferred';
import { String } from 'typescript-string-operations';
import * as component from '../../common/component';
SparkSnail's avatar
SparkSnail committed
32
import { NNIError, NNIErrorNames } from '../../common/errors';
33
import { getExperimentId } from '../../common/experimentStartupInfo';
Deshui Yu's avatar
Deshui Yu committed
34
35
36
import { getLogger, Logger } from '../../common/log';
import { ObservableTimer } from '../../common/observableTimer';
import {
37
    HyperParameters, NNIManagerIpConfig, TrainingService, TrialJobApplicationForm,
38
    TrialJobDetail, TrialJobMetric
Deshui Yu's avatar
Deshui Yu committed
39
} from '../../common/trainingService';
40
41
42
43
44
45
import {
    delay, generateParamFileName, getExperimentRootDir, getIPV4Address, getJobCancelStatus, getRemoteTmpDir,
    getVersion, uniqueString, unixPathJoin
} from '../../common/utils';
import { CONTAINER_INSTALL_NNI_SHELL_FORMAT } from '../common/containerJobData';
import { GPU_INFO_COLLECTOR_FORMAT_LINUX, GPUSummary } from '../common/gpuData';
46
47
import { TrialConfig } from '../common/trialConfig';
import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
48
import { execCopydir, execMkdir, execRemove, validateCodeDir } from '../common/util';
Deshui Yu's avatar
Deshui Yu committed
49
50
import { GPUScheduler } from './gpuScheduler';
import {
51
52
53
    HOST_JOB_SHELL_FORMAT, RemoteCommandResult, REMOTEMACHINE_TRIAL_COMMAND_FORMAT, RemoteMachineMeta,
    RemoteMachineScheduleInfo, RemoteMachineScheduleResult, RemoteMachineTrialJobDetail,
    ScheduleResultType, SSHClient, SSHClientManager
Deshui Yu's avatar
Deshui Yu committed
54
} from './remoteMachineData';
SparkSnail's avatar
SparkSnail committed
55
import { RemoteMachineJobRestServer } from './remoteMachineJobRestServer';
56
import { SSHClientUtility } from './sshClientUtility';
Deshui Yu's avatar
Deshui Yu committed
57
58
59
60

/**
 * Training Service implementation for Remote Machine (Linux)
 */
SparkSnail's avatar
SparkSnail committed
61
@component.Singleton
Deshui Yu's avatar
Deshui Yu committed
62
class RemoteMachineTrainingService implements TrainingService {
63
64
65
66
67
68
    private readonly machineSSHClientMap: Map<RemoteMachineMeta, SSHClientManager>; //machine ssh client map
    private readonly trialSSHClientMap: Map<string, Client>; //trial ssh client map
    private readonly trialJobsMap: Map<string, RemoteMachineTrialJobDetail>;
    private readonly MAX_TRIAL_NUMBER_PER_SSHCONNECTION: number = 5; // every ssh client has a max trial concurrency number
    private readonly expRootDir: string;
    private readonly remoteExpRootDir: string;
69
    private trialConfig: TrialConfig | undefined;
70
71
72
    private readonly gpuScheduler: GPUScheduler;
    private readonly jobQueue: string[];
    private readonly timer: ObservableTimer;
Deshui Yu's avatar
Deshui Yu committed
73
    private stopping: boolean = false;
74
75
    private readonly metricsEmitter: EventEmitter;
    private readonly log: Logger;
76
    private isMultiPhase: boolean = false;
77
    private trialSequenceId: number;
SparkSnail's avatar
SparkSnail committed
78
    private remoteRestServerPort?: number;
79
    private readonly remoteOS: string;
SparkSnail's avatar
SparkSnail committed
80
    private nniManagerIpConfig?: NNIManagerIpConfig;
81
    private versionCheck: boolean = true;
SparkSnail's avatar
SparkSnail committed
82
    private logCollection: string;
Deshui Yu's avatar
Deshui Yu committed
83
84

    constructor(@component.Inject timer: ObservableTimer) {
85
        this.remoteOS = 'linux';
Deshui Yu's avatar
Deshui Yu committed
86
87
        this.metricsEmitter = new EventEmitter();
        this.trialJobsMap = new Map<string, RemoteMachineTrialJobDetail>();
SparkSnail's avatar
SparkSnail committed
88
89
        this.trialSSHClientMap = new Map<string, Client>();
        this.machineSSHClientMap = new Map<RemoteMachineMeta, SSHClientManager>();
Deshui Yu's avatar
Deshui Yu committed
90
91
92
        this.gpuScheduler = new GPUScheduler(this.machineSSHClientMap);
        this.jobQueue = [];
        this.expRootDir = getExperimentRootDir();
93
        this.remoteExpRootDir = this.getRemoteExperimentRootDir();
Deshui Yu's avatar
Deshui Yu committed
94
95
        this.timer = timer;
        this.log = getLogger();
96
        this.trialSequenceId = -1;
SparkSnail's avatar
SparkSnail committed
97
        this.logCollection = 'none';
chicm-ms's avatar
chicm-ms committed
98
        this.log.info('Construct remote machine training service.');
Deshui Yu's avatar
Deshui Yu committed
99
100
101
102
103
104
    }

    /**
     * Loop to launch trial jobs and collect trial metrics
     */
    public async run(): Promise<void> {
SparkSnail's avatar
SparkSnail committed
105
106
        const restServer: RemoteMachineJobRestServer = component.get(RemoteMachineJobRestServer);
        await restServer.start();
107
        restServer.setEnableVersionCheck = this.versionCheck;
chicm-ms's avatar
chicm-ms committed
108
        this.log.info('Run remote machine training service.');
Deshui Yu's avatar
Deshui Yu committed
109
110
        while (!this.stopping) {
            while (this.jobQueue.length > 0) {
SparkSnail's avatar
SparkSnail committed
111
                this.updateGpuReservation();
Deshui Yu's avatar
Deshui Yu committed
112
113
114
115
116
117
                const trialJobId: string = this.jobQueue[0];
                const prepareResult : boolean = await this.prepareTrialJob(trialJobId);
                if (prepareResult) {
                    // Remove trial job with trialJobId from job queue
                    this.jobQueue.shift();
                } else {
118
                    // Break the while loop since no GPU resource is available right now,
Deshui Yu's avatar
Deshui Yu committed
119
120
121
                    // Wait to schedule job in next time iteration
                    break;
                }
122
            }
123
            if (restServer.getErrorMessage !== undefined) {
124
125
126
                throw new Error(restServer.getErrorMessage);
                this.stopping = true;
            }
Deshui Yu's avatar
Deshui Yu committed
127
128
            await delay(3000);
        }
chicm-ms's avatar
chicm-ms committed
129
        this.log.info('Remote machine training service exit.');
Deshui Yu's avatar
Deshui Yu committed
130
    }
131

SparkSnail's avatar
SparkSnail committed
132
133
    /**
     * give trial a ssh connection
134
     * @param trial remote machine trial job detail
SparkSnail's avatar
SparkSnail committed
135
136
137
     */
    public async allocateSSHClientForTrial(trial: RemoteMachineTrialJobDetail): Promise<void> {
        const deferred: Deferred<void> = new Deferred<void>();
138
        if (trial.rmMeta === undefined) {
SparkSnail's avatar
SparkSnail committed
139
140
            throw new Error(`rmMeta not set in trial ${trial.id}`);
        }
141
142
        const sshClientManager: SSHClientManager | undefined = this.machineSSHClientMap.get(trial.rmMeta);
        if (sshClientManager === undefined) {
SparkSnail's avatar
SparkSnail committed
143
144
            throw new Error(`remoteSSHClient not initialized`);
        }
145
        const sshClient: Client = await sshClientManager.getAvailableSSHClient();
SparkSnail's avatar
SparkSnail committed
146
147
        this.trialSSHClientMap.set(trial.id, sshClient);
        deferred.resolve();
148

SparkSnail's avatar
SparkSnail committed
149
150
        return deferred.promise;
    }
151

SparkSnail's avatar
SparkSnail committed
152
153
    /**
     * If a trial is finished, release the connection resource
154
     * @param trial remote machine trial job detail
SparkSnail's avatar
SparkSnail committed
155
156
     */
    public releaseTrialSSHClient(trial: RemoteMachineTrialJobDetail): void {
157
        if (trial.rmMeta === undefined) {
SparkSnail's avatar
SparkSnail committed
158
159
            throw new Error(`rmMeta not set in trial ${trial.id}`);
        }
160
161
        const sshClientManager: SSHClientManager | undefined = this.machineSSHClientMap.get(trial.rmMeta);
        if (sshClientManager === undefined) {
SparkSnail's avatar
SparkSnail committed
162
163
164
165
            throw new Error(`sshClientManager not initialized`);
        }
        sshClientManager.releaseConnection(this.trialSSHClientMap.get(trial.id));
    }
Deshui Yu's avatar
Deshui Yu committed
166
167
168
169

    /**
     * List submitted trial jobs
     */
170
    public async listTrialJobs(): Promise<TrialJobDetail[]> {
Deshui Yu's avatar
Deshui Yu committed
171
172
173
        const jobs: TrialJobDetail[] = [];
        const deferred: Deferred<TrialJobDetail[]> = new Deferred<TrialJobDetail[]>();

174
        for (const [key, value] of this.trialJobsMap) {
175
            jobs.push(await this.getTrialJob(key));
176
        }
Deshui Yu's avatar
Deshui Yu committed
177
178
179
180
181
182
183
184
185
186
187
        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);
188
        if (trialJob === undefined) {
Deshui Yu's avatar
Deshui Yu committed
189
190
191
192
193
194
195
196
            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') {
            // Get ssh client where the job is running
            if (trialJob.rmMeta === undefined) {
                throw new Error(`rmMeta not set for submitted job ${trialJobId}`);
            }
SparkSnail's avatar
SparkSnail committed
197
            const sshClient: Client | undefined  = this.trialSSHClientMap.get(trialJob.id);
198
            if (sshClient === undefined) {
Deshui Yu's avatar
Deshui Yu committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
                throw new Error(`Invalid job id: ${trialJobId}, cannot find ssh client`);
            }

            return this.updateTrialJobStatus(trialJob, sshClient);
        } 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
     */
228
    // tslint:disable-next-line:informative-docs
229
    public async submitTrialJob(form: TrialJobApplicationForm): Promise<TrialJobDetail> {
230
        if (this.trialConfig === undefined) {
Deshui Yu's avatar
Deshui Yu committed
231
232
233
            throw new Error('trial config is not initialized');
        }

234
235
236
        // 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
237

238
239
240
241
242
243
244
245
246
        const trialJobDetail: RemoteMachineTrialJobDetail = new RemoteMachineTrialJobDetail(
            trialJobId,
            'WAITING',
            Date.now(),
            trialWorkingFolder,
            form
        );
        this.jobQueue.push(trialJobId);
        this.trialJobsMap.set(trialJobId, trialJobDetail);
247

248
        return Promise.resolve(trialJobDetail);
Deshui Yu's avatar
Deshui Yu committed
249
250
    }

251
252
253
254
255
    /**
     * Update trial job for multi-phase
     * @param trialJobId trial job id
     * @param form job application form
     */
256
    public async updateTrialJob(trialJobId: string, form: TrialJobApplicationForm): Promise<TrialJobDetail> {
chicm-ms's avatar
chicm-ms committed
257
258
259
260
        const trialJobDetail: undefined | TrialJobDetail = this.trialJobsMap.get(trialJobId);
        if (trialJobDetail === undefined) {
            throw new Error(`updateTrialJob failed: ${trialJobId} not found`);
        }
261
262
263
        const rmMeta: RemoteMachineMeta | undefined = (<RemoteMachineTrialJobDetail>trialJobDetail).rmMeta;
        if (rmMeta !== undefined) {
            await this.writeParameterFile(trialJobId, form.hyperParameters, rmMeta);
chicm-ms's avatar
chicm-ms committed
264
        } else {
265
            throw new Error(`updateTrialJob failed: ${trialJobId} rmMeta not found`);
chicm-ms's avatar
chicm-ms committed
266
267
268
        }

        return trialJobDetail;
269
    }
270

271
272
273
274
    /**
     * Is multiphase job supported in current training service
     */
    public get isMultiPhaseJobSupported(): boolean {
275
        return true;
276
277
    }

Deshui Yu's avatar
Deshui Yu committed
278
279
280
281
    /**
     * Cancel trial job
     * @param trialJobId ID of trial job
     */
282
    // tslint:disable:informative-docs no-unsafe-any
QuanluZhang's avatar
QuanluZhang committed
283
    public async cancelTrialJob(trialJobId: string, isEarlyStopped: boolean = false): Promise<void> {
Deshui Yu's avatar
Deshui Yu committed
284
285
        const deferred: Deferred<void> = new Deferred<void>();
        const trialJob: RemoteMachineTrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);
286
        if (trialJob === undefined) {
Deshui Yu's avatar
Deshui Yu committed
287
288
289
290
291
292
            deferred.reject();
            throw new Error(`trial job id ${trialJobId} not found`);
        }

        // Remove the job with trialJobId from job queue
        const index : number = this.jobQueue.indexOf(trialJobId);
293
        if (index >= 0) {
Deshui Yu's avatar
Deshui Yu committed
294
295
296
297
298
299
            this.jobQueue.splice(index, 1);
        }

        // Get ssh client where the job is running
        if (trialJob.rmMeta !== undefined) {
            // If the trial job is already scheduled, check its status and kill the trial process in remote machine
SparkSnail's avatar
SparkSnail committed
300
            const sshClient: Client | undefined = this.trialSSHClientMap.get(trialJob.id);
301
            if (sshClient === undefined) {
Deshui Yu's avatar
Deshui Yu committed
302
303
304
305
306
307
                deferred.reject();
                throw new Error(`Invalid job id ${trialJobId}, cannot find ssh client`);
            }

            const jobpidPath: string = this.getJobPidPath(trialJob.id);
            try {
308
309
                // Mark the toEarlyStop tag here
                trialJob.isEarlyStopped = isEarlyStopped;
Deshui Yu's avatar
Deshui Yu committed
310
                await SSHClientUtility.remoteExeCommand(`pkill -P \`cat ${jobpidPath}\``, sshClient);
SparkSnail's avatar
SparkSnail committed
311
                this.releaseTrialSSHClient(trialJob);
Deshui Yu's avatar
Deshui Yu committed
312
313
314
315
316
317
            } 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
318
319
            assert(isEarlyStopped === false, 'isEarlyStopped is not supposed to be true here.');
            trialJob.status = getJobCancelStatus(isEarlyStopped);
Deshui Yu's avatar
Deshui Yu committed
320
321
322
323
324
325
326
327
328
329
330
331
        }
    }

    /**
     * Set culster metadata
     * @param key metadata key
     * //1. MACHINE_LIST -- create ssh client connect of machine list
     * //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
332
333
334
            case TrialConfigMetadataKey.NNI_MANAGER_IP:
                this.nniManagerIpConfig = <NNIManagerIpConfig>JSON.parse(value);
                break;
335
            case TrialConfigMetadataKey.MACHINE_LIST:
Deshui Yu's avatar
Deshui Yu committed
336
                await this.setupConnections(value);
SparkSnail's avatar
SparkSnail committed
337
                //remove local temp files
338
                await execRemove(this.getLocalGpuMetricCollectorDir());
Deshui Yu's avatar
Deshui Yu committed
339
                break;
340
341
            case TrialConfigMetadataKey.TRIAL_CONFIG:
                const remoteMachineTrailConfig: TrialConfig = <TrialConfig>JSON.parse(value);
Deshui Yu's avatar
Deshui Yu committed
342
                // Parse trial config failed, throw Error
343
                if (remoteMachineTrailConfig === undefined) {
Deshui Yu's avatar
Deshui Yu committed
344
345
346
                    throw new Error('trial config parsed failed');
                }
                // codeDir is not a valid directory, throw Error
347
348
349
                // tslint:disable-next-line:non-literal-fs-path
                if (!fs.lstatSync(remoteMachineTrailConfig.codeDir)
                  .isDirectory()) {
Deshui Yu's avatar
Deshui Yu committed
350
351
                    throw new Error(`codeDir ${remoteMachineTrailConfig.codeDir} is not a directory`);
                }
352
353
354
355

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

359
                    return Promise.reject(new Error(error));
360
361
                }

Deshui Yu's avatar
Deshui Yu committed
362
363
                this.trialConfig = remoteMachineTrailConfig;
                break;
364
365
366
            case TrialConfigMetadataKey.MULTI_PHASE:
                this.isMultiPhase = (value === 'true' || value === 'True');
                break;
367
368
369
            case TrialConfigMetadataKey.VERSION_CHECK:
                this.versionCheck = (value === 'true' || value === 'True');
                break;
SparkSnail's avatar
SparkSnail committed
370
371
372
            case TrialConfigMetadataKey.LOG_COLLECTION:
                this.logCollection = value;
                break;
Deshui Yu's avatar
Deshui Yu committed
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
            default:
                //Reject for unknown keys
                throw new Error(`Uknown key: ${key}`);
        }
    }

    /**
     * Get culster metadata
     * @param key metadata key
     */
    public getClusterMetadata(key: string): Promise<string> {
        const deferred: Deferred<string> = new Deferred<string>();

        return deferred.promise;
    }
388

SparkSnail's avatar
SparkSnail committed
389
    /**
390
     * cleanup() has a time out of 10s to clean remote connections
SparkSnail's avatar
SparkSnail committed
391
392
     */
    public async cleanUp(): Promise<void> {
chicm-ms's avatar
chicm-ms committed
393
        this.log.info('Stopping remote machine training service...');
Deshui Yu's avatar
Deshui Yu committed
394
        this.stopping = true;
SparkSnail's avatar
SparkSnail committed
395
396
        await Promise.race([delay(10000), this.cleanupConnections()]);
    }
397

398
399
400
401
402
403
404
405
406
407
408
    /**
     * remove gpu reversion when job is not running
     */
    private updateGpuReservation(): void {
        for (const [key, value] of this.trialJobsMap) {
            if (!['WAITING', 'RUNNING'].includes(value.status)) {
                this.gpuScheduler.removeGpuReservation(key, this.trialJobsMap);
            }
        }
    }

SparkSnail's avatar
SparkSnail committed
409
410
411
412
    /**
     * stop gpu_metric_collector process in remote machine and remove unused scripts
     */
    private async cleanupConnections(): Promise<void> {
413
        try {
SparkSnail's avatar
SparkSnail committed
414
            for (const [rmMeta, sshClientManager] of this.machineSSHClientMap.entries()) {
415
416
417
                const jobpidPath: string = unixPathJoin(this.getRemoteScriptsPath(rmMeta.username), 'pid');
                const client: Client | undefined = sshClientManager.getFirstSSHClient();
                if (client !== undefined) {
SparkSnail's avatar
SparkSnail committed
418
419
420
421
                    await SSHClientUtility.remoteExeCommand(`pkill -P \`cat ${jobpidPath}\``, client);
                    await SSHClientUtility.remoteExeCommand(`rm -rf ${this.getRemoteScriptsPath(rmMeta.username)}`, client);
                }
                sshClientManager.closeAllSSHClient();
SparkSnail's avatar
SparkSnail committed
422
            }
423
        } catch (error) {
SparkSnail's avatar
SparkSnail committed
424
425
426
            //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
427
428

        return Promise.resolve();
429
430
    }

SparkSnail's avatar
SparkSnail committed
431
432
433
434
    /**
     * Generate gpu metric collector directory to store temp gpu metric collector script files
     */
    private getLocalGpuMetricCollectorDir(): string {
435
436
        const userName: string = path.basename(os.homedir()); //get current user name of os

437
        return path.join(os.tmpdir(), userName, 'nni', 'scripts');
SparkSnail's avatar
SparkSnail committed
438
439
440
    }

    /**
441
442
     * Generate gpu metric collector shell script in local machine,
     * used to run in remote machine, and will be deleted after uploaded from local.
SparkSnail's avatar
SparkSnail committed
443
444
     */
    private async generateGpuMetricsCollectorScript(userName: string): Promise<void> {
445
        const gpuMetricCollectorScriptFolder : string = this.getLocalGpuMetricCollectorDir();
446
        await execMkdir(path.join(gpuMetricCollectorScriptFolder, userName));
SparkSnail's avatar
SparkSnail committed
447
        //generate gpu_metrics_collector.sh
448
449
450
        const gpuMetricsCollectorScriptPath: string = path.join(gpuMetricCollectorScriptFolder, userName, 'gpu_metrics_collector.sh');
        // This directory is used to store gpu_metrics and pid created by script
        const remoteGPUScriptsDir: string = this.getRemoteScriptsPath(userName);
SparkSnail's avatar
SparkSnail committed
451
        const gpuMetricsCollectorScriptContent: string = String.Format(
452
453
            GPU_INFO_COLLECTOR_FORMAT_LINUX,
            remoteGPUScriptsDir,
454
            unixPathJoin(remoteGPUScriptsDir, 'pid')
SparkSnail's avatar
SparkSnail committed
455
456
        );
        await fs.promises.writeFile(gpuMetricsCollectorScriptPath, gpuMetricsCollectorScriptContent, { encoding: 'utf8' });
Deshui Yu's avatar
Deshui Yu committed
457
458
459
    }

    private async setupConnections(machineList: string): Promise<void> {
chicm-ms's avatar
chicm-ms committed
460
        this.log.debug(`Connecting to remote machines: ${machineList}`);
Deshui Yu's avatar
Deshui Yu committed
461
462
463
464
        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
465

SparkSnail's avatar
SparkSnail committed
466
        rmMetaList.forEach(async (rmMeta: RemoteMachineMeta) => {
467
            rmMeta.occupiedGpuIndexMap = new Map<number, number>();
468
469
            const sshClientManager: SSHClientManager = new SSHClientManager([], this.MAX_TRIAL_NUMBER_PER_SSHCONNECTION, rmMeta);
            const sshClient: Client = await sshClientManager.getAvailableSSHClient();
SparkSnail's avatar
SparkSnail committed
470
471
472
473
            this.machineSSHClientMap.set(rmMeta, sshClientManager);
            await this.initRemoteMachineOnConnected(rmMeta, sshClient);
            if (++connectedRMNum === rmMetaList.length) {
                deferred.resolve();
474
            }
Deshui Yu's avatar
Deshui Yu committed
475
        });
476

Deshui Yu's avatar
Deshui Yu committed
477
478
479
480
481
        return deferred.promise;
    }

    private async initRemoteMachineOnConnected(rmMeta: RemoteMachineMeta, conn: Client): Promise<void> {
        // Create root working directory after ssh connection is ready
482
483
        // generate gpu script in local machine first, will copy to remote machine later
        await this.generateGpuMetricsCollectorScript(rmMeta.username);
484
        const nniRootDir: string = unixPathJoin(getRemoteTmpDir(this.remoteOS), 'nni');
Deshui Yu's avatar
Deshui Yu committed
485
        await SSHClientUtility.remoteExeCommand(`mkdir -p ${this.remoteExpRootDir}`, conn);
486

Deshui Yu's avatar
Deshui Yu committed
487
        // Copy NNI scripts to remote expeirment working directory
SparkSnail's avatar
SparkSnail committed
488
        const localGpuScriptCollectorDir: string = this.getLocalGpuMetricCollectorDir();
489
490
        // the directory to store temp scripts in remote machine
        const remoteGpuScriptCollectorDir: string = this.getRemoteScriptsPath(rmMeta.username);
SparkSnail's avatar
SparkSnail committed
491
        await SSHClientUtility.remoteExeCommand(`mkdir -p ${remoteGpuScriptCollectorDir}`, conn);
Deshui Yu's avatar
Deshui Yu committed
492
        await SSHClientUtility.remoteExeCommand(`chmod 777 ${nniRootDir} ${nniRootDir}/* ${nniRootDir}/scripts/*`, conn);
SparkSnail's avatar
SparkSnail committed
493
        //copy gpu_metrics_collector.sh to remote
494
495
        await SSHClientUtility.copyFileToRemote(path.join(localGpuScriptCollectorDir, rmMeta.username, 'gpu_metrics_collector.sh'),
                                                unixPathJoin(remoteGpuScriptCollectorDir, 'gpu_metrics_collector.sh'), conn);
496

Deshui Yu's avatar
Deshui Yu committed
497
        //Begin to execute gpu_metrics_collection scripts
498
        // tslint:disable-next-line: no-floating-promises
499
        SSHClientUtility.remoteExeCommand(`bash ${unixPathJoin(remoteGpuScriptCollectorDir, 'gpu_metrics_collector.sh')}`, conn);
500

501
        const disposable: Rx.IDisposable = this.timer.subscribe(
Deshui Yu's avatar
Deshui Yu committed
502
503
            async (tick: number) => {
                const cmdresult: RemoteCommandResult = await SSHClientUtility.remoteExeCommand(
504
                    `tail -n 1 ${unixPathJoin(remoteGpuScriptCollectorDir, 'gpu_metrics')}`, conn);
505
                if (cmdresult !== undefined && cmdresult.stdout !== undefined) {
Deshui Yu's avatar
Deshui Yu committed
506
                    rmMeta.gpuSummary = <GPUSummary>JSON.parse(cmdresult.stdout);
507
508
509
510
                    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
511
512
513
514
515
516
517
518
                }
            }
        );
    }

    private async prepareTrialJob(trialJobId: string): Promise<boolean> {
        const deferred : Deferred<boolean> = new Deferred<boolean>();

519
        if (this.trialConfig === undefined) {
Deshui Yu's avatar
Deshui Yu committed
520
521
522
523
524
525
            throw new Error('trial config is not initialized');
        }
        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}`);
        }
526
527
528
        // If job is not WATIING, Don't prepare and resolve true immediately
        if (trialJobDetail.status !== 'WAITING') {
            deferred.resolve(true);
529

530
531
            return deferred.promise;
        }
Deshui Yu's avatar
Deshui Yu committed
532
        // get an ssh client from scheduler
533
        const rmScheduleResult: RemoteMachineScheduleResult = this.gpuScheduler.scheduleMachine(this.trialConfig.gpuNum, trialJobDetail);
Deshui Yu's avatar
Deshui Yu committed
534
535
536
537
538
        if (rmScheduleResult.resultType === ScheduleResultType.REQUIRE_EXCEED_TOTAL) {
            const errorMessage : string = `Required GPU number ${this.trialConfig.gpuNum} is too large, no machine can meet`;
            this.log.error(errorMessage);
            deferred.reject();
            throw new NNIError(NNIErrorNames.RESOURCE_NOT_AVAILABLE, errorMessage);
539
        } else if (rmScheduleResult.resultType === ScheduleResultType.SUCCEED
Deshui Yu's avatar
Deshui Yu committed
540
541
            && rmScheduleResult.scheduleInfo !== undefined) {
            const rmScheduleInfo : RemoteMachineScheduleInfo = rmScheduleResult.scheduleInfo;
542
            const trialWorkingFolder: string = unixPathJoin(this.remoteExpRootDir, 'trials', trialJobId);
SparkSnail's avatar
SparkSnail committed
543
544
545
546

            trialJobDetail.rmMeta = rmScheduleInfo.rmMeta;

            await this.allocateSSHClientForTrial(trialJobDetail);
Deshui Yu's avatar
Deshui Yu committed
547
            await this.launchTrialOnScheduledMachine(
548
                trialJobId, trialWorkingFolder, trialJobDetail.form, rmScheduleInfo);
Deshui Yu's avatar
Deshui Yu committed
549
550
551

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

554
            this.trialJobsMap.set(trialJobId, trialJobDetail);
Deshui Yu's avatar
Deshui Yu committed
555
            deferred.resolve(true);
556
        } else if (rmScheduleResult.resultType === ScheduleResultType.TMP_NO_AVAILABLE_GPU) {
Deshui Yu's avatar
Deshui Yu committed
557
558
559
            this.log.info(`Right now no available GPU can be allocated for trial ${trialJobId}, will try to schedule later`);
            deferred.resolve(false);
        } else {
560
            deferred.reject(`Invalid schedule resutl type: ${rmScheduleResult.resultType}`);
Deshui Yu's avatar
Deshui Yu committed
561
562
563
564
565
566
567
        }

        return deferred.promise;
    }

    private async launchTrialOnScheduledMachine(trialJobId: string, trialWorkingFolder: string, form: TrialJobApplicationForm,
                                                rmScheduleInfo: RemoteMachineScheduleInfo): Promise<void> {
568
        if (this.trialConfig === undefined) {
Deshui Yu's avatar
Deshui Yu committed
569
570
571
            throw new Error('trial config is not initialized');
        }
        const cuda_visible_device: string = rmScheduleInfo.cuda_visible_device;
SparkSnail's avatar
SparkSnail committed
572
        const sshClient: Client | undefined = this.trialSSHClientMap.get(trialJobId);
573
574
575
576
577
578
        if (sshClient === undefined) {
            assert(false, 'sshClient is undefined.');

            // for lint
            return;
        }
579
580
581
582
583
        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
584
585
586
        const trialLocalTempFolder: string = path.join(this.expRootDir, 'trials-local', trialJobId);

        await SSHClientUtility.remoteExeCommand(`mkdir -p ${trialWorkingFolder}`, sshClient);
587
        await SSHClientUtility.remoteExeCommand(`mkdir -p ${unixPathJoin(trialWorkingFolder, '.nni')}`, sshClient);
Deshui Yu's avatar
Deshui Yu committed
588
589
590

        // RemoteMachineRunShellFormat is the run shell format string,
        // See definition in remoteMachineData.ts
SparkSnail's avatar
SparkSnail committed
591
592
593
594

        let command: string;
        // Set CUDA_VISIBLE_DEVICES environment variable based on cuda_visible_device
        // If no valid cuda_visible_device is defined, set CUDA_VISIBLE_DEVICES to empty string to hide GPU device
SparkSnail's avatar
SparkSnail committed
595
596
597
        // 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
598
        } else {
SparkSnail's avatar
SparkSnail committed
599
600
601
602
603
            if (typeof cuda_visible_device === 'string' && cuda_visible_device.length > 0) {
                command = `CUDA_VISIBLE_DEVICES=${cuda_visible_device} ${this.trialConfig.command}`;
            } else {
                command = `CUDA_VISIBLE_DEVICES=" " ${this.trialConfig.command}`;
            }
SparkSnail's avatar
SparkSnail committed
604
        }
605
606
607
        // tslint:disable-next-line: strict-boolean-expressions
        const nniManagerIp: string = this.nniManagerIpConfig ? this.nniManagerIpConfig.nniManagerIp : getIPV4Address();
        if (this.remoteRestServerPort === undefined) {
SparkSnail's avatar
SparkSnail committed
608
609
610
            const restServer: RemoteMachineJobRestServer = component.get(RemoteMachineJobRestServer);
            this.remoteRestServerPort = restServer.clusterRestServerPort;
        }
611
        const version: string = this.versionCheck ? await getVersion() : '';
SparkSnail's avatar
SparkSnail committed
612
613
614
        const runScriptTrialContent: string = String.Format(
            REMOTEMACHINE_TRIAL_COMMAND_FORMAT,
            trialWorkingFolder,
Deshui Yu's avatar
Deshui Yu committed
615
616
            trialWorkingFolder,
            trialJobId,
SparkSnail's avatar
SparkSnail committed
617
            getExperimentId(),
618
            trialJobDetail.form.sequenceId.toString(),
619
            this.isMultiPhase,
620
            unixPathJoin(trialWorkingFolder, '.nni', 'jobpid'),
SparkSnail's avatar
SparkSnail committed
621
622
623
            command,
            nniManagerIp,
            this.remoteRestServerPort,
624
            version,
SparkSnail's avatar
SparkSnail committed
625
            this.logCollection,
626
            unixPathJoin(trialWorkingFolder, '.nni', 'code')
627
        );
Deshui Yu's avatar
Deshui Yu committed
628
629

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

SparkSnail's avatar
SparkSnail committed
632
        //create tmp trial working folder locally.
633
        await execCopydir(path.join(this.trialConfig.codeDir, '*'), trialLocalTempFolder);
SparkSnail's avatar
SparkSnail committed
634
635
636
        const installScriptContent : string = CONTAINER_INSTALL_NNI_SHELL_FORMAT;
        // Write NNI installation file to local tmp files
        await fs.promises.writeFile(path.join(trialLocalTempFolder, 'install_nni.sh'), installScriptContent, { encoding: 'utf8' });
637
        // Write file content ( run.sh and parameter.cfg ) to local tmp files
SparkSnail's avatar
SparkSnail committed
638
        await fs.promises.writeFile(path.join(trialLocalTempFolder, 'run.sh'), runScriptTrialContent, { encoding: 'utf8' });
chicm-ms's avatar
chicm-ms committed
639
        await this.writeParameterFile(trialJobId, form.hyperParameters, rmScheduleInfo.rmMeta);
Deshui Yu's avatar
Deshui Yu committed
640
        // Copy files in codeDir to remote working directory
SparkSnail's avatar
SparkSnail committed
641
        await SSHClientUtility.copyDirectoryToRemote(trialLocalTempFolder, trialWorkingFolder, sshClient, this.remoteOS);
Deshui Yu's avatar
Deshui Yu committed
642
        // Execute command in remote machine
643
        // tslint:disable-next-line: no-floating-promises
644
        SSHClientUtility.remoteExeCommand(`bash ${unixPathJoin(trialWorkingFolder, 'run.sh')}`, sshClient);
Deshui Yu's avatar
Deshui Yu committed
645
646
647
648
649
650
651
652
653
654
655
656
657
658
    }

    private getRmMetaByHost(host: string): RemoteMachineMeta {
        for (const [rmMeta, client] of this.machineSSHClientMap.entries()) {
            if (rmMeta.ip === host) {
                return rmMeta;
            }
        }
        throw new Error(`Host not found: ${host}`);
    }

    private async updateTrialJobStatus(trialJob: RemoteMachineTrialJobDetail, sshClient: Client): Promise<TrialJobDetail> {
        const deferred: Deferred<TrialJobDetail> = new Deferred<TrialJobDetail>();
        const jobpidPath: string = this.getJobPidPath(trialJob.id);
659
        const trialReturnCodeFilePath: string = unixPathJoin(this.remoteExpRootDir, 'trials', trialJob.id, '.nni', 'code');
Deshui Yu's avatar
Deshui Yu committed
660
661
662
663
664
665
        try {
            const killResult: number = (await SSHClientUtility.remoteExeCommand(`kill -0 \`cat ${jobpidPath}\``, sshClient)).exitCode;
            // if the process of jobpid is not alive any more
            if (killResult !== 0) {
                const trailReturnCode: string = await SSHClientUtility.getRemoteFileContent(trialReturnCodeFilePath, sshClient);
                this.log.debug(`trailjob ${trialJob.id} return code: ${trailReturnCode}`);
666
667
668
                const match: RegExpMatchArray | null = trailReturnCode.trim()
                  .match(/^(\d+)\s+(\d+)$/);
                if (match !== null) {
Deshui Yu's avatar
Deshui Yu committed
669
670
671
672
673
                    const { 1: code, 2: timestamp } = match;
                    // Update trial job's status based on result code
                    if (parseInt(code, 10) === 0) {
                        trialJob.status = 'SUCCEEDED';
                    } else {
674
675
676
677
678
679
                        // 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
680
                    }
681
                    trialJob.endTime = parseInt(timestamp, 10);
SparkSnail's avatar
SparkSnail committed
682
                    this.releaseTrialSSHClient(trialJob);
Deshui Yu's avatar
Deshui Yu committed
683
                }
chicm-ms's avatar
chicm-ms committed
684
                this.log.debug(`trailJob status update: ${trialJob.id}, ${trialJob.status}`);
Deshui Yu's avatar
Deshui Yu committed
685
686
687
688
689
690
691
692
693
694
695
            }
            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);
            }
        }
696

Deshui Yu's avatar
Deshui Yu committed
697
698
699
        return deferred.promise;
    }

SparkSnail's avatar
SparkSnail committed
700
    private getRemoteScriptsPath(userName: string): string {
701
        return unixPathJoin(getRemoteTmpDir(this.remoteOS), userName, 'nni', 'scripts');
Deshui Yu's avatar
Deshui Yu committed
702
703
704
    }

    private getHostJobRemoteDir(jobId: string): string {
705
        return unixPathJoin(this.remoteExpRootDir, 'hostjobs', jobId);
Deshui Yu's avatar
Deshui Yu committed
706
707
    }

708
    private getRemoteExperimentRootDir(): string {
709
        return unixPathJoin(getRemoteTmpDir(this.remoteOS), 'nni', 'experiments', getExperimentId());
Deshui Yu's avatar
Deshui Yu committed
710
711
    }

SparkSnail's avatar
SparkSnail committed
712
713
714
715
    public get MetricsEmitter() : EventEmitter {
        return this.metricsEmitter;
    }

Deshui Yu's avatar
Deshui Yu committed
716
717
718
719
720
721
722
    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}`);
        }

        let jobpidPath: string;
723
        jobpidPath = unixPathJoin(trialJobDetail.workingDirectory, '.nni', 'jobpid');
Deshui Yu's avatar
Deshui Yu committed
724
725
726

        return jobpidPath;
    }
chicm-ms's avatar
chicm-ms committed
727
728

    private async writeParameterFile(trialJobId: string, hyperParameters: HyperParameters, rmMeta: RemoteMachineMeta): Promise<void> {
SparkSnail's avatar
SparkSnail committed
729
        const sshClient: Client | undefined = this.trialSSHClientMap.get(trialJobId);
chicm-ms's avatar
chicm-ms committed
730
731
732
733
        if (sshClient === undefined) {
            throw new Error('sshClient is undefined.');
        }

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

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

741
        await SSHClientUtility.copyFileToRemote(localFilepath, unixPathJoin(trialWorkingFolder, fileName), sshClient);
chicm-ms's avatar
chicm-ms committed
742
    }
Deshui Yu's avatar
Deshui Yu committed
743
744
745
}

export { RemoteMachineTrainingService };