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

4
5
import cpp from 'child-process-promise';
import path from 'path';
6

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

chicm-ms's avatar
chicm-ms committed
24
const fs = require('fs');
25

26
27
28
/**
 * Training Service implementation for Kubernetes
 */
29
30
31
32
33
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>;
34
    //  experiment root dir in NFS
35
    protected readonly trialLocalTempFolder: string;
36
    protected stopping: boolean = false;
chicm-ms's avatar
chicm-ms committed
37
    protected experimentId!: string;
38
39
40
41
42
43
44
45
46
47
48
    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;
49
    protected versionCheck: boolean = true;
SparkSnail's avatar
SparkSnail committed
50
    protected logCollection: string;
51
52
    protected copyExpCodeDirPromise?: Promise<string>;
    protected expContainerCodeFolder: string;
53

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

66
    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
    }
78

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

82
        for (const key of this.trialJobsMap.keys()) {
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);
    }

Yuge Zhang's avatar
Yuge Zhang committed
100
    public async getTrialFile(_trialJobId: string, _filename: string): Promise<string | Buffer> {
101
102
103
        throw new MethodNotImplementedError();
    }

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

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

112
113
114
115
    public get isMultiPhaseJobSupported(): boolean {
        return false;
    }

116
    public getClusterMetadata(_key: string): Promise<string> {
117
118
119
        return Promise.resolve('');
    }

chicm-ms's avatar
chicm-ms committed
120
    public get MetricsEmitter(): EventEmitter {
121
122
123
        return this.metricsEmitter;
    }

124
    public async cancelTrialJob(trialJobId: string, isEarlyStopped: boolean = false): Promise<void> {
125
        const trialJobDetail: KubernetesTrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);
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
        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) {
169
                    // DONT throw error during cleanup
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
                }
                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 {
192
            await cpp.exec(`sudo umount ${this.trialLocalTempFolder}`);
193
        } catch (error) {
194
            this.log.error(`Unmount ${this.trialLocalTempFolder} failed, error is ${error}`);
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
        }

        // 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) {
            this.log.error(`Kubernetes Training service rest server stopped failed, error: ${error.message}`);

            return Promise.reject(error);
        }

        return Promise.resolve();
    }

chicm-ms's avatar
chicm-ms committed
213
    protected async createAzureStorage(vaultName: string, valutKeyName: string): Promise<void> {
214
        try {
215
216
            const result: any = await cpp.exec(`az keyvault secret show --name ${valutKeyName} --vault-name ${vaultName}`);
            if (result.stderr) {
217
218
                const errorMessage: string = result.stderr;
                this.log.error(errorMessage);
219

220
221
                return Promise.reject(errorMessage);
            }
222
223
224
225
            const storageAccountKey: any = JSON.parse(result.stdout).value;
            if (this.azureStorageAccountName === undefined) {
                throw new Error('azureStorageAccountName not initialized!');
            }
226
            //create storage client
227
            this.azureStorageClient = azureStorage.createFileService(this.azureStorageAccountName, storageAccountKey);
228
229
            await AzureStorageClientUtility.createShare(this.azureStorageClient, this.azureStorageShare);
            //create sotrage secret
230
            this.azureStorageSecretName = String.Format('nni-secret-{0}', uniqueString(8)
231
232
                .toLowerCase());

J-shang's avatar
J-shang committed
233
            const namespace = this.genericK8sClient.getNamespace ?? "default";
234
235
236
237
            await this.genericK8sClient.createSecret(
                {
                    apiVersion: 'v1',
                    kind: 'Secret',
238
                    metadata: {
239
                        name: this.azureStorageSecretName,
240
                        namespace: namespace,
241
242
243
244
245
246
247
                        labels: {
                            app: this.NNI_KUBERNETES_TRIAL_LABEL,
                            expId: getExperimentId()
                        }
                    },
                    type: 'Opaque',
                    data: {
248
249
                        azurestorageaccountname: Base64.encode(this.azureStorageAccountName),
                        azurestorageaccountkey: Base64.encode(storageAccountKey)
250
251
252
                    }
                }
            );
253
        } catch (error) {
254
            this.log.error(error);
255

256
257
            return Promise.reject(error);
        }
258

259
260
        return Promise.resolve();
    }
261

262
    /**
263
264
265
     * Genereate run script for different roles(like worker or ps)
     * @param trialJobId trial job id
     * @param trialWorkingFolder working folder
266
     * @param command command
267
268
     * @param trialSequenceId sequence id
     */
269
    protected async generateRunScript(platform: string, trialJobId: string, trialWorkingFolder: string,
270
        command: string, trialSequenceId: string, roleName: string, gpuNum: number): Promise<string> {
271
        let nvidiaScript: string = '';
272
273
274
        // 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
275
        if (gpuNum === 0) {
276
            nvidiaScript = 'export CUDA_VISIBLE_DEVICES=';
277
        }
liuzhe-lz's avatar
liuzhe-lz committed
278
        const nniManagerIp: string = this.nniManagerIpConfig ? this.nniManagerIpConfig.nniManagerIp : await getIPV4Address();
279
        const version: string = this.versionCheck ? await getVersion() : '';
280
        const runScript: string = String.Format(
281
            kubernetesScriptFormat,
282
            platform,
283
            trialWorkingFolder,
284
285
286
            path.join(trialWorkingFolder, 'output', `${roleName}_output`),
            trialJobId,
            getExperimentId(),
287
            this.expContainerCodeFolder,
288
            trialSequenceId,
289
            nvidiaScript,
290
291
            command,
            nniManagerIp,
292
            this.kubernetesRestServerPort,
SparkSnail's avatar
SparkSnail committed
293
294
            version,
            this.logCollection
295
        );
296

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

307
            return Promise.reject(mountError);
308
309
310
311
        }

        return Promise.resolve();
    }
312
313
314
315
316
317
318
319
320
321
322
323
324
    protected async createPVCStorage(pvcPath: string): Promise<void> {
        try {
            await cpp.exec(`mkdir -p ${pvcPath}`);
            await cpp.exec(`sudo ln -s ${pvcPath} ${this.trialLocalTempFolder}`);
        } catch (error) {
            const linkError: string = `Linking ${pvcPath} to ${this.trialLocalTempFolder} failed, error is ${error}`;
            this.log.error(linkError);

            return Promise.reject(linkError);
        }

        return Promise.resolve();
    }
325
326

    protected async createRegistrySecret(filePath: string | undefined): Promise<string | undefined> {
327
        if (filePath === undefined || filePath === '') {
328
329
            return undefined;
        }
chicm-ms's avatar
chicm-ms committed
330
331
        const body = fs.readFileSync(filePath).toString('base64');
        const registrySecretName = String.Format('nni-secret-{0}', uniqueString(8)
332
            .toLowerCase());
J-shang's avatar
J-shang committed
333
        const namespace = this.genericK8sClient.getNamespace ?? "default";
334
335
336
337
338
339
        await this.genericK8sClient.createSecret(
            {
                apiVersion: 'v1',
                kind: 'Secret',
                metadata: {
                    name: registrySecretName,
340
                    namespace: namespace,
341
342
343
344
345
346
347
348
349
350
351
352
353
                    labels: {
                        app: this.NNI_KUBERNETES_TRIAL_LABEL,
                        expId: getExperimentId()
                    }
                },
                type: 'kubernetes.io/dockerconfigjson',
                data: {
                    '.dockerconfigjson': body
                }
            }
        );
        return registrySecretName;
    }
354

355
356
357
358
359
360
361
    /**
     * upload local directory to azureStorage
     * @param srcDirectory the source directory of local folder
     * @param destDirectory the target directory in azure
     * @param uploadRetryCount the retry time when upload failed
     */
    protected async uploadFolderToAzureStorage(srcDirectory: string, destDirectory: string, uploadRetryCount: number | undefined): Promise<string> {
362
363
364
365
        if (this.azureStorageClient === undefined) {
            throw new Error('azureStorageClient is not initialized');
        }
        let retryCount: number = 1;
366
        if (uploadRetryCount) {
367
368
            retryCount = uploadRetryCount;
        }
369
370
        let uploadSuccess: boolean = false;
        let folderUriInAzure = '';
371
372
        try {
            do {
373
374
                uploadSuccess = await AzureStorageClientUtility.uploadDirectory(
                    this.azureStorageClient,
375
                    `${destDirectory}`,
376
377
378
                    this.azureStorageShare,
                    `${srcDirectory}`);
                if (!uploadSuccess) {
379
380
381
                    //wait for 5 seconds to re-upload files
                    await delay(5000);
                    this.log.info('Upload failed, Retry: upload files to azure-storage');
382
383
384
                } else {
                    folderUriInAzure = `https://${this.azureStorageAccountName}.file.core.windows.net/${this.azureStorageShare}/${destDirectory}`;
                    break;
385
386
387
388
389
                }
            } while (retryCount-- >= 0)
        } catch (error) {
            this.log.error(error);
            //return a empty url when got error
390
            return Promise.resolve('');
391
        }
392
        return Promise.resolve(folderUriInAzure);
393
    }
J-shang's avatar
J-shang committed
394
395
396
397
398
399
400
401

    public getTrialOutputLocalPath(_trialJobId: string): Promise<string> {
        throw new MethodNotImplementedError();
    }

    public fetchTrialOutput(_trialJobId: string, _subpath: string): Promise<void> {
        throw new MethodNotImplementedError();
    }
402
}
403
export {KubernetesTrainingService};