trialDispatcher.ts 44.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

'use strict';

import { EventEmitter } from 'events';
import * as fs from 'fs';
import * as path from 'path';
import { Writable } from 'stream';
10
import { Container, Scope } from 'typescript-ioc';
11
12
import { String } from 'typescript-string-operations';
import * as component from '../../common/component';
13
import { NNIError, NNIErrorNames, MethodNotImplementedError } from '../../common/errors';
14
import { getBasePort, getExperimentId } from '../../common/experimentStartupInfo';
15
import { getLogger, Logger } from '../../common/log';
16
import { TrainingService, TrialJobApplicationForm, TrialJobMetric, TrialJobStatus, LogType } from '../../common/trainingService';
17
import { delay, getExperimentRootDir, getIPV4Address, getLogLevel, getVersion, mkDirPSync, randomSelect, uniqueString } from '../../common/utils';
18
import { ExperimentConfig, SharedStorageConfig } from '../../common/experimentConfig';
19
import { GPU_INFO, INITIALIZED, KILL_TRIAL_JOB, NEW_TRIAL_JOB, REPORT_METRIC_DATA, SEND_TRIAL_JOB_PARAMETER, STDOUT, TRIAL_END, VERSION_CHECK } from '../../core/commands';
20
import { ScheduleResultType } from '../../training_service/common/gpuData';
21
import { CONTAINER_INSTALL_NNI_SHELL_FORMAT } from '../common/containerJobData';
Ni Hao's avatar
Ni Hao committed
22
import { CONTAINER_INSTALL_NNI_SHELL_FORMAT_FOR_WIN } from '../common/containerJobData';
23
24
25
import { TrialConfig } from '../common/trialConfig';
import { validateCodeDir } from '../common/util';
import { Command, CommandChannel } from './commandChannel';
J-shang's avatar
J-shang committed
26
import { EnvironmentInformation, EnvironmentService, NodeInformation, RunnerSettings, TrialGpuSummary } from './environment';
27
import { createEnvironmentService } from './environments/environmentServiceFactory';
28
import { GpuScheduler } from './gpuScheduler';
SparkSnail's avatar
SparkSnail committed
29
import { MountedStorageService } from './storages/mountedStorageService';
30
import { StorageService } from './storageService';
31
import { SharedStorageService } from './sharedStorage';
32
33
import { NFSSharedStorageService } from './shared_storages/nfsStorageService'
import { AzureBlobSharedStorageService } from './shared_storages/azureblobStorageService'
34
35
36
37
38
39
40
41
import { TrialDetail } from './trial';

/**
 * It uses to manage jobs on training platforms 
 * and expose trial as trial job to upper level.
**/
@component.Singleton
class TrialDispatcher implements TrainingService {
42
43
    private log: Logger;
    private isDeveloping: boolean = false;
44
45
    private stopping: boolean = false;

46
47
48
    private metricsEmitter: EventEmitter;
    private experimentId: string;
    private experimentRootDir: string;
49
50
51
52
53

    private enableVersionCheck: boolean = true;

    private trialConfig: TrialConfig | undefined;

54
55
    private trials: Map<string, TrialDetail>;
    private environments: Map<string, EnvironmentInformation>;
56
57
58
59
60
61
    // make public for ut
    public environmentServiceList: EnvironmentService[] = [];
    public commandChannelSet: Set<CommandChannel>;
    public commandEmitter: EventEmitter;
    public environmentMaintenceLoopInterval: number = -1;

62
63
64
65
66
67
68
69
70
    // uses to accelerate trial manager loop
    // true means there is updates, and trial loop should run a cycle immediately.
    private shouldUpdateTrials: boolean = true;
    // uses to decide environment assign strategy.
    // true means use gpu scheduler to decide if there is free resource for new trial.
    // false means one env run one trial in same time.
    private enableGpuScheduler: boolean = false;
    // uses to save if user like to reuse environment
    private reuseEnvironment: boolean = true;
71
    private logCollection: string = 'none';
72
73
74
75
76
77
78

    private gpuScheduler: GpuScheduler;

    // uses to reduce log count.
    private isLoggedNoMoreEnvironment: boolean = false;
    private isLoggedNoGpuAvailable: boolean = false;

79
80
81
82
    // uses to mark whether to use shared storage
    private useSharedStorage: boolean = false;
    private fileCopyCompleted: boolean = false;

83
84
    private config: ExperimentConfig;

85
86
87
88
89
90
91
92
    public static async construct(config: ExperimentConfig): Promise<TrialDispatcher> {
        const instance = new TrialDispatcher(config);
        await instance.asyncConstructor(config);
        return instance;
    }

    private constructor(config: ExperimentConfig) {
        this.log = getLogger('TrialDispatcher');
93
94
95
96
        this.trials = new Map<string, TrialDetail>();
        this.environments = new Map<string, EnvironmentInformation>();
        this.metricsEmitter = new EventEmitter();
        this.experimentId = getExperimentId();
SparkSnail's avatar
SparkSnail committed
97
        this.experimentRootDir = getExperimentRootDir();
98
        this.commandChannelSet = new Set<CommandChannel>();
99
100
101
102
103
104
105
106

        const logLevel = getLogLevel();
        this.log.debug(`current folder ${__dirname}`);
        // different source folder in Linux and Windows
        if (logLevel == "debug" && (fs.existsSync("../../../src/nni_manager") || __dirname.endsWith("src\\nni_manager\\dist\\training_service\\reusable"))) {
            this.log.debug("log level is debug, and exist code folder, so set to developing mode.");
            this.isDeveloping = true;
        }
107

108
109
        this.commandEmitter = new EventEmitter();

110
        this.gpuScheduler = new GpuScheduler();
111
112
113
114
115
116
117

        this.config = config;

        this.enableGpuScheduler = !!config.trialGpuNumber;
        if (this.enableGpuScheduler) {
            this.log.info(`TrialDispatcher: GPU scheduler is enabled.`)
        }
118
    }
119

120
121
    private async asyncConstructor(config: ExperimentConfig): Promise<void> {
        await validateCodeDir(config.trialCodeDirectory);
122

123
124
125
        const serviceConfigs = Array.isArray(config.trainingService) ? config.trainingService : [ config.trainingService ];
        const servicePromises = serviceConfigs.map(serviceConfig => createEnvironmentService(serviceConfig.platform, config));
        this.environmentServiceList = await Promise.all(servicePromises);
126
127
128
129
130
131
132
133
134
135
136

        this.environmentMaintenceLoopInterval = Math.max(
            ...this.environmentServiceList.map((env) => env.environmentMaintenceLoopInterval)
        );

        for (const env of this.environmentServiceList) {
            env.initCommandChannel(this.commandEmitter);
            this.commandChannelSet.add(env.getCommandChannel);
        }

        if (this.config.sharedStorage !== undefined) {
137
            await this.initializeSharedStorage(this.config.sharedStorage);
138
        }
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
    }

    public async listTrialJobs(): Promise<TrialDetail[]> {
        const trials: TrialDetail[] = [];

        for (const key of this.trials.keys()) {
            trials.push(await this.getTrialJob(key));
        }

        return trials;
    }

    public async getTrialJob(trialJobId: string): Promise<TrialDetail> {
        const trial: TrialDetail | undefined = this.trials.get(trialJobId);
        if (trial === undefined) {
            throw new Error(`trial job ${trialJobId} not found`);
        }

        return trial;
    }

160
161
162
163
    public async getTrialLog(_trialJobId: string, _logType: LogType): Promise<string> {
        throw new MethodNotImplementedError();
    }

164
165
166
    public async submitTrialJob(form: TrialJobApplicationForm): Promise<TrialDetail> {
        const trialId: string = uniqueString(5);

167
        const trialJobDetail: TrialDetail = new TrialDetail(trialId, "WAITING", Date.now(), "", form);
168
169
170
171
172
173
174
175
176
177
178
179
180

        this.trials.set(trialId, trialJobDetail);

        return trialJobDetail;
    }

    // to support multi phase
    public async updateTrialJob(trialJobId: string, form: TrialJobApplicationForm): Promise<TrialDetail> {
        const trialDetail = await this.getTrialJob(trialJobId);
        const environment = trialDetail.environment;
        if (environment === undefined) {
            throw new Error(`TrialDispatcher: trial ${trialJobId}'s env shouldn't be undefined in updateTrialJob.`);
        }
181
182
        if (environment.environmentService === undefined) {
            throw new Error(`Environment ${environment.id} does not assigned environment service.`);
183
184
185
186
187
188
        }

        const message = {
            "trialId": trialJobId,
            "parameters": form.hyperParameters,
        }
189
        await environment.environmentService.getCommandChannel.sendCommand(environment, SEND_TRIAL_JOB_PARAMETER, message);
190
191
192
193
194
195
196
197
198
199
200
201

        return trialDetail;
    }

    public async cancelTrialJob(trialJobId: string, isEarlyStopped?: boolean | undefined): Promise<void> {
        const trial = await this.getTrialJob(trialJobId);
        switch (trial.status) {
            case "RUNNING":
            case "WAITING":
            case "UNKNOWN":
                {
                    const environment = trial.environment;
202
203
                    if (environment && environment.environmentService) {
                        await environment.environmentService.getCommandChannel.sendCommand(environment, KILL_TRIAL_JOB, trial.id);
204
205
206
207
208
209
210
211
212
213
214
                        trial.isEarlyStopped = isEarlyStopped;
                        trial.status = trial.isEarlyStopped === true ?
                            'EARLY_STOPPED' : 'USER_CANCELED';
                        this.releaseEnvironment(trial);
                    }
                }
                break;
        }
    }

    public async run(): Promise<void> {
liuzhe-lz's avatar
liuzhe-lz committed
215
        await Promise.all(this.environmentServiceList.map(env => env.init()));
216
217
218
        for(const environmentService of this.environmentServiceList) {
            
            const runnerSettings: RunnerSettings = new RunnerSettings();
219
            runnerSettings.nniManagerIP = this.config.nniManagerIp === undefined? getIPV4Address() : this.config.nniManagerIp;
220
221
222
            runnerSettings.nniManagerPort = getBasePort() + 1;
            runnerSettings.commandChannel = environmentService.getCommandChannel.channelName;
            runnerSettings.enableGpuCollector = this.enableGpuScheduler;
223
            runnerSettings.command = this.config.trialCommand;
224
225
226
227
228
229
230
231
232
233
            runnerSettings.nniManagerVersion = this.enableVersionCheck ? await getVersion() : '';
            runnerSettings.logCollection = this.logCollection;
            runnerSettings.platform = environmentService.getName;
            runnerSettings.experimentId = this.experimentId;

            await environmentService.getCommandChannel.start();
            this.log.info(`TrialDispatcher: started channel: ${environmentService.getCommandChannel.constructor.name}`);
    
            this.log.info(`TrialDispatcher: copying code and settings.`);
            let storageService: StorageService;
234
235
236
237
238
239
240
241
            if (this.useSharedStorage) {
                if (this.fileCopyCompleted) {
                    this.log.debug(`TrialDispatcher: file already copy to shared storage.`);
                    continue;
                }
                this.log.debug(`TrialDispatcher: use shared storage service.`);
                storageService = component.get<SharedStorageService>(SharedStorageService).storageService;
            } else if (environmentService.hasStorageService) {
242
243
244
245
246
                this.log.debug(`TrialDispatcher: use existing storage service.`);
                storageService = component.get<StorageService>(StorageService);
            } else {
                this.log.debug(`TrialDispatcher: create temp storage service to temp folder.`);
                storageService = new MountedStorageService();
SparkSnail's avatar
SparkSnail committed
247
                const environmentLocalTempFolder = path.join(this.experimentRootDir, "environment-temp");
248
                storageService.initialize(this.config.trialCodeDirectory, environmentLocalTempFolder);
249
250
            }
            // Copy the compressed file to remoteDirectory and delete it
251
            const codeDir = path.resolve(this.config.trialCodeDirectory);
252
253
254
255
            const envDir = storageService.joinPath("envs");
            const codeFileName = await storageService.copyDirectory(codeDir, envDir, true);
            storageService.rename(codeFileName, "nni-code.tar.gz");

Ni Hao's avatar
Ni Hao committed
256
257
            const installFileName = storageService.joinPath(envDir, `install_nni.sh`);
            const installFileNameForWin = storageService.joinPath(envDir, `install_nni.ps1`);
258
            await storageService.save(CONTAINER_INSTALL_NNI_SHELL_FORMAT, installFileName);
Ni Hao's avatar
Ni Hao committed
259
            await storageService.save(CONTAINER_INSTALL_NNI_SHELL_FORMAT_FOR_WIN, installFileNameForWin);
260
261
262
263
264
265
266
267
268
269
270

            const runnerSettingsConfig = storageService.joinPath(envDir, "settings.json");
            await storageService.save(JSON.stringify(runnerSettings), runnerSettingsConfig);

            if (this.isDeveloping) {
                let trialToolsPath = path.join(__dirname, "../../../../../tools/nni_trial_tool");
                if (false === fs.existsSync(trialToolsPath)) {
                    trialToolsPath = path.join(__dirname, "..\\..\\..\\..\\..\\tools\\nni_trial_tool");
                }
                await storageService.copyDirectory(trialToolsPath, envDir, true);
            }
271
272
273
274

            if (this.useSharedStorage) {
                this.fileCopyCompleted = true;
            }
275
        }
276
277
278
279
280
281
        // start channel
        this.commandEmitter.on("command", (command: Command): void => {
            this.handleCommand(command).catch((err: Error) => {
                this.log.error(`TrialDispatcher: error on handle env ${command.environment.id} command: ${command.command}, data: ${command.data}, error: ${err}`);
            })
        });
282
        await this.prefetchEnvironments();
283
        this.log.info(`TrialDispatcher: run loop started.`);
284
285
286
287
288
289
290
        const promiseList: Promise<void>[] = [];
        for(const commandChannel of this.commandChannelSet) {
            promiseList.push(commandChannel.run());
        }
        promiseList.push(this.environmentMaintenanceLoop());
        promiseList.push(this.trialManagementLoop());
        await Promise.all(promiseList);
291
292
293
294
295
296
297
298
299
300
    }

    public addTrialJobMetricListener(listener: (metric: TrialJobMetric) => void): void {
        this.metricsEmitter.on('metric', listener);
    }

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

301
302
    public async setClusterMetadata(_key: string, _value: string): Promise<void> { return; }
    public async getClusterMetadata(_key: string): Promise<string> { return ""; }
303
304
305
306
307
308

    public async cleanUp(): Promise<void> {
        if (this.commandEmitter === undefined) {
            throw new Error(`TrialDispatcher: commandEmitter shouldn't be undefined in cleanUp.`);
        }
        this.stopping = true;
309
        this.shouldUpdateTrials = true;
310
311
312
313
        const environments = [...this.environments.values()];

        for (let index = 0; index < environments.length; index++) {
            const environment = environments[index];
314
315
316
            this.log.info(`stopping environment ${environment.id}...`);
            if (environment.environmentService === undefined) {
                throw new Error(`${environment.id} do not have environmentService!`);
317
            }
318
319
            await environment.environmentService.stopEnvironment(environment);
            this.log.info(`stopped environment ${environment.id}.`);
320
321
        }
        this.commandEmitter.off("command", this.handleCommand);
322
323
324
        for (const commandChannel of this.commandChannelSet) {
            await commandChannel.stop();
        }
325
326
327
328
329
        if (this.useSharedStorage) {
            this.log.info(`stopping shared storage...`)
            await component.get<SharedStorageService>(SharedStorageService).cleanUp();
            this.log.info(`shared storage stopped.`)
        }
330
331
332
333
334
335
336
337
338
    }

    private async environmentMaintenanceLoop(): Promise<void> {
        while (!this.stopping) {
            const environments: EnvironmentInformation[] = [];
            for (const environment of this.environments.values()) {
                if (environment.isAlive === true) {
                    environments.push(environment);
                } else {
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
                    if (environment.environmentService === undefined) {
                        throw new Error(`${environment.id} do not have environment service!`);
                    }
                    await environment.environmentService.getCommandChannel.close(environment);
                }
            }
            // Group environments according to environmentService
            const environmentServiceDict: Map<EnvironmentService, EnvironmentInformation[]> =
                new Map<EnvironmentService, EnvironmentInformation[]>();
            for (const environment of environments) {
                if (environment.environmentService === undefined) {
                    throw new Error(`${environment.id} do not have environment service!`);
                }
                if (!environmentServiceDict.has(environment.environmentService)) {
                    environmentServiceDict.set(environment.environmentService, [environment]);
                } else {
                    const environmentsList: EnvironmentInformation[] | undefined = environmentServiceDict.get(environment.environmentService);
                    if (environmentsList === undefined) {
                        throw new Error(`Environment list not initialized!`);
                    }
                    environmentsList.push(environment);
                    environmentServiceDict.set(environment.environmentService, environmentsList);
                }
            }
            // Refresh all environments
            const taskList: Promise<void>[] = [];
            for (const environmentService of environmentServiceDict.keys()) {
                const environmentsList: EnvironmentInformation[] | undefined = environmentServiceDict.get(environmentService);
                if (environmentsList) {
                    taskList.push(environmentService.refreshEnvironmentsStatus(environmentsList));
369
370
                }
            }
371
            await Promise.all(taskList);
372

373
374
375
376
            for (const environment of environments) {
                if (environment.environmentService === undefined) {
                    throw new Error(`${environment.id} do not have environment service!`);
                }
377
378
379
380
381
382
383
384
385
386
387
388
389
390
                const oldIsAlive = environment.isAlive;
                switch (environment.status) {
                    case 'WAITING':
                    case 'RUNNING':
                    case 'UNKNOWN':
                        environment.isAlive = true;
                        break;
                    default:
                        environment.isAlive = false;
                        break;
                }
                if (oldIsAlive !== environment.isAlive) {
                    this.log.debug(`set environment ${environment.id} isAlive from ${oldIsAlive} to ${environment.isAlive} due to status is ${environment.status}.`);
                }
391
            }
392
            this.shouldUpdateTrials = true;
393
394
395
396
            if (this.environmentMaintenceLoopInterval === -1) {
                throw new Error("EnvironmentMaintenceLoopInterval not initialized!");
            }
            await delay(this.environmentMaintenceLoopInterval);
397
398
399
400
        }
    }

    private async trialManagementLoop(): Promise<void> {
401
        const interval = 1;
402
403

        while (!this.stopping) {
404
405
406
407
408
409
410
411
412
            let totalInterval = 1000;
            while (totalInterval > 0) {
                if (this.shouldUpdateTrials) {
                    this.shouldUpdateTrials = false;
                    break;
                }
                totalInterval -= interval;
                await delay(interval);
            }
413
414
415
416
417
418
419
420
421
422
423
424

            const toRefreshedTrials: TrialDetail[] = [];
            for (const trial of this.trials.values()) {
                if (trial.status === "RUNNING" || trial.status === "WAITING" || trial.status === "UNKNOWN") {
                    toRefreshedTrials.push(trial);
                }
            }

            if (toRefreshedTrials.length == 0) {
                continue;
            }

425
            let waitingTrials: TrialDetail[] = [];
426
427
428
429
430
431
432
433
434
435
436
437
438
439
            let liveTrialsCount = 0;
            for (const trial of toRefreshedTrials) {
                const currentStatus = trial.status;
                switch (currentStatus) {
                    case "RUNNING":
                        {
                            const environment = trial.environment;

                            if (environment === undefined) {
                                this.log.error(`found running trial ${trial.id} has no environment, set trial to UNKNOWN.`);
                                trial.status = "UNKNOWN";
                                liveTrialsCount++;
                                continue;
                            }
440
441
442
443

                            if (environment.environmentService === undefined) {
                                throw new Error(`${environment.id} does not has environment service!`);
                            }
SparkSnail's avatar
SparkSnail committed
444
                            trial.url = environment.trackingUrl;
445
446
447
448
449
450
451
452
453
454
455
456
457
                            const environmentStatus = environment.status;

                            // any node exit, then make sure the whole trial stopped.
                            if (trial.nodes.size > 0) {
                                const completedCount = trial.nodes.size;
                                let finalStatus: TrialJobStatus = "SUCCEEDED";
                                let lastTimestamp: number | undefined;
                                this.log.debug(`found ${completedCount} completed trial node(s), nodeCount: ${environment.nodeCount}`);

                                // if some trial processes doesn't exit, kill it for next one.
                                // for example, in horovod, it's just sleep command, has no impact on trial result.
                                if (environment.nodeCount > completedCount) {
                                    this.log.info(`stop partial completed trial ${trial.id}`);
458
                                    await environment.environmentService.getCommandChannel.sendCommand(environment, KILL_TRIAL_JOB, trial.id);
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
                                }
                                for (const node of trial.nodes.values()) {
                                    if (node.status === "FAILED") {
                                        finalStatus = "FAILED";
                                    }
                                    if (node.endTime !== undefined) {
                                        if (lastTimestamp === undefined) {
                                            lastTimestamp = node.endTime
                                        } else {
                                            lastTimestamp = Math.max(node.endTime, lastTimestamp);
                                        }
                                    }
                                }
                                trial.status = finalStatus;
                                if (lastTimestamp === undefined) {
                                    trial.endTime = lastTimestamp;
                                }
                                this.releaseEnvironment(trial);
                            } else if (environmentStatus !== "RUNNING") {
478
                                this.log.error(`found running trial ${trial.id} on '${environment.envId}' with '${environmentStatus}', set trial to environment status.`);
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
                                this.releaseEnvironment(trial);
                                trial.status = environmentStatus;
                            } else {
                                liveTrialsCount++;
                            }
                        }
                        break;
                    case "WAITING":
                    case "UNKNOWN":
                        // deal it later, if there is free environment.
                        waitingTrials.push(trial);
                        liveTrialsCount++;
                        break;
                }
            }
494

495
            let liveEnvironmentsCount = 0;
496
497
            const reusableEnvironments: EnvironmentInformation[] = [];
            for (const environment of this.environments.values()) {
498
499
                if (environment.isAlive === true) {
                    liveEnvironmentsCount++;
500
501
                    if (environment.status === "RUNNING" && environment.isRunnerReady) {
                        // if environment is not reusable and used, stop and not count as idle;
liuzhe-lz's avatar
liuzhe-lz committed
502
                        const reuseMode = Array.isArray(this.config.trainingService) || (this.config.trainingService as any).reuseMode;
503
504
                        if (
                            0 === environment.runningTrialCount &&
liuzhe-lz's avatar
liuzhe-lz committed
505
                            reuseMode === false &&
506
507
                            environment.assignedTrialCount > 0
                        ) {
508
509
510
511
                            if (environment.environmentService === undefined) {
                                throw new Error(`${environment.id} does not has environment service!`);
                            }
                            await environment.environmentService.stopEnvironment(environment);
512
513
514
515
516
517
518
519
520
                            continue;
                        }

                        // if gpu scheduler is not enabled, and there is running trial, skip it.
                        if (false === this.enableGpuScheduler && environment.runningTrialCount > 0) {
                            continue;
                        }

                        reusableEnvironments.push(environment);
521
522
                    }
                }
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
            }

            let neededEnvironmentCount = 0;
            if (true === this.enableGpuScheduler) {
                let noGpuAvailable: boolean = false;
                while (waitingTrials.length > 0) {
                    // skip following trials, if first trial doesn't find available GPU.
                    if (true === noGpuAvailable) {
                        // break loop to try next time.
                        break;
                    }
                    const trial = waitingTrials.shift();
                    if (undefined === trial) {
                        throw new Error(`TrialDispatcher: waiting trial shouldn't be undefined!`);
                    }
538
                    const gpuNum = this.config.trialGpuNumber;
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
                    const result = this.gpuScheduler.scheduleMachine(reusableEnvironments, gpuNum, trial);
                    switch (result.resultType) {
                        case ScheduleResultType.REQUIRE_EXCEED_TOTAL:
                            {
                                if (liveEnvironmentsCount == 0) {
                                    this.log.debug(`TrialDispatcher: no live environment, so request one.`);
                                    neededEnvironmentCount = 1;
                                    waitingTrials = [];
                                    this.isLoggedNoGpuAvailable = false;
                                } else if (reusableEnvironments.length > 0) {
                                    const errorMessage: string = `TrialDispatcher: REQUIRE_EXCEED_TOTAL Required GPU number ${gpuNum} is too large, no machine can meet`;
                                    this.log.error(errorMessage);
                                    throw new NNIError(NNIErrorNames.RESOURCE_NOT_AVAILABLE, errorMessage);
                                } else {
                                    if (false === this.isLoggedNoGpuAvailable) {
                                        this.log.debug(`TrialDispatcher: wait GPU, live environment ${liveEnvironmentsCount}, no reusable, REQUIRE_EXCEED_TOTAL.`)
                                        this.isLoggedNoGpuAvailable = true;
                                    }
                                }
                                break;
                            }
                        case ScheduleResultType.TMP_NO_AVAILABLE_GPU:
                            {
                                if (false === this.isLoggedNoGpuAvailable) {
                                    this.log.debug(`TrialDispatcher: wait GPU, live environment ${liveEnvironmentsCount}, reusable ${reusableEnvironments.length}, TMP_NO_AVAILABLE_GPU.`)
                                    this.isLoggedNoGpuAvailable = true;
                                }

                                // if some environment is alive, but not ready, no need to create more.
                                if (liveEnvironmentsCount <= reusableEnvironments.length) {
                                    neededEnvironmentCount = 1;
                                    this.isLoggedNoGpuAvailable = false;
                                    this.log.info(`TrialDispatcher: ${liveEnvironmentsCount} live env, and ${reusableEnvironments.length} reusable, but no GPU available so request a new one.`);
                                }
                                noGpuAvailable = true;
                            }
                            break
                        case ScheduleResultType.SUCCEED:
                            {
                                const environment = result.environment;
                                if (undefined === environment) {
                                    throw new Error(`TrialDispatcher: scheduled env shouldn't be undefined!`);
                                }
                                trial.assignedGpus = result.gpuIndices;
                                await this.allocateEnvironment(trial, environment);
                                this.isLoggedNoGpuAvailable = false;
                            }
                            break
                        default:
                            throw new Error(`TrialDispatcher: Unknown gpu schecduler type: ${result.resultType}`);
                    }
                }
            } else {
                while (reusableEnvironments.length > 0 && waitingTrials.length > 0) {
                    const trial = waitingTrials.shift();
                    const idleEnvironment = reusableEnvironments.shift();
                    if (trial !== undefined && idleEnvironment != undefined) {
                        await this.allocateEnvironment(trial, idleEnvironment);
                    }
598
                }
599
                neededEnvironmentCount = liveTrialsCount - liveEnvironmentsCount;
600
601
            }

602
603
            if (neededEnvironmentCount > 0) {
                let requestedCount = 0;
604
                let hasMoreEnvironments = false;
605
                for (let index = 0; index < neededEnvironmentCount; index++) {
606
607
608
609
                    const environmentService: EnvironmentService | undefined = this.selectEnvironmentService();
                    if (environmentService !== undefined) {
                        hasMoreEnvironments = true;
                        await this.requestEnvironment(environmentService);
610
611
612
613
614
615
616
617
618
                        requestedCount++;
                        this.isLoggedNoMoreEnvironment = false;
                    } else {
                        if (this.isLoggedNoMoreEnvironment === false) {
                            this.isLoggedNoMoreEnvironment = true;
                            this.log.info(`no more environment so far, so skip to request environment.`)
                        }
                    }
                }
619
                if (hasMoreEnvironments === true || requestedCount > 0) {
620
621
622
                    this.log.info(`requested new environment, live trials: ${liveTrialsCount}, ` +
                        `live environments: ${liveEnvironmentsCount}, neededEnvironmentCount: ${neededEnvironmentCount}, ` +
                        `requestedCount: ${requestedCount}`);
623
624
                }
            }
625

626
627
        }
    }
628

629
630
631
632
633
634
635
636
637
638
    // Schedule a environment platform for environment
    private selectEnvironmentService(): EnvironmentService | undefined {
        const validEnvironmentServiceList = [];
        for(const environmentService of this.environmentServiceList){
            if (environmentService.hasMoreEnvironments) {
                validEnvironmentServiceList.push(environmentService);
            }
        }
        if (validEnvironmentServiceList.length === 0) {
            return undefined;
639
        }
640
641
        // Random scheduler
        return randomSelect(validEnvironmentServiceList);
642
    }
643

644
645
646
647
648
649
650
    private async prefetchEnvironments (): Promise<void> {
        for (const environmentService of this.environmentServiceList) {
            const number = environmentService.prefetchedEnvironmentCount;
            this.log.info(`Initialize environments total number: ${number}`);
            for (let index = 0; index < number; index++) {
                await this.requestEnvironment(environmentService);
            }
651
        }
652
    }
653

654
    private async requestEnvironment(environmentService: EnvironmentService): Promise<void> {
655
656
        const envId = uniqueString(5);
        const envName = `nni_exp_${this.experimentId}_env_${envId}`;
J-shang's avatar
J-shang committed
657
        const environment = environmentService.createEnvironmentInformation(envId, envName);
658
659
        environment.environmentService = environmentService;
        this.log.info(`Assign environment service ${environmentService.getName} to environment ${envId}`);
660
        environment.command = `sh ../install_nni.sh && python3 -m nni.tools.trial_tool.trial_runner`;
661
662
663
664
665

        if (this.isDeveloping) {
            environment.command = "[ -d \"nni_trial_tool\" ] && echo \"nni_trial_tool exists already\" || (mkdir ./nni_trial_tool && tar -xof ../nni_trial_tool.tar.gz -C ./nni_trial_tool) && pip3 install websockets && " + environment.command;
        }

SparkSnail's avatar
SparkSnail committed
666
        environment.command = `mkdir -p envs/${envId} && cd envs/${envId} && ${environment.command}`;
667

668
669
        environment.useSharedStorage = this.useSharedStorage;

670
        await environmentService.startEnvironment(environment);
SparkSnail's avatar
SparkSnail committed
671
        this.environments.set(environment.id, environment);
672
673
674

        if (environment.status === "FAILED") {
            environment.isAlive = false;
675
            throw new Error(`error on request environment ${environment.envId}, please check log for more details.`);
676
677
678
        } else {
            environment.isAlive = true;
        }
679
        await environment.environmentService.getCommandChannel.open(environment);
680
        this.log.info(`requested environment ${environment.id} and job id is ${environment.envId}.`);
681
682
    }

683
    private async allocateEnvironment(trial: TrialDetail, environment: EnvironmentInformation): Promise<void> {
684
        if (trial.environment) {
685
            throw new Error(`TrialDispatcher: trial ${trial.id} has assigned environment ${trial.environment.id} already, not assign to ${environment.id}!`);
686
        }
687
688
        if (environment.runningTrialCount > 0 && false === this.enableGpuScheduler) {
            throw new Error(`TrialDispatcher: environment ${environment.id} has running trial, and gpu scheduler is not enabled, it cannot be assigned again!`);
689
690
691
        }
        this.log.info(`assigning environment ${environment.id} to trial ${trial.id}.`);

692
693
694
        // convert assigned gpus to string for nvidia visible settings
        // undefined means no constraint, [] means no gpu visible.
        let gpuIndices: string | undefined = undefined;
695
        if (undefined !== this.config.trialGpuNumber) {
696
697
698
699
700
701
702
703
704
705
706
            const gpuArray: number[] = [];
            if (undefined !== trial.assignedGpus) {
                trial.assignedGpus.map((value) => {
                    gpuArray.push(value.index);
                });
            }
            gpuIndices = gpuArray.join(',');
        }

        environment.runningTrialCount++;
        environment.assignedTrialCount++;
707
        trial.environment = environment;
708
709
710
        if (environment.environmentService === undefined) {
            throw new Error(`${environment.id} environmentService not initialized!`);
        }
711
        trial.message = `Platform: ${environment.environmentService.getName}, environment: ${environment.id}`;
712
713
714
715
        if (environment.environmentService.hasStorageService) {	
            const storageService = component.get<StorageService>(StorageService);	
            trial.workingDirectory = storageService.joinPath('trials', trial.id);
        }	
716
717
        trial.settings = {
            trialId: trial.id,
718
            gpuIndices: gpuIndices,
719
720
721
722
723
            sequenceId: trial.form.sequenceId,
            parameter: trial.form.hyperParameters,
        }
        trial.startTime = Date.now();
        trial.status = "RUNNING";
724
725
726
727
        if (environment.environmentService === undefined) {
            throw new Error(`${environment.id} does not have environment service!`);
        }
        await environment.environmentService.getCommandChannel.sendCommand(trial.environment, NEW_TRIAL_JOB, trial.settings);
728
    }
SparkSnail's avatar
SparkSnail committed
729
730
731
732
733
    
    /**
     * release the trial assigned environment resources
     * @param trial 
     */
734
    private releaseEnvironment(trial: TrialDetail): void {
SparkSnail's avatar
SparkSnail committed
735
736
737
738
739
        if (trial.environment !== undefined) {
            if (trial.environment.runningTrialCount <= 0) {
                throw new Error(`TrialDispatcher: environment ${trial.environment.id} has no counted running trial!`);
            }
            trial.environment.runningTrialCount--;
740
            trial.environment.latestTrialReleasedTime = Date.now();
SparkSnail's avatar
SparkSnail committed
741
            trial.environment = undefined;
742
        }
743
744
        if (true === this.enableGpuScheduler) {
            this.gpuScheduler.removeGpuReservation(trial);
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
        }
    }

    private async handleMetricData(trialId: string, data: any): Promise<void> {
        if (Array.isArray(data)) {
            for (const subItem of data) {
                this.metricsEmitter.emit('metric', {
                    id: trialId,
                    data: subItem
                });
            }
        } else {
            this.metricsEmitter.emit('metric', {
                id: trialId,
                data: data
            });
        }
    }

    private async handleStdout(commandData: any): Promise<void> {
765
        const metricPattern: RegExp = /NNISDK_MEb'(?<metrics>.*a?)'$/gm;
766
767
768
769
770
771
        const trialLogDir: string = path.join(getExperimentRootDir(), 'trials', commandData["trial"]);
        mkDirPSync(trialLogDir);
        const trialLogPath: string = path.join(trialLogDir, 'stdout_log_collection.log');
        try {
            let skipLogging: boolean = false;
            if (commandData["tag"] === 'trial' && commandData["msg"] !== undefined) {
772
773
774
                const message: string = commandData["msg"];
                let metricsContent = metricPattern.exec(message);
                while (metricsContent && metricsContent.groups) {
775
776
                    const key: string = 'metrics';
                    const data = metricsContent.groups[key];
777
778
                    await this.handleMetricData(commandData["trial"], data);
                    metricsContent = metricPattern.exec(message);
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
                    skipLogging = true;
                }
            }

            if (!skipLogging) {
                // Construct write stream to write remote trial's log into local file
                const writeStream: Writable = fs.createWriteStream(trialLogPath, {
                    flags: 'a+',
                    encoding: 'utf8',
                    autoClose: true
                });

                writeStream.write(String.Format('{0}\n', commandData["msg"]));
                writeStream.end();
            }
        } catch (err) {
            this.log.error(`TrialDispatcher: handleStdout error: ${err}`);
        }
    }

    private async handleCommand(command: Command): Promise<void> {
800
        this.log.debug(`TrialDispatcher: env ${command.environment.id} received command ${command.command}.`);
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
        const environment = command.environment;
        const data = command.data;
        const nodeId = data["node"];
        switch (command.command) {
            case REPORT_METRIC_DATA:
                this.log.error(`TrialDispatcher: TODO: not implement to handle direct REPORT_METRIC_DATA command yet.`);
                break;
            case STDOUT:
                await this.handleStdout(data);
                break;
            case INITIALIZED:
                {
                    let isAllReady = true;
                    if (environment.nodeCount > 1) {
                        let node = environment.nodes.get(nodeId);
                        if (node === undefined) {
J-shang's avatar
J-shang committed
817
                            node = new NodeInformation(nodeId);
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
                            environment.nodes.set(nodeId, node);
                        }
                        const oldNodeStatus = node.status;
                        if (oldNodeStatus === "UNKNOWN" || oldNodeStatus === "WAITING") {
                            node.status = "RUNNING";
                        }

                        if (environment.nodes.size === environment.nodeCount) {
                            for (const node of environment.nodes.values()) {
                                if (node.status !== "RUNNING") {
                                    isAllReady = false;
                                    break;
                                }
                            }
                        } else {
                            isAllReady = false;
                        }
                    }

                    // single node is always ready to set env status
838
839
840
                    if (isAllReady) {
                        environment.isRunnerReady = true;
                        this.log.info(`TrialDispatcher: env ${environment.id} received initialized message and runner is ready, env status: ${environment.status}.`);
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
                    }
                }
                break;
            case VERSION_CHECK:
                {
                    if (this.enableVersionCheck) {
                        const checkResultSuccess: boolean = data["tag"] === 'VCSuccess' ? true : false;
                        if (checkResultSuccess) {
                            this.log.info(`TrialDispatcher: Version check in trialKeeper success!`);
                        } else {
                            const errorMessage = `TrialDispatcher: Version check error, ${data["msg"]}!`;
                            this.log.error(errorMessage);
                        }
                    }
                }
                break;
            case GPU_INFO:
858
859
860
861
                {
                    const gpuData = <TrialGpuSummary>(data);
                    environment.setGpuSummary(nodeId, gpuData);
                }
862
863
864
865
866
867
868
869
870
871
872
873
874
875
                break;
            case TRIAL_END:
                {
                    const trialId = data["trial"];
                    const trial = await this.getTrialJob(trialId);
                    const code = parseInt(data["code"]);
                    const timestamp = parseInt(data["time"]);
                    let exitStatus: TrialJobStatus = "SUCCEEDED";
                    if (code !== 0) {
                        exitStatus = "FAILED";
                    }

                    let node = environment.nodes.get(nodeId);
                    if (node === undefined) {
J-shang's avatar
J-shang committed
876
                        node = new NodeInformation(nodeId);
877
878
879
880
881
882
883
884
885
886
                        trial.nodes.set(nodeId, node);
                    }
                    if (undefined === node) {
                        throw new Error("node is impossible to be undefined (see above code), but make eslint happy!");
                    }
                    node.status = exitStatus;
                    node.endTime = timestamp;
                }
                break;
        }
887
        this.shouldUpdateTrials = true;
888
    }
889

890
891
    private async initializeSharedStorage(config: SharedStorageConfig): Promise<void> {
        switch (config.storageType) {
892
893
894
895
896
897
898
899
900
901
902
            case 'NFS':
                Container.bind(SharedStorageService)
                         .to(NFSSharedStorageService)
                         .scope(Scope.Singleton);
                break;
            case 'AzureBlob':
                Container.bind(SharedStorageService)
                         .to(AzureBlobSharedStorageService)
                         .scope(Scope.Singleton);
                break;
            default: {
903
                const errorMessage = `Shared storage type '${config.storageType}' not support.`;
904
905
906
907
                this.log.error(errorMessage)
                return Promise.reject(errorMessage);
            }
        }
908
        await component.get<SharedStorageService>(SharedStorageService).config(config);
909
910
911
        this.useSharedStorage = true;
        return Promise.resolve();
    }
J-shang's avatar
J-shang committed
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934

    public async getTrialOutputLocalPath(trialJobId: string): Promise<string> {
        // TODO: support non shared storage
        if (this.useSharedStorage) {
            const localWorkingRoot = component.get<SharedStorageService>(SharedStorageService).localWorkingRoot;
            return Promise.resolve(path.join(localWorkingRoot, 'trials', trialJobId));
        } else {
            return Promise.reject(new Error('Only support shared storage right now.'));
        }
    }

    public async fetchTrialOutput(trialJobId: string, subpath: string | undefined): Promise<void> {
        // TODO: support non shared storage
        let trialLocalPath = await this.getTrialOutputLocalPath(trialJobId);
        if (subpath !== undefined) {
            trialLocalPath = path.join(trialLocalPath, subpath);
        }
        if (fs.existsSync(trialLocalPath)) {
            return Promise.resolve();
        } else {
            return Promise.reject(new Error('Trial local path not exist.'));
        }
    }
935
936
937
}

export { TrialDispatcher };