trialDispatcher.ts 44.2 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 { EnvironmentServiceFactory } 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
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
    // 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
85
    private config: ExperimentConfig;

    constructor(config: ExperimentConfig) {
86
87
88
89
90
        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
126
127
128
129
130
131
132
133
134
135
136
137

        this.config = config;

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

        validateCodeDir(config.trialCodeDirectory);

        if (Array.isArray(config.trainingService)) {
            config.trainingService.forEach(trainingService => {
                const env = EnvironmentServiceFactory.createEnvironmentService(trainingService.platform, config);
                this.environmentServiceList.push(env);
            });
        } else {
            const env = EnvironmentServiceFactory.createEnvironmentService(config.trainingService.platform, config);
            this.environmentServiceList.push(env);
        }

        // FIXME: max?
        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) {
            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
    }

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

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

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

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

        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.`);
        }
180
181
        if (environment.environmentService === undefined) {
            throw new Error(`Environment ${environment.id} does not assigned environment service.`);
182
183
184
185
186
187
        }

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

        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;
201
202
                    if (environment && environment.environmentService) {
                        await environment.environmentService.getCommandChannel.sendCommand(environment, KILL_TRIAL_JOB, trial.id);
203
204
205
206
207
208
209
210
211
212
213
                        trial.isEarlyStopped = isEarlyStopped;
                        trial.status = trial.isEarlyStopped === true ?
                            'EARLY_STOPPED' : 'USER_CANCELED';
                        this.releaseEnvironment(trial);
                    }
                }
                break;
        }
    }

    public async run(): Promise<void> {
214
215
216
        for(const environmentService of this.environmentServiceList) {
            
            const runnerSettings: RunnerSettings = new RunnerSettings();
217
            runnerSettings.nniManagerIP = this.config.nniManagerIp === undefined? getIPV4Address() : this.config.nniManagerIp;
218
219
220
            runnerSettings.nniManagerPort = getBasePort() + 1;
            runnerSettings.commandChannel = environmentService.getCommandChannel.channelName;
            runnerSettings.enableGpuCollector = this.enableGpuScheduler;
221
            runnerSettings.command = this.config.trialCommand;
222
223
224
225
226
227
228
229
230
231
            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;
232
233
234
235
236
237
238
239
            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) {
240
241
242
243
244
                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
245
                const environmentLocalTempFolder = path.join(this.experimentRootDir, "environment-temp");
246
                storageService.initialize(this.config.trialCodeDirectory, environmentLocalTempFolder);
247
248
            }
            // Copy the compressed file to remoteDirectory and delete it
249
            const codeDir = path.resolve(this.config.trialCodeDirectory);
250
251
252
253
            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
254
255
            const installFileName = storageService.joinPath(envDir, `install_nni.sh`);
            const installFileNameForWin = storageService.joinPath(envDir, `install_nni.ps1`);
256
            await storageService.save(CONTAINER_INSTALL_NNI_SHELL_FORMAT, installFileName);
Ni Hao's avatar
Ni Hao committed
257
            await storageService.save(CONTAINER_INSTALL_NNI_SHELL_FORMAT_FOR_WIN, installFileNameForWin);
258
259
260
261
262
263
264
265
266
267
268

            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);
            }
269
270
271
272

            if (this.useSharedStorage) {
                this.fileCopyCompleted = true;
            }
273
        }
274
275
276
277
278
279
        // 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}`);
            })
        });
280
        await this.prefetchEnvironments();
281
        this.log.info(`TrialDispatcher: run loop started.`);
282
283
284
285
286
287
288
        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);
289
290
291
292
293
294
295
296
297
298
    }

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

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

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

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

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

    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 {
337
338
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
                    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));
367
368
                }
            }
369
            await Promise.all(taskList);
370

371
372
373
374
            for (const environment of environments) {
                if (environment.environmentService === undefined) {
                    throw new Error(`${environment.id} do not have environment service!`);
                }
375
376
377
378
379
380
381
382
383
384
385
386
387
388
                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}.`);
                }
389
            }
390
            this.shouldUpdateTrials = true;
391
392
393
394
            if (this.environmentMaintenceLoopInterval === -1) {
                throw new Error("EnvironmentMaintenceLoopInterval not initialized!");
            }
            await delay(this.environmentMaintenceLoopInterval);
395
396
397
398
        }
    }

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

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

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

423
            let waitingTrials: TrialDetail[] = [];
424
425
426
427
428
429
430
431
432
433
434
435
436
437
            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;
                            }
438
439
440
441

                            if (environment.environmentService === undefined) {
                                throw new Error(`${environment.id} does not has environment service!`);
                            }
SparkSnail's avatar
SparkSnail committed
442
                            trial.url = environment.trackingUrl;
443
444
445
446
447
448
449
450
451
452
453
454
455
                            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}`);
456
                                    await environment.environmentService.getCommandChannel.sendCommand(environment, KILL_TRIAL_JOB, trial.id);
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
                                }
                                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") {
476
                                this.log.error(`found running trial ${trial.id} on '${environment.envId}' with '${environmentStatus}', set trial to environment status.`);
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
                                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;
                }
            }
492

493
            let liveEnvironmentsCount = 0;
494
495
            const reusableEnvironments: EnvironmentInformation[] = [];
            for (const environment of this.environments.values()) {
496
497
                if (environment.isAlive === true) {
                    liveEnvironmentsCount++;
498
499
500
501
                    if (environment.status === "RUNNING" && environment.isRunnerReady) {
                        // if environment is not reusable and used, stop and not count as idle;
                        if (
                            0 === environment.runningTrialCount &&
502
                            !(this.config as any).reuseMode &&
503
504
                            environment.assignedTrialCount > 0
                        ) {
505
506
507
508
                            if (environment.environmentService === undefined) {
                                throw new Error(`${environment.id} does not has environment service!`);
                            }
                            await environment.environmentService.stopEnvironment(environment);
509
510
511
512
513
514
515
516
517
                            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);
518
519
                    }
                }
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
            }

            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!`);
                    }
535
                    const gpuNum = this.config.trialGpuNumber;
536
537
538
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
                    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);
                    }
595
                }
596
                neededEnvironmentCount = liveTrialsCount - liveEnvironmentsCount;
597
598
            }

599
600
            if (neededEnvironmentCount > 0) {
                let requestedCount = 0;
601
                let hasMoreEnvironments = false;
602
                for (let index = 0; index < neededEnvironmentCount; index++) {
603
604
605
606
                    const environmentService: EnvironmentService | undefined = this.selectEnvironmentService();
                    if (environmentService !== undefined) {
                        hasMoreEnvironments = true;
                        await this.requestEnvironment(environmentService);
607
608
609
610
611
612
613
614
615
                        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.`)
                        }
                    }
                }
616
                if (hasMoreEnvironments === true || requestedCount > 0) {
617
618
619
                    this.log.info(`requested new environment, live trials: ${liveTrialsCount}, ` +
                        `live environments: ${liveEnvironmentsCount}, neededEnvironmentCount: ${neededEnvironmentCount}, ` +
                        `requestedCount: ${requestedCount}`);
620
621
                }
            }
622

623
624
        }
    }
625

626
627
628
629
630
631
632
633
634
635
    // 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;
636
        }
637
638
        // Random scheduler
        return randomSelect(validEnvironmentServiceList);
639
    }
640

641
642
643
644
645
646
647
    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);
            }
648
        }
649
    }
650

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

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

665
666
        environment.useSharedStorage = this.useSharedStorage;

667
        await environmentService.startEnvironment(environment);
SparkSnail's avatar
SparkSnail committed
668
        this.environments.set(environment.id, environment);
669
670
671

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

680
    private async allocateEnvironment(trial: TrialDetail, environment: EnvironmentInformation): Promise<void> {
681
        if (trial.environment) {
682
            throw new Error(`TrialDispatcher: trial ${trial.id} has assigned environment ${trial.environment.id} already, not assign to ${environment.id}!`);
683
        }
684
685
        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!`);
686
687
688
        }
        this.log.info(`assigning environment ${environment.id} to trial ${trial.id}.`);

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

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

    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> {
762
        const metricPattern: RegExp = /NNISDK_MEb'(?<metrics>.*a?)'$/gm;
763
764
765
766
767
768
        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) {
769
770
771
                const message: string = commandData["msg"];
                let metricsContent = metricPattern.exec(message);
                while (metricsContent && metricsContent.groups) {
772
773
                    const key: string = 'metrics';
                    const data = metricsContent.groups[key];
774
775
                    await this.handleMetricData(commandData["trial"], data);
                    metricsContent = metricPattern.exec(message);
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
                    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> {
797
        this.log.debug(`TrialDispatcher: env ${command.environment.id} received command ${command.command}.`);
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
        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
814
                            node = new NodeInformation(nodeId);
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
                            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
835
836
837
                    if (isAllReady) {
                        environment.isRunnerReady = true;
                        this.log.info(`TrialDispatcher: env ${environment.id} received initialized message and runner is ready, env status: ${environment.status}.`);
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
                    }
                }
                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:
855
856
857
858
                {
                    const gpuData = <TrialGpuSummary>(data);
                    environment.setGpuSummary(nodeId, gpuData);
                }
859
860
861
862
863
864
865
866
867
868
869
870
871
872
                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
873
                        node = new NodeInformation(nodeId);
874
875
876
877
878
879
880
881
882
883
                        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;
        }
884
        this.shouldUpdateTrials = true;
885
    }
886

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

    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.'));
        }
    }
932
933
934
}

export { TrialDispatcher };