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 { NNIManagerIpConfig, TrainingService, TrialJobApplicationForm, TrialJobMetric, TrialJobStatus, LogType } from '../../common/trainingService';
17
import { delay, getExperimentRootDir, getIPV4Address, getLogLevel, getVersion, mkDirPSync, randomSelect, uniqueString } from '../../common/utils';
18
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';
19
import { ScheduleResultType } from '../../training_service/common/gpuData';
20
21
22
23
24
import { CONTAINER_INSTALL_NNI_SHELL_FORMAT } from '../common/containerJobData';
import { TrialConfig } from '../common/trialConfig';
import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
import { validateCodeDir } from '../common/util';
import { Command, CommandChannel } from './commandChannel';
J-shang's avatar
J-shang committed
25
import { EnvironmentInformation, EnvironmentService, NodeInformation, RunnerSettings, TrialGpuSummary } from './environment';
26
import { EnvironmentServiceFactory } from './environments/environmentServiceFactory';
27
import { GpuScheduler } from './gpuScheduler';
SparkSnail's avatar
SparkSnail committed
28
import { MountedStorageService } from './storages/mountedStorageService';
29
import { StorageService } from './storageService';
30
31
32
import { SharedStorageService, SharedStorageConfig } from './sharedStorage';
import { NFSSharedStorageService } from './shared_storages/nfsStorageService'
import { AzureBlobSharedStorageService } from './shared_storages/azureblobStorageService'
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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 {
    private readonly log: Logger;
    private readonly isDeveloping: boolean = false;
    private stopping: boolean = false;

    private readonly metricsEmitter: EventEmitter;
    private readonly experimentId: string;
SparkSnail's avatar
SparkSnail committed
48
    private readonly experimentRootDir: string;
49
50
51
52
53
54
55

    private enableVersionCheck: boolean = true;

    private trialConfig: TrialConfig | undefined;

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

    private nniManagerIp: string | undefined;
63

64
65
66
67
68
69
70
71
72
    // 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;
73
    private logCollection: string = '';
74
75
76
77
78
79
80

    private gpuScheduler: GpuScheduler;

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

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

85
86
87
88
89
90
    constructor() {
        this.log = getLogger();
        this.trials = new Map<string, TrialDetail>();
        this.environments = new Map<string, EnvironmentInformation>();
        this.metricsEmitter = new EventEmitter();
        this.experimentId = getExperimentId();
SparkSnail's avatar
SparkSnail committed
91
        this.experimentRootDir = getExperimentRootDir();
92
        this.commandChannelSet = new Set<CommandChannel>();
93
94
95
96
97
98
99
100

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

102
103
        this.commandEmitter = new EventEmitter();

104
        this.gpuScheduler = new GpuScheduler();
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
    }

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

126
127
128
129
    public async getTrialLog(_trialJobId: string, _logType: LogType): Promise<string> {
        throw new MethodNotImplementedError();
    }

130
131
132
133
134
135
136
    public async submitTrialJob(form: TrialJobApplicationForm): Promise<TrialDetail> {
        if (this.trialConfig === undefined) {
            throw new Error(`trialConfig not initialized!`);
        }

        const trialId: string = uniqueString(5);

137
        const trialJobDetail: TrialDetail = new TrialDetail(trialId, "WAITING", Date.now(), "", form);
138
139
140
141
142
143
144
145
146
147
148
149
150

        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.`);
        }
151
152
        if (environment.environmentService === undefined) {
            throw new Error(`Environment ${environment.id} does not assigned environment service.`);
153
154
155
156
157
158
        }

        const message = {
            "trialId": trialJobId,
            "parameters": form.hyperParameters,
        }
159
        await environment.environmentService.getCommandChannel.sendCommand(environment, SEND_TRIAL_JOB_PARAMETER, message);
160
161
162
163
164
165
166
167
168
169
170
171

        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;
172
173
                    if (environment && environment.environmentService) {
                        await environment.environmentService.getCommandChannel.sendCommand(environment, KILL_TRIAL_JOB, trial.id);
174
175
176
177
178
179
180
181
182
183
184
                        trial.isEarlyStopped = isEarlyStopped;
                        trial.status = trial.isEarlyStopped === true ?
                            'EARLY_STOPPED' : 'USER_CANCELED';
                        this.releaseEnvironment(trial);
                    }
                }
                break;
        }
    }

    public async run(): Promise<void> {
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
        if (this.trialConfig === undefined) {
            throw new Error(`trial config shouldn't be undefined in run()`);
        }
        for(const environmentService of this.environmentServiceList) {
            
            const runnerSettings: RunnerSettings = new RunnerSettings();
            runnerSettings.nniManagerIP = this.nniManagerIp === undefined? getIPV4Address() : this.nniManagerIp;
            runnerSettings.nniManagerPort = getBasePort() + 1;
            runnerSettings.commandChannel = environmentService.getCommandChannel.channelName;
            runnerSettings.enableGpuCollector = this.enableGpuScheduler;
            runnerSettings.command = this.trialConfig.command;
            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;
206
207
208
209
210
211
212
213
            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) {
214
215
216
217
218
                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
219
                const environmentLocalTempFolder = path.join(this.experimentRootDir, "environment-temp");
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
                storageService.initialize(this.trialConfig.codeDir, environmentLocalTempFolder);
            }
            // Copy the compressed file to remoteDirectory and delete it
            const codeDir = path.resolve(this.trialConfig.codeDir);
            const envDir = storageService.joinPath("envs");
            const codeFileName = await storageService.copyDirectory(codeDir, envDir, true);
            storageService.rename(codeFileName, "nni-code.tar.gz");

            const installFileName = storageService.joinPath(envDir, 'install_nni.sh');
            await storageService.save(CONTAINER_INSTALL_NNI_SHELL_FORMAT, installFileName);

            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);
            }
241
242
243
244

            if (this.useSharedStorage) {
                this.fileCopyCompleted = true;
            }
245
        }
246
247
248
249
250
251
        // 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}`);
            })
        });
252
        await this.prefetchEnvironments();
253
        this.log.info(`TrialDispatcher: run loop started.`);
254
255
256
257
258
259
260
        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);
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
    }

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

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

    public get isMultiPhaseJobSupported(): boolean {
        return true;
    }

    public async setClusterMetadata(key: string, value: string): Promise<void> {
        switch (key) {
            case TrialConfigMetadataKey.NNI_MANAGER_IP:
278
                this.nniManagerIp = (<NNIManagerIpConfig>JSON.parse(value)).nniManagerIp;
279
280
281
282
283
                break;
            case TrialConfigMetadataKey.VERSION_CHECK:
                this.enableVersionCheck = (value === 'true' || value === 'True');
                break;
            case TrialConfigMetadataKey.LOG_COLLECTION:
284
                this.logCollection = value;
285
286
287
288
                break;
            case TrialConfigMetadataKey.TRIAL_CONFIG:
                this.trialConfig = <TrialConfig>JSON.parse(value);

289
290
291
292
293
294
295
296
                if (this.trialConfig.reuseEnvironment !== undefined) {
                    this.reuseEnvironment = this.trialConfig.reuseEnvironment;
                }
                if (this.trialConfig.gpuNum !== undefined && this.trialConfig.gpuNum > 0) {
                    this.log.info(`TrialDispatcher: GPU scheduler is enabled.`)
                    this.enableGpuScheduler = true;
                }

297
298
299
                // Validate to make sure codeDir doesn't have too many files
                await validateCodeDir(this.trialConfig.codeDir);
                break;
300
301
302
303
304
305
306
307
308
309
            case TrialConfigMetadataKey.PLATFORM_LIST: {
                const platforms: string[] = value.split(",");
                for(const platform of platforms) {
                    const environmentService: EnvironmentService = EnvironmentServiceFactory.createEnvironmentService(platform);
                    environmentService.initCommandChannel(this.commandEmitter);
                    this.environmentMaintenceLoopInterval =
                      Math.max(environmentService.environmentMaintenceLoopInterval, this.environmentMaintenceLoopInterval);
                    this.commandChannelSet.add(environmentService.getCommandChannel);
                    this.environmentServiceList.push(environmentService);
                }
310
                break;
311
            }
312
313
314
315
316
317
318
319
            case TrialConfigMetadataKey.SHARED_STORAGE_CONFIG:
                if (this.useSharedStorage === false) {
                    await this.initializeSharedStorage(key, value);
                } else {
                    const errorMessage = `Already has set shared storage.`;
                    this.log.error(errorMessage);
                }
                break;
320
321
322
        }
        for(const environmentService of this.environmentServiceList) {
            await environmentService.config(key, value);
323
324
325
326
327
328
329
330
331
332
333
334
        }
    }

    public getClusterMetadata(_key: string): Promise<string> {
        throw new Error('Not implemented!');
    }

    public async cleanUp(): Promise<void> {
        if (this.commandEmitter === undefined) {
            throw new Error(`TrialDispatcher: commandEmitter shouldn't be undefined in cleanUp.`);
        }
        this.stopping = true;
335
        this.shouldUpdateTrials = true;
336
337
338
339
        const environments = [...this.environments.values()];

        for (let index = 0; index < environments.length; index++) {
            const environment = environments[index];
340
341
342
            this.log.info(`stopping environment ${environment.id}...`);
            if (environment.environmentService === undefined) {
                throw new Error(`${environment.id} do not have environmentService!`);
343
            }
344
345
            await environment.environmentService.stopEnvironment(environment);
            this.log.info(`stopped environment ${environment.id}.`);
346
347
        }
        this.commandEmitter.off("command", this.handleCommand);
348
349
350
        for (const commandChannel of this.commandChannelSet) {
            await commandChannel.stop();
        }
351
352
353
354
355
356
357
358
359
    }

    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 {
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
                    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));
390
391
                }
            }
392
            await Promise.all(taskList);
393

394
395
396
397
            for (const environment of environments) {
                if (environment.environmentService === undefined) {
                    throw new Error(`${environment.id} do not have environment service!`);
                }
398
399
400
401
402
403
404
405
406
407
408
409
410
411
                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}.`);
                }
412
            }
413
            this.shouldUpdateTrials = true;
414
415
416
417
            if (this.environmentMaintenceLoopInterval === -1) {
                throw new Error("EnvironmentMaintenceLoopInterval not initialized!");
            }
            await delay(this.environmentMaintenceLoopInterval);
418
419
420
421
        }
    }

    private async trialManagementLoop(): Promise<void> {
422
        const interval = 1;
423
424

        while (!this.stopping) {
425
426
427
428
429
430
431
432
433
            let totalInterval = 1000;
            while (totalInterval > 0) {
                if (this.shouldUpdateTrials) {
                    this.shouldUpdateTrials = false;
                    break;
                }
                totalInterval -= interval;
                await delay(interval);
            }
434
435
436
437
438
439
440
441
442
443
444
445

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

446
            let waitingTrials: TrialDetail[] = [];
447
448
449
450
451
452
453
454
455
456
457
458
459
460
            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;
                            }
461
462
463
464

                            if (environment.environmentService === undefined) {
                                throw new Error(`${environment.id} does not has environment service!`);
                            }
SparkSnail's avatar
SparkSnail committed
465
                            trial.url = environment.trackingUrl;
466
467
468
469
470
471
472
473
474
475
476
477
478
                            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}`);
479
                                    await environment.environmentService.getCommandChannel.sendCommand(environment, KILL_TRIAL_JOB, trial.id);
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
                                }
                                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") {
499
                                this.log.error(`found running trial ${trial.id} on '${environment.envId}' with '${environmentStatus}', set trial to environment status.`);
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
                                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;
                }
            }
515

516
            let liveEnvironmentsCount = 0;
517
518
            const reusableEnvironments: EnvironmentInformation[] = [];
            for (const environment of this.environments.values()) {
519
520
                if (environment.isAlive === true) {
                    liveEnvironmentsCount++;
521
522
523
524
525
526
527
                    if (environment.status === "RUNNING" && environment.isRunnerReady) {
                        // if environment is not reusable and used, stop and not count as idle;
                        if (
                            0 === environment.runningTrialCount &&
                            false === this.reuseEnvironment &&
                            environment.assignedTrialCount > 0
                        ) {
528
529
530
531
                            if (environment.environmentService === undefined) {
                                throw new Error(`${environment.id} does not has environment service!`);
                            }
                            await environment.environmentService.stopEnvironment(environment);
532
533
534
535
536
537
538
539
540
                            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);
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
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
            }

            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!`);
                    }
                    const gpuNum = this.trialConfig ? this.trialConfig.gpuNum : undefined;
                    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);
                    }
618
                }
619
                neededEnvironmentCount = liveTrialsCount - liveEnvironmentsCount;
620
621
            }

622
623
            if (neededEnvironmentCount > 0) {
                let requestedCount = 0;
624
                let hasMoreEnvironments = false;
625
                for (let index = 0; index < neededEnvironmentCount; index++) {
626
627
628
629
                    const environmentService: EnvironmentService | undefined = this.selectEnvironmentService();
                    if (environmentService !== undefined) {
                        hasMoreEnvironments = true;
                        await this.requestEnvironment(environmentService);
630
631
632
633
634
635
636
637
638
                        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.`)
                        }
                    }
                }
639
                if (hasMoreEnvironments === true || requestedCount > 0) {
640
641
642
                    this.log.info(`requested new environment, live trials: ${liveTrialsCount}, ` +
                        `live environments: ${liveEnvironmentsCount}, neededEnvironmentCount: ${neededEnvironmentCount}, ` +
                        `requestedCount: ${requestedCount}`);
643
644
                }
            }
645

646
647
        }
    }
648

649
650
651
652
653
654
655
656
657
658
    // 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;
659
        }
660
661
        // Random scheduler
        return randomSelect(validEnvironmentServiceList);
662
    }
663

664
665
666
667
668
669
670
    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);
            }
671
        }
672
    }
673

674
    private async requestEnvironment(environmentService: EnvironmentService): Promise<void> {
675
676
        const envId = uniqueString(5);
        const envName = `nni_exp_${this.experimentId}_env_${envId}`;
J-shang's avatar
J-shang committed
677
        const environment = environmentService.createEnvironmentInformation(envId, envName);
678
679
        environment.environmentService = environmentService;
        this.log.info(`Assign environment service ${environmentService.getName} to environment ${envId}`);
680
        environment.command = `sh ../install_nni.sh && python3 -m nni.tools.trial_tool.trial_runner`;
681
682
683
684
685

        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
686
        environment.command = `mkdir -p envs/${envId} && cd envs/${envId} && ${environment.command}`;
687

688
689
        environment.useSharedStorage = this.useSharedStorage;

690
        await environmentService.startEnvironment(environment);
SparkSnail's avatar
SparkSnail committed
691
        this.environments.set(environment.id, environment);
692
693
694

        if (environment.status === "FAILED") {
            environment.isAlive = false;
695
            throw new Error(`error on request environment ${environment.envId}, please check log for more details.`);
696
697
698
        } else {
            environment.isAlive = true;
        }
699
        await environment.environmentService.getCommandChannel.open(environment);
700
        this.log.info(`requested environment ${environment.id} and job id is ${environment.envId}.`);
701
702
    }

703
704
705
    private async allocateEnvironment(trial: TrialDetail, environment: EnvironmentInformation): Promise<void> {
        if (this.trialConfig === undefined) {
            throw new Error(`TrialDispatcher: trialConfig shouldn't be undefined in allocateEnvironment.`);
706
707
708
        }

        if (trial.environment) {
709
            throw new Error(`TrialDispatcher: trial ${trial.id} has assigned environment ${trial.environment.id} already, not assign to ${environment.id}!`);
710
        }
711
712
        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!`);
713
714
715
        }
        this.log.info(`assigning environment ${environment.id} to trial ${trial.id}.`);

716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
        // convert assigned gpus to string for nvidia visible settings
        // undefined means no constraint, [] means no gpu visible.
        let gpuIndices: string | undefined = undefined;
        if (undefined !== this.trialConfig.gpuNum) {
            const gpuArray: number[] = [];
            if (undefined !== trial.assignedGpus) {
                trial.assignedGpus.map((value) => {
                    gpuArray.push(value.index);
                });
            }
            gpuIndices = gpuArray.join(',');
        }

        environment.runningTrialCount++;
        environment.assignedTrialCount++;
731
        trial.environment = environment;
732
733
734
        if (environment.environmentService === undefined) {
            throw new Error(`${environment.id} environmentService not initialized!`);
        }
735
        trial.message = `Platform: ${environment.environmentService.getName}, environment: ${environment.id}`;
736
737
738
739
        if (environment.environmentService.hasStorageService) {	
            const storageService = component.get<StorageService>(StorageService);	
            trial.workingDirectory = storageService.joinPath('trials', trial.id);
        }	
740
741
        trial.settings = {
            trialId: trial.id,
742
            gpuIndices: gpuIndices,
743
744
745
746
747
            sequenceId: trial.form.sequenceId,
            parameter: trial.form.hyperParameters,
        }
        trial.startTime = Date.now();
        trial.status = "RUNNING";
748
749
750
751
        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);
752
    }
SparkSnail's avatar
SparkSnail committed
753
754
755
756
757
    
    /**
     * release the trial assigned environment resources
     * @param trial 
     */
758
    private releaseEnvironment(trial: TrialDetail): void {
SparkSnail's avatar
SparkSnail committed
759
760
761
762
763
        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--;
764
            trial.environment.latestTrialReleasedTime = Date.now();
SparkSnail's avatar
SparkSnail committed
765
            trial.environment = undefined;
766
        }
767
768
        if (true === this.enableGpuScheduler) {
            this.gpuScheduler.removeGpuReservation(trial);
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
        }
    }

    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> {
789
        const metricPattern: RegExp = /NNISDK_MEb'(?<metrics>.*a?)'$/gm;
790
791
792
793
794
795
        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) {
796
797
798
                const message: string = commandData["msg"];
                let metricsContent = metricPattern.exec(message);
                while (metricsContent && metricsContent.groups) {
799
800
                    const key: string = 'metrics';
                    const data = metricsContent.groups[key];
801
802
                    await this.handleMetricData(commandData["trial"], data);
                    metricsContent = metricPattern.exec(message);
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
                    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> {
824
        this.log.debug(`TrialDispatcher: env ${command.environment.id} received command ${command.command}.`);
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
        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
841
                            node = new NodeInformation(nodeId);
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
                            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
862
863
864
                    if (isAllReady) {
                        environment.isRunnerReady = true;
                        this.log.info(`TrialDispatcher: env ${environment.id} received initialized message and runner is ready, env status: ${environment.status}.`);
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
                    }
                }
                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:
882
883
884
885
                {
                    const gpuData = <TrialGpuSummary>(data);
                    environment.setGpuSummary(nodeId, gpuData);
                }
886
887
888
889
890
891
892
893
894
895
896
897
898
899
                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
900
                        node = new NodeInformation(nodeId);
901
902
903
904
905
906
907
908
909
910
                        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;
        }
911
        this.shouldUpdateTrials = true;
912
    }
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936

    private async initializeSharedStorage(key: string, value: string): Promise<void> {
        const storageType = (<SharedStorageConfig>JSON.parse(value)).storageType;
        switch (storageType) {
            case 'NFS':
                Container.bind(SharedStorageService)
                         .to(NFSSharedStorageService)
                         .scope(Scope.Singleton);
                break;
            case 'AzureBlob':
                Container.bind(SharedStorageService)
                         .to(AzureBlobSharedStorageService)
                         .scope(Scope.Singleton);
                break;
            default: {
                const errorMessage = `Shared storage type '${storageType}' not support.`;
                this.log.error(errorMessage)
                return Promise.reject(errorMessage);
            }
        }
        await component.get<SharedStorageService>(SharedStorageService).config(key, value);
        this.useSharedStorage = true;
        return Promise.resolve();
    }
937
938
939
}

export { TrialDispatcher };