kubernetesTrainingService.ts 13.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
/**
 * Copyright (c) Microsoft Corporation
 * All rights reserved.
 *
 * MIT License
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
 * documentation files (the "Software"), to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
 * to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
 * BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

20
'use strict';
21
22
23
24

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

25
import * as azureStorage from 'azure-storage';
26
import { EventEmitter } from 'events';
27
28
import { Base64 } from 'js-base64';
import { String } from 'typescript-string-operations';
29
30
31
import { getExperimentId, getInitTrialSequenceId } from '../../common/experimentStartupInfo';
import { getLogger, Logger } from '../../common/log';
import {
32
    NNIManagerIpConfig, TrialJobDetail, TrialJobMetric
33
} from '../../common/trainingService';
34
import { getExperimentRootDir, getIPV4Address, getJobCancelStatus, getVersion, uniqueString } from '../../common/utils';
35
import { AzureStorageClientUtility } from './azureStorageClientUtils';
36
37
38
import { GeneralK8sClient, KubernetesCRDClient } from './kubernetesApiClient';
import { KubernetesClusterConfig } from './kubernetesConfig';
import { kubernetesScriptFormat, KubernetesTrialJobDetail } from './kubernetesData';
39
40
import { KubernetesJobRestServer } from './kubernetesJobRestServer';

41
42
43
/**
 * Training Service implementation for Kubernetes
 */
44
45
46
47
48
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>;
49
    //  experiment root dir in NFS
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    protected readonly trialLocalNFSTempFolder: string;
    protected stopping: boolean = false;
    protected experimentId! : string;
    protected nextTrialSequenceId: number;
    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;
65
    protected versionCheck: boolean = true;
SparkSnail's avatar
SparkSnail committed
66
    protected logCollection: string;
67

68
69
70
71
72
    constructor() {
        this.log = getLogger();
        this.metricsEmitter = new EventEmitter();
        this.trialJobsMap = new Map<string, KubernetesTrialJobDetail>();
        this.trialLocalNFSTempFolder = path.join(getExperimentRootDir(), 'trials-nfs-tmp');
73
        this.experimentId = getExperimentId();
74
75
76
        this.nextTrialSequenceId = -1;
        this.CONTAINER_MOUNT_PATH = '/tmp/mount';
        this.genericK8sClient = new GeneralK8sClient();
SparkSnail's avatar
SparkSnail committed
77
        this.logCollection = 'none';
78
79
    }

80
81
    // tslint:disable:no-any
    public generatePodResource(memory: number, cpuNum: number, gpuNum: number): any {
82
        return {
83
84
            memory: `${memory}Mi`,
            cpu: `${cpuNum}`,
85
            'nvidia.com/gpu': `${gpuNum}`
86
87
        };
    } // tslint:enable:no-any
88

89
    public async listTrialJobs(): Promise<TrialJobDetail[]> {
90
        const jobs: TrialJobDetail[] = [];
91
92

        for (const [key, value] of this.trialJobsMap) {
93
94
95
            if (value.form.jobType === 'TRIAL') {
                jobs.push(await this.getTrialJob(key));
            }
96
        }
97
98
99
100

        return Promise.resolve(jobs);
    }

101
    public async getTrialJob(trialJobId: string): Promise<TrialJobDetail> {
102
103
104

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

105
106
        if (kubernetesTrialJob === undefined) {
            return Promise.reject(`trial job ${trialJobId} not found`);
107
        }
108
109
110
111

        return Promise.resolve(kubernetesTrialJob);
    }

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

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

120
121
122
123
124
125
126
127
128
129
130
131
    public get isMultiPhaseJobSupported(): boolean {
        return false;
    }

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

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

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
211
212
213
214
215
216
217
218
219
220
221
    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();
    }

222
223
224
225
226
227
228
229
    protected generateSequenceId(): number {
        if (this.nextTrialSequenceId === -1) {
            this.nextTrialSequenceId = getInitTrialSequenceId();
        }

        return this.nextTrialSequenceId++;
    }

230
    // tslint:disable: no-unsafe-any no-any
231
232
    protected async createAzureStorage(vaultName: string, valutKeyName: string, accountName: string, azureShare: string): Promise<void> {
        try {
233
234
            const result: any = await cpp.exec(`az keyvault secret show --name ${valutKeyName} --vault-name ${vaultName}`);
            if (result.stderr) {
235
236
                const errorMessage: string = result.stderr;
                this.log.error(errorMessage);
237

238
239
                return Promise.reject(errorMessage);
            }
240
241
242
243
            const storageAccountKey: any = JSON.parse(result.stdout).value;
            if (this.azureStorageAccountName === undefined) {
                throw new Error('azureStorageAccountName not initialized!');
            }
244
            //create storage client
245
            this.azureStorageClient = azureStorage.createFileService(this.azureStorageAccountName, storageAccountKey);
246
247
            await AzureStorageClientUtility.createShare(this.azureStorageClient, this.azureStorageShare);
            //create sotrage secret
248
249
            this.azureStorageSecretName = String.Format('nni-secret-{0}', uniqueString(8)
                                                                            .toLowerCase());
250
251
252
253
            await this.genericK8sClient.createSecret(
                {
                    apiVersion: 'v1',
                    kind: 'Secret',
254
                    metadata: {
255
256
257
258
259
260
261
262
263
                        name: this.azureStorageSecretName,
                        namespace: 'default',
                        labels: {
                            app: this.NNI_KUBERNETES_TRIAL_LABEL,
                            expId: getExperimentId()
                        }
                    },
                    type: 'Opaque',
                    data: {
264
265
                        azurestorageaccountname: Base64.encode(this.azureStorageAccountName),
                        azurestorageaccountkey: Base64.encode(storageAccountKey)
266
267
268
                    }
                }
            );
269
        } catch (error) {
270
            this.log.error(error);
271

272
273
            return Promise.reject(error);
        }
274

275
276
        return Promise.resolve();
    }
277
    // tslint:enable: no-unsafe-any no-any
278

279
    /**
280
281
282
     * Genereate run script for different roles(like worker or ps)
     * @param trialJobId trial job id
     * @param trialWorkingFolder working folder
283
     * @param command command
284
285
     * @param trialSequenceId sequence id
     */
286
    protected async generateRunScript(platform: string, trialJobId: string, trialWorkingFolder: string,
287
288
                                      command: string, trialSequenceId: string, roleName: string, gpuNum: number): Promise<string> {
        let nvidiaScript: string = '';
289
290
291
        // 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
292
293
        if (gpuNum === 0) {
            nvidiaScript = `export CUDA_VISIBLE_DEVICES='0'`;
294
        }
295
296
297
        // tslint:disable-next-line: strict-boolean-expressions
        const nniManagerIp: string = this.nniManagerIpConfig ? this.nniManagerIpConfig.nniManagerIp : getIPV4Address();
        const version: string = this.versionCheck ? await getVersion() : '';
298
        const runScript: string = String.Format(
299
            kubernetesScriptFormat,
300
301
302
303
304
305
306
            platform,
            trialJobId,
            path.join(trialWorkingFolder, 'output', `${roleName}_output`),
            trialJobId,
            getExperimentId(),
            trialWorkingFolder,
            trialSequenceId,
307
            nvidiaScript,
308
309
            command,
            nniManagerIp,
310
            this.kubernetesRestServerPort,
SparkSnail's avatar
SparkSnail committed
311
312
            version,
            this.logCollection
313
        );
314

315
        return Promise.resolve(runScript);
316
317
318
319
320
    }
    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}`);
321
        } catch (error) {
322
323
324
            const mountError: string = `Mount NFS ${nfsServer}:${nfsPath} to ${this.trialLocalNFSTempFolder} failed, error is ${error}`;
            this.log.error(mountError);

325
            return Promise.reject(mountError);
326
327
328
329
330
        }

        return Promise.resolve();
    }
}
331
export { KubernetesTrainingService };