kubeflowTrainingService.ts 33.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
/**
 * 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.
 */

'use strict'

import * as assert from 'assert';
import * as component from '../../common/component';
import * as cpp from 'child-process-promise';
import * as fs from 'fs';
import * as path from 'path';

import { CONTAINER_INSTALL_NNI_SHELL_FORMAT } from '../common/containerJobData';
import { EventEmitter } from 'events';
import { getExperimentId, getInitTrialSequenceId } from '../../common/experimentStartupInfo';
import { getLogger, Logger } from '../../common/log';
import { MethodNotImplementedError } from '../../common/errors';
import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
import {
    JobApplicationForm, TrainingService, TrialJobApplicationForm,
36
    TrialJobDetail, TrialJobMetric, NNIManagerIpConfig
37
} from '../../common/trainingService';
QuanluZhang's avatar
QuanluZhang committed
38
import { delay, generateParamFileName, getExperimentRootDir, getIPV4Address, uniqueString, getJobCancelStatus } from '../../common/utils';
39
40
import { KubeflowClusterConfigBase, KubeflowClusterConfigNFS, KubeflowClusterConfigAzure, kubeflowOperatorMap, KubeflowTrialConfigBase,
     KubeflowTrialConfigPytorch, KubeflowTrialConfigTensorflow, NFSConfig, kubeflowOperatorJobKindMap } from './kubeflowConfig';
41
import { KubeflowTrialJobDetail } from './kubeflowData';
42
43
import { KubeflowJobRestServer } from './kubeflowJobRestServer';
import { KubeflowJobInfoCollector } from './kubeflowJobInfoCollector';
44
import { validateCodeDir } from '../common/util';
SparkSnail's avatar
SparkSnail committed
45
46
import { AzureStorageClientUtility } from './azureStorageClientUtils';
import * as azureStorage from 'azure-storage';
47
48

var yaml = require('node-yaml');
SparkSnail's avatar
SparkSnail committed
49
var azure = require('azure-storage');
50

51
type DistTrainRole = 'worker' | 'ps' | 'master';
52

53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
/**
 * Training Service implementation for Kubeflow
 * Refer https://github.com/kubeflow/kubeflow for more info about Kubeflow
 */
@component.Singleton
class KubeflowTrainingService implements TrainingService {
    private readonly NNI_KUBEFLOW_TRIAL_LABEL = 'nni-kubeflow-trial';
    private readonly log!: Logger;
    private readonly metricsEmitter: EventEmitter;
    private readonly trialJobsMap: Map<string, KubeflowTrialJobDetail>;
    /**  experiment root dir in NFS */
    private readonly trialLocalNFSTempFolder: string;
    private stopping: boolean = false;
    private experimentId! : string;
    private nextTrialSequenceId: number;
68
69
    private kubeflowClusterConfig?: KubeflowClusterConfigBase;
    private kubeflowTrialConfig?: KubeflowTrialConfigBase;
70
71
72
    private kubeflowJobInfoCollector: KubeflowJobInfoCollector;
    private kubeflowRestServerPort?: number;
    private kubeflowJobPlural?: string;
73
    private kubeflowJobKind?: string;
74
    private readonly CONTAINER_MOUNT_PATH: string;
SparkSnail's avatar
SparkSnail committed
75
76
77
78
    private azureStorageClient?: azureStorage.FileService;
    private azureStorageShare?: string;
    private azureStorageSecretName?: string;
    private azureStorageAccountName?: string;
79
    private nniManagerIpConfig?: NNIManagerIpConfig;
80
81
82
83
84
85
86
87
88
    
    constructor() {        
        this.log = getLogger();
        this.metricsEmitter = new EventEmitter();
        this.trialJobsMap = new Map<string, KubeflowTrialJobDetail>();
        this.kubeflowJobInfoCollector = new KubeflowJobInfoCollector(this.trialJobsMap);
        this.trialLocalNFSTempFolder = path.join(getExperimentRootDir(), 'trials-nfs-tmp');
        this.experimentId = getExperimentId();      
        this.nextTrialSequenceId = -1;
SparkSnail's avatar
SparkSnail committed
89
        this.CONTAINER_MOUNT_PATH = '/tmp/mount';
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
    }

    public async run(): Promise<void> {
        const restServer: KubeflowJobRestServer = component.get(KubeflowJobRestServer);
        await restServer.start();
        this.log.info(`Kubeflow Training service rest server listening on: ${restServer.endPoint}`);
        while (!this.stopping) {
            // collect metrics by calling 'kubectl get' command on Kubeflow jobs 
            await delay(3000);
            await this.kubeflowJobInfoCollector.retrieveTrialStatus();            
        }
    }

    public async submitTrialJob(form: JobApplicationForm): Promise<TrialJobDetail> {
        if(!this.kubeflowClusterConfig) {
            throw new Error('Kubeflow Cluster config is not initialized');
        }

108
109
        if(!this.kubeflowTrialConfig) {
            throw new Error('Kubeflow trial config is not initialized');
110
111
112
113
114
115
116
117
118
119
        }

        if(!this.kubeflowJobPlural) {
            throw new Error('Kubeflow job plural name is undefined');
        }

        if(!this.kubeflowRestServerPort) {
            const restServer: KubeflowJobRestServer = component.get(KubeflowJobRestServer);
            this.kubeflowRestServerPort = restServer.clusterRestServerPort;
        }
120
121
122
123
124
125
126
127
128
        // initialize kubeflow trial config to specific type
        let kubeflowTrialConfig;
        if(this.kubeflowClusterConfig.operator === 'tf-operator') {
            kubeflowTrialConfig = <KubeflowTrialConfigTensorflow>this.kubeflowTrialConfig;
        }else if(this.kubeflowClusterConfig.operator === 'pytorch-operator'){
            kubeflowTrialConfig = <KubeflowTrialConfigPytorch>this.kubeflowTrialConfig;
        }else {
            throw Error(`operator ${this.kubeflowClusterConfig.operator} is invalid`)
        }
129
130
131
132
133
134
135
136

        const trialJobId: string = uniqueString(5);
        const curTrialSequenceId: number = this.generateSequenceId();
        // Set trial's NFS working folder
        const trialWorkingFolder: string = path.join(this.CONTAINER_MOUNT_PATH, 'nni', getExperimentId(), trialJobId);
        const trialLocalTempFolder: string = path.join(getExperimentRootDir(), 'trials-local', trialJobId);
        //create tmp trial working folder locally.
        await cpp.exec(`mkdir -p ${path.dirname(trialLocalTempFolder)}`);
137
        await cpp.exec(`cp -r ${kubeflowTrialConfig.codeDir} ${trialLocalTempFolder}`);
138
139
140
        const runScriptContent : string = CONTAINER_INSTALL_NNI_SHELL_FORMAT;
        // Write NNI installation file to local tmp files
        await fs.promises.writeFile(path.join(trialLocalTempFolder, 'install_nni.sh'), runScriptContent, { encoding: 'utf8' });
141
        // Create tmp trial working folder locally.
142
143
        await cpp.exec(`mkdir -p ${trialLocalTempFolder}`);

144
        // Write worker file content run_worker.sh to local tmp folders
145
146
147
        if(kubeflowTrialConfig.worker) {
            const workerRunScriptContent: string = this.generateRunScript(trialJobId, trialWorkingFolder, 
                    kubeflowTrialConfig.worker.command, curTrialSequenceId.toString(), 'worker', kubeflowTrialConfig.worker.gpuNum);
148
149
150
151

            await fs.promises.writeFile(path.join(trialLocalTempFolder, 'run_worker.sh'), workerRunScriptContent, { encoding: 'utf8' });
        }
        // Write parameter server file content run_ps.sh to local tmp folders
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
        if(this.kubeflowClusterConfig.operator === 'tf-operator') {
            let tensorflowTrialConfig: KubeflowTrialConfigTensorflow = <KubeflowTrialConfigTensorflow>this.kubeflowTrialConfig;
            if(tensorflowTrialConfig.ps){
                const psRunScriptContent: string = this.generateRunScript(trialJobId, trialWorkingFolder, 
                    tensorflowTrialConfig.ps.command, curTrialSequenceId.toString(), 'ps', tensorflowTrialConfig.ps.gpuNum);
                await fs.promises.writeFile(path.join(trialLocalTempFolder, 'run_ps.sh'), psRunScriptContent, { encoding: 'utf8' });
            }
        }
        else if(this.kubeflowClusterConfig.operator === 'pytorch-operator') {
            let pytorchTrialConfig: KubeflowTrialConfigPytorch = <KubeflowTrialConfigPytorch>this.kubeflowTrialConfig;
            if(pytorchTrialConfig.master){
                const masterRunScriptContent: string = this.generateRunScript(trialJobId, trialWorkingFolder, 
                    pytorchTrialConfig.master.command, curTrialSequenceId.toString(), 'master', pytorchTrialConfig.master.gpuNum);
                await fs.promises.writeFile(path.join(trialLocalTempFolder, 'run_master.sh'), masterRunScriptContent, { encoding: 'utf8' });
            }
167
        }
168
169
170
171
172
        // Write file content ( parameter.cfg ) to local tmp folders
        const trialForm : TrialJobApplicationForm = (<TrialJobApplicationForm>form)
        if(trialForm && trialForm.hyperParameters) {
            await fs.promises.writeFile(path.join(trialLocalTempFolder, generateParamFileName(trialForm.hyperParameters)), 
                            trialForm.hyperParameters.value, { encoding: 'utf8' });
173
        }
174
175
        const kubeflowJobYamlPath = path.join(trialLocalTempFolder, `kubeflow-job-${trialJobId}.yaml`);
        const kubeflowJobName = `nni-exp-${this.experimentId}-trial-${trialJobId}`.toLowerCase();
176
        
177
        const workerPodResources : any = {};
178
179
        workerPodResources.requests = this.generatePodResource(kubeflowTrialConfig.worker.memoryMB, kubeflowTrialConfig.worker.cpuNum, 
            kubeflowTrialConfig.worker.gpuNum)
180
181
        workerPodResources.limits = Object.assign({}, workerPodResources.requests);

182
183
184
185
186
187
188
        let nonWorkerResources : any = {};
        if(this.kubeflowClusterConfig.operator === 'tf-operator') {
            let tensorflowTrialConfig: KubeflowTrialConfigTensorflow = <KubeflowTrialConfigTensorflow>this.kubeflowTrialConfig;
            if (tensorflowTrialConfig.ps) {
                nonWorkerResources.requests = this.generatePodResource(tensorflowTrialConfig.ps.memoryMB, tensorflowTrialConfig.ps.cpuNum, 
                    tensorflowTrialConfig.ps.gpuNum)
                    nonWorkerResources.limits = Object.assign({}, nonWorkerResources.requests); 
189
            }
190
191
192
193
194
195
196
197
        }else if(this.kubeflowClusterConfig.operator === 'pytorch-operator'){
            let pyTorchTrialConfig: KubeflowTrialConfigPytorch = <KubeflowTrialConfigPytorch>this.kubeflowTrialConfig;
            if (pyTorchTrialConfig.master) {
                nonWorkerResources.requests = this.generatePodResource(pyTorchTrialConfig.master.memoryMB, pyTorchTrialConfig.master.cpuNum, 
                    pyTorchTrialConfig.master.gpuNum)
                    nonWorkerResources.limits = Object.assign({}, nonWorkerResources.requests); 
            }
        }       
198
199
200
        // Generate kubeflow job resource yaml file for K8S
        yaml.write(
            kubeflowJobYamlPath,
201
            this.generateKubeflowJobConfig(trialJobId, trialWorkingFolder, kubeflowJobName, workerPodResources, nonWorkerResources),
202
203
            'utf-8'
        );
SparkSnail's avatar
SparkSnail committed
204
205
206
        let trialJobDetail: KubeflowTrialJobDetail;
        //The url used in trialJobDetail
        let trialJobDetailUrl: string;
207
        if(this.kubeflowClusterConfig.storage && this.kubeflowClusterConfig.storage === 'azureStorage') {
SparkSnail's avatar
SparkSnail committed
208
209
210
211
212
213
214
215
216
217
            try{
                //upload local files to azure storage
                await AzureStorageClientUtility.uploadDirectory(this.azureStorageClient, 
                    `nni/${getExperimentId()}/${trialJobId}`, this.azureStorageShare, `${trialLocalTempFolder}`);

                trialJobDetailUrl = `https://${this.azureStorageAccountName}.file.core.windows.net/${this.azureStorageShare}/${path.join('nni', getExperimentId(), trialJobId, 'output')}`
            }catch(error){
                this.log.error(error);
                return Promise.reject(error);
            }
218
219
220
221
222
223
224
225
226
227
228
229
230
        } else if(this.kubeflowClusterConfig.storage && (this.kubeflowClusterConfig.storage === 'nfs' || this.kubeflowClusterConfig.storage === undefined)) {
            let nfsKubeflowClusterConfig: KubeflowClusterConfigNFS = <KubeflowClusterConfigNFS>this.kubeflowClusterConfig;
            // Creat work dir for current trial in NFS directory 
            await cpp.exec(`mkdir -p ${this.trialLocalNFSTempFolder}/nni/${getExperimentId()}/${trialJobId}`);
            // Copy code files from local dir to NFS mounted dir
            await cpp.exec(`cp -r ${trialLocalTempFolder}/* ${this.trialLocalNFSTempFolder}/nni/${getExperimentId()}/${trialJobId}/.`);
        
            const nfsConfig: NFSConfig = nfsKubeflowClusterConfig.nfs;
            trialJobDetailUrl = `nfs://${nfsConfig.server}:${path.join(nfsConfig.path, 'nni', getExperimentId(), trialJobId, 'output')}`
        } else {
            const error: string = `kubeflowClusterConfig format error!`;
            this.log.error(error);
            throw new Error(error);
SparkSnail's avatar
SparkSnail committed
231
        }
232

SparkSnail's avatar
SparkSnail committed
233
        trialJobDetail = new KubeflowTrialJobDetail(
234
235
236
237
238
239
240
            trialJobId,
            'WAITING',
            Date.now(),
            trialWorkingFolder,
            form,
            kubeflowJobName,
            curTrialSequenceId,
SparkSnail's avatar
SparkSnail committed
241
            trialJobDetailUrl, 
242
            this.kubeflowJobPlural
SparkSnail's avatar
SparkSnail committed
243
        );
244
        
245
246
247
248
249
250
251
252
        // Create kubeflow training jobs
        await cpp.exec(`kubectl create -f ${kubeflowJobYamlPath}`);
        // Set trial job detail until kubectl create resource successfully 
        this.trialJobsMap.set(trialJobId, trialJobDetail);

        return Promise.resolve(trialJobDetail);
    }

253
254
255
256
257
258
259
260
    public generatePodResource(memory: number, cpuNum: number, gpuNum: number) {
        return {
            'memory': `${memory}Mi`,
            'cpu': `${cpuNum}`,
            'nvidia.com/gpu': `${gpuNum}`
        }
    }

261
262
263
264
265
    public updateTrialJob(trialJobId: string, form: JobApplicationForm): Promise<TrialJobDetail> {
        throw new MethodNotImplementedError();
    }

    public listTrialJobs(): Promise<TrialJobDetail[]> {
266
267
268
269
270
271
272
273
274
        const jobs: TrialJobDetail[] = [];
        
        this.trialJobsMap.forEach(async (value: KubeflowTrialJobDetail, key: string) => {
            if (value.form.jobType === 'TRIAL') {
                jobs.push(await this.getTrialJob(key));
            }
        });

        return Promise.resolve(jobs);
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
    }

    public getTrialJob(trialJobId: string): Promise<TrialJobDetail> {
        if(!this.kubeflowClusterConfig) {
            throw new Error('Kubeflow Cluster config is not initialized');
        }

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

        if (!kubeflowTrialJob) {
            return Promise.reject(`trial job ${trialJobId} not found`)
        }        

        return Promise.resolve(kubeflowTrialJob);
    }

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

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

QuanluZhang's avatar
QuanluZhang committed
303
    public async cancelTrialJob(trialJobId: string, isEarlyStopped: boolean = false): Promise<void> {
304
305
306
307
308
309
310
311
312
313
314
315
        const trialJobDetail : KubeflowTrialJobDetail | undefined =  this.trialJobsMap.get(trialJobId);
        if(!trialJobDetail) {
            const errorMessage: string = `CancelTrialJob: trial job id ${trialJobId} not found`;
            this.log.error(errorMessage);
            return Promise.reject(errorMessage);
        }
        if(!this.kubeflowJobPlural) {
            const errorMessage: string = `CancelTrialJob: trial job id ${trialJobId} failed because kubeflowJobPlural is undefined`;
            this.log.error(errorMessage);
            return Promise.reject(errorMessage);
        }

316
317
        const result: cpp.childProcessPromise.Result = await cpp.exec(`kubectl delete ${this.kubeflowJobPlural} -l ` 
                                        + `app=${this.NNI_KUBEFLOW_TRIAL_LABEL},expId=${getExperimentId()},trialId=${trialJobId}`);
318
319
320
321
322
323
324
        if(result.stderr) {
            const errorMessage: string = `kubectl delete ${this.kubeflowJobPlural} for trial ${trialJobId} failed: ${result.stderr}`;
            this.log.error(errorMessage);
            return Promise.reject(errorMessage);
        }

        trialJobDetail.endTime = Date.now();
QuanluZhang's avatar
QuanluZhang committed
325
        trialJobDetail.status = getJobCancelStatus(isEarlyStopped);
326
327
328
329
330
331

        return Promise.resolve();
    }

    public async setClusterMetadata(key: string, value: string): Promise<void> {
        switch (key) {
332
333
334
335
            case TrialConfigMetadataKey.NNI_MANAGER_IP:
                this.nniManagerIpConfig = <NNIManagerIpConfig>JSON.parse(value);
                break;
            
336
            case TrialConfigMetadataKey.KUBEFLOW_CLUSTER_CONFIG:
337
338
339
340
341
342
343
344
345
346
347
                let kubeflowClusterJsonObject = JSON.parse(value);
                let kubeflowClusterConfigBase: KubeflowClusterConfigBase = new KubeflowClusterConfigBase(kubeflowClusterJsonObject.operator, kubeflowClusterJsonObject.storage);
                
                if(kubeflowClusterConfigBase && kubeflowClusterConfigBase.storage === 'azureStorage') {
                    this.kubeflowClusterConfig = new KubeflowClusterConfigAzure(kubeflowClusterJsonObject.operator, kubeflowClusterJsonObject.keyvault, 
                        kubeflowClusterJsonObject.azureStorage, kubeflowClusterJsonObject.storage)
                    let azureKubeflowClusterConfig = <KubeflowClusterConfigAzure>this.kubeflowClusterConfig;
                    const vaultName = azureKubeflowClusterConfig.keyVault.vaultName;
                    const valutKeyName = azureKubeflowClusterConfig.keyVault.name;
                    this.azureStorageAccountName = azureKubeflowClusterConfig.azureStorage.accountName;
                    this.azureStorageShare = azureKubeflowClusterConfig.azureStorage.azureShare;
SparkSnail's avatar
SparkSnail committed
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
                    try{
                        const result = await cpp.exec(`az keyvault secret show --name ${valutKeyName} --vault-name ${vaultName}`);
                        if(result.stderr) {
                            const errorMessage: string = result.stderr;
                            this.log.error(errorMessage);
                            return Promise.reject(errorMessage);
                        }
                        const storageAccountKey =JSON.parse(result.stdout).value;
                        //create storage client
                        this.azureStorageClient = azure.createFileService(this.azureStorageAccountName, storageAccountKey);
                        await AzureStorageClientUtility.createShare(this.azureStorageClient, this.azureStorageShare);
                        //create sotrage secret
                        this.azureStorageSecretName = 'nni-secret-' + uniqueString(8).toLowerCase();
                        await cpp.exec(`kubectl create secret generic ${this.azureStorageSecretName} `
                        + `--from-literal=azurestorageaccountname=${this.azureStorageAccountName} `
                        + `--from-literal=azurestorageaccountkey=${storageAccountKey}`)
364
365
                    }catch(error) {
                        this.log.error(error);
SparkSnail's avatar
SparkSnail committed
366
367
                        throw new Error(error);
                    }
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
                } else if(kubeflowClusterConfigBase && (kubeflowClusterConfigBase.storage === 'nfs' || kubeflowClusterConfigBase.storage === undefined)) {
                    //Check and mount NFS mount point here
                    //If storage is undefined, the default value is nfs
                    this.kubeflowClusterConfig = new KubeflowClusterConfigNFS(kubeflowClusterJsonObject.operator, kubeflowClusterJsonObject.nfs, 
                        kubeflowClusterJsonObject.storage)
                    let nfsKubeflowClusterConfig = <KubeflowClusterConfigNFS>this.kubeflowClusterConfig;
                    await cpp.exec(`mkdir -p ${this.trialLocalNFSTempFolder}`);
                    const nfsServer: string = nfsKubeflowClusterConfig.nfs.server;
                    const nfsPath: string = nfsKubeflowClusterConfig.nfs.path;

                    try {
                        await cpp.exec(`sudo mount ${nfsServer}:${nfsPath} ${this.trialLocalNFSTempFolder}`);
                    } catch(error) {
                        const mountError: string = `Mount NFS ${nfsServer}:${nfsPath} to ${this.trialLocalNFSTempFolder} failed, error is ${error}`;
                        this.log.error(mountError);
                        throw new Error(mountError);
                    }
                } else {
                    const error: string = `kubeflowClusterConfig format error!`;
                    this.log.error(error);
                    throw new Error(error);
389
390
391
                }

                this.kubeflowJobPlural = kubeflowOperatorMap.get(this.kubeflowClusterConfig.operator);
392
                this.kubeflowJobKind = kubeflowOperatorJobKindMap.get(this.kubeflowClusterConfig.operator)
393
394
395
396
397
398
399
400
                break;

            case TrialConfigMetadataKey.TRIAL_CONFIG:
                if (!this.kubeflowClusterConfig){
                    this.log.error('kubeflow cluster config is not initialized');
                    return Promise.reject(new Error('kubeflow cluster config is not initialized'));                    
                }

401
402
403
404
405
406
407
408
409
410
411
412
413
414
                assert(this.kubeflowClusterConfig !== undefined)
                let kubeflowTrialJsonObjsect = JSON.parse(value);
                if(this.kubeflowClusterConfig.operator === 'tf-operator'){
                    this.kubeflowTrialConfig = new KubeflowTrialConfigTensorflow(kubeflowTrialJsonObjsect.codeDir, 
                        kubeflowTrialJsonObjsect.worker, kubeflowTrialJsonObjsect.ps);
                }else if(this.kubeflowClusterConfig.operator === 'pytorch-operator'){
                    this.kubeflowTrialConfig = new KubeflowTrialConfigPytorch(kubeflowTrialJsonObjsect.codeDir, 
                        kubeflowTrialJsonObjsect.worker, kubeflowTrialJsonObjsect.master);
                }

                if (!this.kubeflowTrialConfig){
                    this.log.error('kubeflow kubeflow TrialConfig is not initialized');
                    return Promise.reject(new Error('kubeflow kubeflow TrialConfig is not initialized'));                    
                }
415
416
417
418
419
420
421
422

                // Validate to make sure codeDir doesn't have too many files
                try {
                    await validateCodeDir(this.kubeflowTrialConfig.codeDir);
                } catch(error) {
                    this.log.error(error);
                    return Promise.reject(new Error(error));                    
                }
423
424
425
426
427
428
429
430
431
                break;
            default:
                break;
        }

        return Promise.resolve();
    }

    public getClusterMetadata(key: string): Promise<string> {
432
        return Promise.resolve('');
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
    }

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

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

        assert(this.kubeflowJobPlural !== undefined);
        
        // Delete all kubeflow jobs whose expId label is current experiment id 
        try {
            await cpp.exec(`kubectl delete ${this.kubeflowJobPlural} -l app=${this.NNI_KUBEFLOW_TRIAL_LABEL},expId=${getExperimentId()}`);
        } catch(error) {
            this.log.error(`Delete ${this.kubeflowJobPlural} with label: app=${this.NNI_KUBEFLOW_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 Kubeflow rest server 
        const restServer: KubeflowJobRestServer = component.get(KubeflowJobRestServer);
        try {
            await restServer.stop();
            this.log.info('Kubeflow Training service rest server stopped successfully.');
        } catch (error) {
            this.log.error(`Kubeflow Training service rest server stopped failed, error: ${error.message}`);
            Promise.reject(error);
        }

        return Promise.resolve();
    }

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

481
482
483
484
485
486
    /**
     * Generate kubeflow resource config file
     * @param trialJobId trial job id
     * @param trialWorkingFolder working folder
     * @param kubeflowJobName job name
     * @param workerPodResources worker pod template
487
     * @param nonWorkerPodResources non-worker pod template, like ps or master
488
     */
489
    private generateKubeflowJobConfig(trialJobId: string, trialWorkingFolder: string, kubeflowJobName : string, workerPodResources : any, nonWorkerPodResources?: any) : any {
490
491
492
493
494
495
496
497
        if(!this.kubeflowClusterConfig) {
            throw new Error('Kubeflow Cluster config is not initialized');
        }

        if(!this.kubeflowTrialConfig) {
            throw new Error('Kubeflow trial config is not initialized');
        }

498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
        if(!this.kubeflowJobPlural) {
            throw new Error('Kubeflow job plural name is undefined');
        }
        const replicaSpecsObj: any = {};
        let replicaSpecsObjMap = new Map<string, object>();

        if(this.kubeflowClusterConfig.operator === 'tf-operator') {
            let tensorflowTrialConfig: KubeflowTrialConfigTensorflow = <KubeflowTrialConfigTensorflow>this.kubeflowTrialConfig;
            replicaSpecsObj.Worker = this.generateReplicaConfig(trialWorkingFolder, tensorflowTrialConfig.worker.replicas, 
                tensorflowTrialConfig.worker.image, 'run_worker.sh', workerPodResources);
            if (tensorflowTrialConfig.ps){
                replicaSpecsObj.Ps = this.generateReplicaConfig(trialWorkingFolder, tensorflowTrialConfig.ps.replicas, 
                    tensorflowTrialConfig.ps.image, 'run_ps.sh', nonWorkerPodResources);
            }
            replicaSpecsObjMap.set(this.kubeflowJobPlural, {'tfReplicaSpecs': replicaSpecsObj})
        }
        else if(this.kubeflowClusterConfig.operator === 'pytorch-operator') {
            let pytorchTrialConfig: KubeflowTrialConfigPytorch = <KubeflowTrialConfigPytorch>this.kubeflowTrialConfig;
            replicaSpecsObj.Worker = this.generateReplicaConfig(trialWorkingFolder, pytorchTrialConfig.worker.replicas, 
                pytorchTrialConfig.worker.image, 'run_worker.sh', workerPodResources);
            if(pytorchTrialConfig.master){
                replicaSpecsObj.Master = this.generateReplicaConfig(trialWorkingFolder, pytorchTrialConfig.master.replicas, 
                    pytorchTrialConfig.master.image, 'run_master.sh', nonWorkerPodResources);
            }
            replicaSpecsObjMap.set(this.kubeflowJobPlural, {'pytorchReplicaSpecs': replicaSpecsObj})
523
524
        }

525
526
        return {
            apiVersion: 'kubeflow.org/v1alpha2',
527
            kind: this.kubeflowJobKind,
528
529
530
531
532
533
534
535
536
            metadata: { 
                name: kubeflowJobName,
                namespace: 'default',
                labels: {
                    app: this.NNI_KUBEFLOW_TRIAL_LABEL,
                    expId: getExperimentId(),
                    trialId: trialJobId
                }
            },
537
538
            spec: replicaSpecsObjMap.get(this.kubeflowJobPlural)
        };     
539
540
    }

541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
    /**
     * Generate tf-operator's tfjobs replica config section
     * @param trialWorkingFolder trial working folder
     * @param replicaNumber replica number
     * @param replicaImage image
     * @param runScriptFile script file name
     * @param podResources pod resource config section
     */
    private generateReplicaConfig(trialWorkingFolder: string, replicaNumber: number, replicaImage: string, runScriptFile: string, podResources: any): any {
        if(!this.kubeflowClusterConfig) {
            throw new Error('Kubeflow Cluster config is not initialized');
        }

        if(!this.kubeflowTrialConfig) {
            throw new Error('Kubeflow trial config is not initialized');
        }

558
559
560
561
        if(!this.kubeflowJobPlural) {
            throw new Error('Kubeflow job plural name is undefined');
        }

SparkSnail's avatar
SparkSnail committed
562
        let volumeSpecMap = new Map<string, object>();
563
        if(this.kubeflowClusterConfig.storage && this.kubeflowClusterConfig.storage === 'azureStorage'){
SparkSnail's avatar
SparkSnail committed
564
565
            volumeSpecMap.set('nniVolumes', [
            {
566
567
568
569
570
571
                    name: 'nni-vol',
                    azureFile: {
                        secretName: `${this.azureStorageSecretName}`,
                        shareName: `${this.azureStorageShare}`,
                        readonly: false
                    }
SparkSnail's avatar
SparkSnail committed
572
            }])
573
574
        }else {
            let nfsKubeflowClusterConfig: KubeflowClusterConfigNFS = <KubeflowClusterConfigNFS> this.kubeflowClusterConfig;
SparkSnail's avatar
SparkSnail committed
575
576
577
            volumeSpecMap.set('nniVolumes', [
            {
                name: 'nni-vol',
578
579
580
                nfs: {
                    server: `${nfsKubeflowClusterConfig.nfs.server}`,
                    path: `${nfsKubeflowClusterConfig.nfs.path}`
SparkSnail's avatar
SparkSnail committed
581
582
                }
            }])
583
584
585
586
587
588
589
        }

        let containerNameMap = new Map<string, string>();
        if(this.kubeflowJobPlural == 'tfjobs'){
            containerNameMap.set(this.kubeflowJobPlural, 'tensorflow');
        }else if(this.kubeflowJobPlural == 'pytorchjobs'){
            containerNameMap.set(this.kubeflowJobPlural, 'pytorch');
SparkSnail's avatar
SparkSnail committed
590
591
        }

592
593
594
595
596
597
598
599
600
601
602
        return {
            replicas: replicaNumber,
            template: {
                metadata: {
                    creationTimestamp: null
                },
                spec: {
                    containers: [
                    {
                        // Kubeflow tensorflow operator requires that containers' name must be tensorflow
                        // TODO: change the name based on operator's type
603
                        name: containerNameMap.get(this.kubeflowJobPlural),
604
605
                        image: replicaImage,
                        args: ["sh", `${path.join(trialWorkingFolder, runScriptFile)}`],
SparkSnail's avatar
SparkSnail committed
606
607
608
                        volumeMounts: [
                        {
                            name: 'nni-vol',
609
610
611
612
613
                            mountPath: this.CONTAINER_MOUNT_PATH
                        }],
                        resources: podResources
                    }],
                    restartPolicy: 'ExitCode',
SparkSnail's avatar
SparkSnail committed
614
                    volumes: volumeSpecMap.get('nniVolumes')
615
616
617
618
619
620
621
622
623
624
625
626
                }
            }
        };
    }

    /**
     * Genereate run script for different roles(like worker or ps)
     * @param trialJobId trial job id
     * @param trialWorkingFolder working folder
     * @param command 
     * @param trialSequenceId sequence id
     */
627
628
    private generateRunScript(trialJobId: string, trialWorkingFolder: string, 
                command: string, trialSequenceId: string, roleType: DistTrainRole, gpuNum: number): string {
629
630
631
632
633
634
635
636
637
638
639
        const runScriptLines: string[] = [];

        runScriptLines.push('#!/bin/bash');
        runScriptLines.push('export NNI_PLATFORM=kubeflow');
        runScriptLines.push(`export NNI_SYS_DIR=$PWD/nni/${trialJobId}`);
        runScriptLines.push(`export NNI_OUTPUT_DIR=${path.join(trialWorkingFolder, 'output', `${roleType}_output`)}`);
        runScriptLines.push('export MULTI_PHASE=false');
        runScriptLines.push(`export NNI_TRIAL_JOB_ID=${trialJobId}`);
        runScriptLines.push(`export NNI_EXP_ID=${getExperimentId()}`);
        runScriptLines.push(`export NNI_CODE_DIR=${trialWorkingFolder}`);
        runScriptLines.push(`export NNI_TRIAL_SEQ_ID=${trialSequenceId}`);
640
        
641
642
643
        // 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
644
645
        if(gpuNum === 0) {
            runScriptLines.push(`export CUDA_VISIBLE_DEVICES=''`);
646
        }
647

648
        const nniManagerIp = this.nniManagerIpConfig?this.nniManagerIpConfig.nniManagerIp:getIPV4Address();
649
650
651
652
653
        runScriptLines.push('mkdir -p $NNI_SYS_DIR');
        runScriptLines.push('mkdir -p $NNI_OUTPUT_DIR');
        runScriptLines.push('cp -rT $NNI_CODE_DIR $NNI_SYS_DIR');
        runScriptLines.push('cd $NNI_SYS_DIR');
        runScriptLines.push('sh install_nni.sh # Check and install NNI pkg');
SparkSnail's avatar
SparkSnail committed
654
655
656
        runScriptLines.push(`python3 -m nni_trial_tool.trial_keeper --trial_command '${command}' `
        + `--nnimanager_ip '${nniManagerIp}' --nnimanager_port '${this.kubeflowRestServerPort}' `
        + `1>$NNI_OUTPUT_DIR/trialkeeper_stdout 2>$NNI_OUTPUT_DIR/trialkeeper_stderr`);
657
658

        return runScriptLines.join('\n');
659
660
    }

661
662
663
664
665
666
667
668
669
    private generateSequenceId(): number {
        if (this.nextTrialSequenceId === -1) {
            this.nextTrialSequenceId = getInitTrialSequenceId();
        }

        return this.nextTrialSequenceId++;
    }
}

QuanluZhang's avatar
QuanluZhang committed
670
export { KubeflowTrainingService }