kubernetesTrainingService.ts 15.5 KB
Newer Older
liuzhe-lz's avatar
liuzhe-lz committed
1
2
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
3

4
'use strict';
5
6
7
8

import * as cpp from 'child-process-promise';
import * as path from 'path';

9
import * as azureStorage from 'azure-storage';
10
import { EventEmitter } from 'events';
11
12
import { Base64 } from 'js-base64';
import { String } from 'typescript-string-operations';
13
import { getExperimentId } from '../../common/experimentStartupInfo';
14
15
import { getLogger, Logger } from '../../common/log';
import {
16
    NNIManagerIpConfig, TrialJobDetail, TrialJobMetric
17
} from '../../common/trainingService';
18
import { delay, getExperimentRootDir, getIPV4Address, getJobCancelStatus, getVersion, uniqueString } from '../../common/utils';
19
import { AzureStorageClientUtility } from './azureStorageClientUtils';
20
21
22
import { GeneralK8sClient, KubernetesCRDClient } from './kubernetesApiClient';
import { KubernetesClusterConfig } from './kubernetesConfig';
import { kubernetesScriptFormat, KubernetesTrialJobDetail } from './kubernetesData';
23
24
import { KubernetesJobRestServer } from './kubernetesJobRestServer';

25
var yaml = require('js-yaml');
26
27
var fs = require('fs');

28
29
30
/**
 * Training Service implementation for Kubernetes
 */
31
32
33
34
35
abstract class KubernetesTrainingService {
    protected readonly NNI_KUBERNETES_TRIAL_LABEL: string = 'nni-kubernetes-trial';
    protected readonly log!: Logger;
    protected readonly metricsEmitter: EventEmitter;
    protected readonly trialJobsMap: Map<string, KubernetesTrialJobDetail>;
36
    //  experiment root dir in NFS
37
38
39
40
41
42
43
44
45
46
47
48
49
50
    protected readonly trialLocalNFSTempFolder: string;
    protected stopping: boolean = false;
    protected experimentId! : string;
    protected kubernetesRestServerPort?: number;
    protected readonly CONTAINER_MOUNT_PATH: string;
    protected azureStorageClient?: azureStorage.FileService;
    protected azureStorageShare?: string;
    protected azureStorageSecretName?: string;
    protected azureStorageAccountName?: string;
    protected nniManagerIpConfig?: NNIManagerIpConfig;
    protected readonly genericK8sClient: GeneralK8sClient;
    protected kubernetesCRDClient?: KubernetesCRDClient;
    protected kubernetesJobRestServer?: KubernetesJobRestServer;
    protected kubernetesClusterConfig?: KubernetesClusterConfig;
51
    protected versionCheck: boolean = true;
SparkSnail's avatar
SparkSnail committed
52
    protected logCollection: string;
53

54
55
56
57
58
    constructor() {
        this.log = getLogger();
        this.metricsEmitter = new EventEmitter();
        this.trialJobsMap = new Map<string, KubernetesTrialJobDetail>();
        this.trialLocalNFSTempFolder = path.join(getExperimentRootDir(), 'trials-nfs-tmp');
59
        this.experimentId = getExperimentId();
60
61
        this.CONTAINER_MOUNT_PATH = '/tmp/mount';
        this.genericK8sClient = new GeneralK8sClient();
SparkSnail's avatar
SparkSnail committed
62
        this.logCollection = 'none';
63
64
    }

65
66
    // tslint:disable:no-any
    public generatePodResource(memory: number, cpuNum: number, gpuNum: number): any {
67
        const resources: any = {
68
            memory: `${memory}Mi`,
69
            cpu: `${cpuNum}`
70
        };
71
72
73
74
75
76

        if (gpuNum !== 0) {
            resources['nvidia.com/gpu'] = `${gpuNum}`;
        }

        return resources;
77
    } // tslint:enable:no-any
78

79
    public async listTrialJobs(): Promise<TrialJobDetail[]> {
80
        const jobs: TrialJobDetail[] = [];
81
82

        for (const [key, value] of this.trialJobsMap) {
83
            jobs.push(await this.getTrialJob(key));
84
        }
85
86
87
88

        return Promise.resolve(jobs);
    }

89
    public async getTrialJob(trialJobId: string): Promise<TrialJobDetail> {
90
91
92

        const kubernetesTrialJob: TrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);

93
94
        if (kubernetesTrialJob === undefined) {
            return Promise.reject(`trial job ${trialJobId} not found`);
95
        }
96
97
98
99

        return Promise.resolve(kubernetesTrialJob);
    }

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

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

108
109
110
111
112
113
114
115
116
117
118
119
    public get isMultiPhaseJobSupported(): boolean {
        return false;
    }

    public getClusterMetadata(key: string): Promise<string> {
        return Promise.resolve('');
    }

    public get MetricsEmitter() : EventEmitter {
        return this.metricsEmitter;
    }

120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
    public async cancelTrialJob(trialJobId: string, isEarlyStopped: boolean = false): Promise<void> {
        const trialJobDetail : KubernetesTrialJobDetail | undefined =  this.trialJobsMap.get(trialJobId);
        if (trialJobDetail === undefined) {
            const errorMessage: string = `CancelTrialJob: trial job id ${trialJobId} not found`;
            this.log.error(errorMessage);

            return Promise.reject(errorMessage);
        }
        if (this.kubernetesCRDClient === undefined) {
            const errorMessage: string = `CancelTrialJob: trial job id ${trialJobId} failed because operatorClient is undefined`;
            this.log.error(errorMessage);

            return Promise.reject(errorMessage);
        }

        try {
            await this.kubernetesCRDClient.deleteKubernetesJob(new Map(
                [
                    ['app', this.NNI_KUBERNETES_TRIAL_LABEL],
                    ['expId', getExperimentId()],
                    ['trialId', trialJobId]
                ]
            ));
        } catch (err) {
            const errorMessage: string = `Delete trial ${trialJobId} failed: ${err}`;
            this.log.error(errorMessage);

            return Promise.reject(errorMessage);
        }

        trialJobDetail.endTime = Date.now();
        trialJobDetail.status = getJobCancelStatus(isEarlyStopped);

        return Promise.resolve();
    }

    public async cleanUp(): Promise<void> {
        this.stopping = true;

        // First, cancel all running kubernetes jobs
        for (const [trialJobId, kubernetesTrialJob] of this.trialJobsMap) {
            if (['RUNNING', 'WAITING', 'UNKNOWN'].includes(kubernetesTrialJob.status)) {
                try {
                    await this.cancelTrialJob(trialJobId);
                } catch (error) {
                  // DONT throw error during cleanup
                }
                kubernetesTrialJob.status = 'SYS_CANCELED';
            }
        }

        // Delete all kubernetes jobs whose expId label is current experiment id
        try {
            if (this.kubernetesCRDClient !== undefined) {
                await this.kubernetesCRDClient.deleteKubernetesJob(new Map(
                    [
                        ['app', this.NNI_KUBERNETES_TRIAL_LABEL],
                        ['expId', getExperimentId()]
                    ]
                ));
            }
        } catch (error) {
            this.log.error(`Delete kubernetes job with label: app=${this.NNI_KUBERNETES_TRIAL_LABEL},\
            expId=${getExperimentId()} failed, error is ${error}`);
        }

        // Unmount NFS
        try {
            await cpp.exec(`sudo umount ${this.trialLocalNFSTempFolder}`);
        } catch (error) {
            this.log.error(`Unmount ${this.trialLocalNFSTempFolder} failed, error is ${error}`);
        }

        // Stop kubernetes rest server
        if (this.kubernetesJobRestServer === undefined) {
            throw new Error('kubernetesJobRestServer not initialized!');
        }
        try {
            await this.kubernetesJobRestServer.stop();
            this.log.info('Kubernetes Training service rest server stopped successfully.');
        } catch (error) {
            // tslint:disable-next-line: no-unsafe-any
            this.log.error(`Kubernetes Training service rest server stopped failed, error: ${error.message}`);

            return Promise.reject(error);
        }

        return Promise.resolve();
    }

    // tslint:disable: no-unsafe-any no-any
211
212
    protected async createAzureStorage(vaultName: string, valutKeyName: string, accountName: string, azureShare: string): Promise<void> {
        try {
213
214
            const result: any = await cpp.exec(`az keyvault secret show --name ${valutKeyName} --vault-name ${vaultName}`);
            if (result.stderr) {
215
216
                const errorMessage: string = result.stderr;
                this.log.error(errorMessage);
217

218
219
                return Promise.reject(errorMessage);
            }
220
221
222
223
            const storageAccountKey: any = JSON.parse(result.stdout).value;
            if (this.azureStorageAccountName === undefined) {
                throw new Error('azureStorageAccountName not initialized!');
            }
224
            //create storage client
225
            this.azureStorageClient = azureStorage.createFileService(this.azureStorageAccountName, storageAccountKey);
226
227
            await AzureStorageClientUtility.createShare(this.azureStorageClient, this.azureStorageShare);
            //create sotrage secret
228
229
            this.azureStorageSecretName = String.Format('nni-secret-{0}', uniqueString(8)
                                                                            .toLowerCase());
230
231
232
233
            await this.genericK8sClient.createSecret(
                {
                    apiVersion: 'v1',
                    kind: 'Secret',
234
                    metadata: {
235
236
237
238
239
240
241
242
243
                        name: this.azureStorageSecretName,
                        namespace: 'default',
                        labels: {
                            app: this.NNI_KUBERNETES_TRIAL_LABEL,
                            expId: getExperimentId()
                        }
                    },
                    type: 'Opaque',
                    data: {
244
245
                        azurestorageaccountname: Base64.encode(this.azureStorageAccountName),
                        azurestorageaccountkey: Base64.encode(storageAccountKey)
246
247
248
                    }
                }
            );
249
        } catch (error) {
250
            this.log.error(error);
251

252
253
            return Promise.reject(error);
        }
254

255
256
        return Promise.resolve();
    }
257
    // tslint:enable: no-unsafe-any no-any
258

259
    /**
260
261
262
     * Genereate run script for different roles(like worker or ps)
     * @param trialJobId trial job id
     * @param trialWorkingFolder working folder
263
     * @param command command
264
265
     * @param trialSequenceId sequence id
     */
266
    protected async generateRunScript(platform: string, trialJobId: string, trialWorkingFolder: string,
267
268
                                      command: string, trialSequenceId: string, roleName: string, gpuNum: number): Promise<string> {
        let nvidiaScript: string = '';
269
270
271
        // Nvidia devcie plugin for K8S has a known issue that requesting zero GPUs allocates all GPUs
        // Refer https://github.com/NVIDIA/k8s-device-plugin/issues/61
        // So we have to explicitly set CUDA_VISIBLE_DEVICES to empty if user sets gpuNum to 0 in NNI config file
272
        if (gpuNum === 0) {
273
            nvidiaScript = 'export CUDA_VISIBLE_DEVICES=';
274
        }
275
276
277
        // tslint:disable-next-line: strict-boolean-expressions
        const nniManagerIp: string = this.nniManagerIpConfig ? this.nniManagerIpConfig.nniManagerIp : getIPV4Address();
        const version: string = this.versionCheck ? await getVersion() : '';
278
        const runScript: string = String.Format(
279
            kubernetesScriptFormat,
280
281
282
283
284
285
286
            platform,
            trialJobId,
            path.join(trialWorkingFolder, 'output', `${roleName}_output`),
            trialJobId,
            getExperimentId(),
            trialWorkingFolder,
            trialSequenceId,
287
            nvidiaScript,
288
289
            command,
            nniManagerIp,
290
            this.kubernetesRestServerPort,
SparkSnail's avatar
SparkSnail committed
291
292
            version,
            this.logCollection
293
        );
294

295
        return Promise.resolve(runScript);
296
297
298
299
300
    }
    protected async createNFSStorage(nfsServer: string, nfsPath: string): Promise<void> {
        await cpp.exec(`mkdir -p ${this.trialLocalNFSTempFolder}`);
        try {
            await cpp.exec(`sudo mount ${nfsServer}:${nfsPath} ${this.trialLocalNFSTempFolder}`);
301
        } catch (error) {
302
303
304
            const mountError: string = `Mount NFS ${nfsServer}:${nfsPath} to ${this.trialLocalNFSTempFolder} failed, error is ${error}`;
            this.log.error(mountError);

305
            return Promise.reject(mountError);
306
307
308
309
        }

        return Promise.resolve();
    }
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337

    protected async createRegistrySecret(filePath: string | undefined): Promise<string | undefined> {
        if(filePath === undefined || filePath === '') {
            return undefined;
        }
        let body = fs.readFileSync(filePath).toString('base64');
        let registrySecretName = String.Format('nni-secret-{0}', uniqueString(8)
                                                                            .toLowerCase());
        await this.genericK8sClient.createSecret(
            {
                apiVersion: 'v1',
                kind: 'Secret',
                metadata: {
                    name: registrySecretName,
                    namespace: 'default',
                    labels: {
                        app: this.NNI_KUBERNETES_TRIAL_LABEL,
                        expId: getExperimentId()
                    }
                },
                type: 'kubernetes.io/dockerconfigjson',
                data: {
                    '.dockerconfigjson': body
                }
            }
        );
        return registrySecretName;
    }
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
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383

    protected async uploadFilesToAzureStorage(trialJobId: string, trialLocalTempFolder: String, codeDir: String, uploadRetryCount: number | undefined): Promise<string> {
        if (this.azureStorageClient === undefined) {
            throw new Error('azureStorageClient is not initialized');
        }
        let trialJobOutputUrl: string = '';
        let retryCount: number = 1;
        if(uploadRetryCount) {
            retryCount = uploadRetryCount;
        }
        let resultUploadNNIScript: boolean = false;
        let resultUploadCodeFile: boolean = false;
        try {
            do {
                //upload local files, including scripts for running the trial and configuration (e.g., hyperparameters) for the trial, to azure storage
                if(!resultUploadNNIScript) {
                    resultUploadNNIScript = await AzureStorageClientUtility.uploadDirectory(this.azureStorageClient,
                        `nni/${getExperimentId()}/${trialJobId}`, this.azureStorageShare,
                        `${trialLocalTempFolder}`);
                }
                //upload code files to azure storage
                if(!resultUploadCodeFile) {
                    resultUploadCodeFile = await AzureStorageClientUtility.uploadDirectory(this.azureStorageClient,
                        `nni/${getExperimentId()}/${trialJobId}`, this.azureStorageShare,
                        `${codeDir}`);
                }
                if (resultUploadNNIScript && resultUploadCodeFile) {
                    trialJobOutputUrl = `https://${this.azureStorageAccountName}.file.core.windows.net/${this.azureStorageShare}` + 
                    `/${path.join('nni', getExperimentId(), trialJobId, 'output')}`;
                    break;
                } else {
                    //wait for 5 seconds to re-upload files
                    await delay(5000);
                    this.log.info('Upload failed, Retry: upload files to azure-storage');
                }
            } while (retryCount-- >= 0)
        } catch (error) {
            this.log.error(error);
            //return a empty url when got error
            return Promise.resolve("");
        }
        if(!trialJobOutputUrl) {
            this.log.info(`Retry-count is used up, upload files to azureStorage for trial ${trialJobId} failed!`);
        }
        return Promise.resolve(trialJobOutputUrl);
    }
384
     
385
}
386
export { KubernetesTrainingService };