Kubeflow TrainingService support, v1

src/nni_manager/main.ts

@@ -36,8 +36,8 @@ import { LocalTrainingServiceForGPU } from './training_service/local/localTraini
 import {
     RemoteMachineTrainingService
 } from './training_service/remote_machine/remoteMachineTrainingService';
-import { PAITrainingService } from './training_service/pai/paiTrainingService'
-
+import { PAITrainingService } from './training_service/pai/paiTrainingService';
+import { KubeflowTrainingService } from './training_service/kubeflow/kubeflowTrainingService';
 
 function initStartupInfo(startExpMode: string, resumeExperimentId: string, basePort: number) {
     const createNew: boolean = (startExpMode === 'new');
@@ -52,6 +52,8 @@ async function initContainer(platformMode: string): Promise<void> {
         Container.bind(TrainingService).to(RemoteMachineTrainingService).scope(Scope.Singleton);
     } else if (platformMode === 'pai') {
         Container.bind(TrainingService).to(PAITrainingService).scope(Scope.Singleton);
+    } else if (platformMode === 'kubeflow') {
+        Container.bind(TrainingService).to(KubeflowTrainingService).scope(Scope.Singleton);
     } else {
         throw new Error(`Error: unsupported mode: ${mode}`);
     }
@@ -76,19 +78,22 @@ if (!strPort || strPort.length === 0) {
 const port: number = parseInt(strPort, 10);
 
 const mode: string = parseArg(['--mode', '-m']);
-if (!['local', 'remote', 'pai'].includes(mode)) {
+if (!['local', 'remote', 'pai', 'kubeflow'].includes(mode)) {
+    console.log(`FATAL: unknown mode: ${mode}`);
     usage();
     process.exit(1);
 }
 
 const startMode: string = parseArg(['--start_mode', '-s']);
 if (!['new', 'resume'].includes(startMode)) {
+    console.log(`FATAL: unknown start_mode: ${startMode}`);
     usage();
     process.exit(1);
 }
 
 const experimentId: string = parseArg(['--experiment_id', '-id']);
 if (startMode === 'resume' && experimentId.trim().length < 1) {
+    console.log(`FATAL: cannot resume experiment, invalid experiment_id: ${experimentId}`);
     usage();
     process.exit(1);
 }

src/nni_manager/package.json

@@ -15,6 +15,7 @@
     "express": "^4.16.3",
     "express-joi-validator": "^2.0.0",
     "node-nvidia-smi": "^1.0.0",
+    "node-yaml": "^3.1.1",
     "rx": "^4.1.0",
     "sqlite3": "^4.0.2",
     "ssh2": "^0.6.1",

src/nni_manager/rest_server/restValidationSchemas.ts

@@ -46,6 +46,14 @@ export namespace ValidationSchemas {
                 userName: joi.string().min(1).required(),
                 passWord: joi.string().min(1).required(),
                 host: joi.string().min(1).required()
+            }),
+            kubeflow_config: joi.object({
+                operator: joi.string().min(1).required(),
+                nfs: joi.object({
+                    server: joi.string().min(1).required(),
+                    path: joi.string().min(1).required()
+                }).required(),
+                kubernetesServer: joi.string().min(1).required()
             })
         }
     };

src/nni_manager/training_service/common/clusterJobRestServer.ts

@@ -0,0 +1,96 @@
+/**
+ * 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 { Request, Response, Router } from 'express';
+import * as bodyParser from 'body-parser';
+import * as component from '../../common/component';
+import { getBasePort, getExperimentId } from '../../common/experimentStartupInfo';
+import { RestServer } from '../../common/restServer'
+
+/**
+ * Cluster Job Training service Rest server, provides rest API to support Cluster job metrics update
+ * 
+ */
+@component.Singleton
+export abstract class ClusterJobRestServer extends RestServer{
+    private readonly API_ROOT_URL: string = '/api/v1/nni-pai';
+
+    private readonly expId: string = getExperimentId();
+
+    /**
+     * constructor to provide NNIRestServer's own rest property, e.g. port
+     */
+    constructor() {
+        super();
+        const basePort: number = getBasePort();
+        assert(basePort && basePort > 1024);
+
+        this.port = basePort + 1;         
+    }
+
+    public get clusterRestServerPort(): number {
+        if(!this.port) {
+            throw new Error('PAI Rest server port is undefined');
+        }
+        return this.port;
+    }
+
+    /**
+     * NNIRestServer's own router registration
+     */
+    protected registerRestHandler(): void {
+        this.app.use(bodyParser.json());
+        this.app.use(this.API_ROOT_URL, this.createRestHandler());
+    }
+
+    private createRestHandler() : Router {
+        const router: Router = Router();
+
+        // tslint:disable-next-line:typedef
+        router.use((req: Request, res: Response, next) => {
+            this.log.info(`${req.method}: ${req.url}: body:\n${JSON.stringify(req.body, undefined, 4)}`);
+            res.setHeader('Content-Type', 'application/json');
+            next();
+        });
+
+        router.post(`/update-metrics/${this.expId}/:trialId`, (req: Request, res: Response) => {
+            try {
+                this.log.info(`Get update-metrics request, trial job id is ${req.params.trialId}`);
+                this.log.info(`update-metrics body is ${JSON.stringify(req.body)}`);
+
+                this.handleTrialMetrics(req.body.jobId, req.body.metrics);
+
+                res.send();
+            }
+            catch(err) {
+                this.log.error(`json parse metrics error: ${err}`);
+                res.status(500);
+                res.send(err.message);
+            }
+        });
+
+        return router;
+    }
+
+    /** Abstract method to handle trial metrics data */
+    protected abstract handleTrialMetrics(jobId : string, trialMetrics : any[]) : void;
+}

src/nni_manager/training_service/common/containerJobData.ts

@@ -0,0 +1,30 @@
+/**
+ * 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';
+
+export const CONTAINER_INSTALL_NNI_SHELL_FORMAT: string = 
+`#!/bin/bash
+if python3 -c 'import nni' > /dev/null 2>&1; then
+  # nni module is already installed, skip
+  return
+else
+  # Install nni
+  python3 -m pip install --user --upgrade nni  
+fi`;

src/nni_manager/training_service/common/trialConfigMetadataKey.ts

@@ -28,5 +28,6 @@ export enum TrialConfigMetadataKey {
     EXPERIMENT_ID = 'experimentId',
     MULTI_PHASE = 'multiPhase',
     RANDOM_SCHEDULER = 'random_scheduler',
-    PAI_CLUSTER_CONFIG = 'pai_config'
+    PAI_CLUSTER_CONFIG = 'pai_config',
+    KUBEFLOW_CLUSTER_CONFIG = 'kubeflow_config'
 }

src/nni_manager/training_service/kubeflow/kubeflowConfig.ts

@@ -0,0 +1,93 @@
+import { TrialConfig } from "../common/trialConfig";
+
+/**
+ * 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';
+
+
+/** operator types that kubeflow supported */
+export type KubeflowOperator = 'tf-operator' | 'pytorch-operator' | 'mxnet-operator' | 'caffe2-operator' | 'chainer-operator' | 'mpi-operator';
+export type KubeflowOperatorPlural = 'tfjobs' | 'pytorchjobs' | 'mxjobs' | 'caffe2jobs' | 'chainerjobs' | 'mpijobs';
+
+/**
+ * map from Kubeflow operator name to its plural name in K8S
+ */
+export const kubeflowOperatorMap : Map<KubeflowOperator, KubeflowOperatorPlural> =  new Map<KubeflowOperator, KubeflowOperatorPlural>([
+    ['tf-operator' , 'tfjobs'],
+    ['pytorch-operator', 'pytorchjobs'],
+    ['mxnet-operator', 'mxjobs'],
+    ['caffe2-operator', 'caffe2jobs'],
+    ['chainer-operator', 'chainerjobs'],
+    ['mpi-operator', 'mpijobs']    
+]);
+
+/**
+ * Kuberflow cluster configuration
+ * 
+ */
+export class KubeflowClusterConfig {
+    /** Name of Kubeflow operator, like tf-operator */
+    public readonly operator: KubeflowOperator;
+    public readonly nfs: NFSConfig;
+    public readonly kubernetesServer: string;
+
+    /**
+     * Constructor
+     * @param userName User name of Kubeflow Cluster
+     * @param passWord password of Kubeflow Cluster
+     * @param host Host IP of Kubeflow Cluster
+     */
+    constructor(operator: KubeflowOperator, nfs : NFSConfig, kubernetesServer : string) {
+        this.operator = operator;
+        this.nfs = nfs;
+        this.kubernetesServer = kubernetesServer;
+    }
+}
+
+/**
+ * NFS configuration to store Kubeflow job related files
+ */
+export class NFSConfig {
+    /** IP Adress of NFS server */
+    public readonly server : string;
+    /** exported NFS path on NFS server */
+    public readonly path : string;
+
+    constructor(server : string, path : string) {
+        this.server = server;
+        this.path = path;
+    }
+}
+
+/**
+ * Trial job configuration for Kubeflow
+ */
+export class KubeflowTrialConfig extends TrialConfig {
+    public readonly cpuNum: number;
+    public readonly memoryMB: number;
+    public readonly image: string;
+
+    constructor(command : string, codeDir : string, gpuNum : number, cpuNum: number, memoryMB: number, image: string) {
+        super(command, codeDir, gpuNum);
+        this.cpuNum = cpuNum;
+        this.memoryMB = memoryMB;
+        this.image = image;
+    }
+}

src/nni_manager/training_service/kubeflow/kubeflowData.ts

@@ -0,0 +1,78 @@
+/**
+ * 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 { JobApplicationForm, TrialJobDetail, TrialJobStatus  } from '../../common/trainingService';
+
+/**
+ * KubeflowTrialJobDetail
+ */
+// tslint:disable-next-line:max-classes-per-file
+export class KubeflowTrialJobDetail implements TrialJobDetail {
+    public id: string;
+    public status: TrialJobStatus;
+    public submitTime: number;
+    public startTime?: number;
+    public endTime?: number;
+    public tags?: string[];
+    public url?: string;
+    public workingDirectory: string;
+    public form: JobApplicationForm;
+    public kubeflowJobName: string;
+    public sequenceId: number;
+    public queryJobFailedCount: number;
+    public k8sPluralName: string
+
+    constructor(id: string, status: TrialJobStatus, submitTime: number,
+                workingDirectory: string, form: JobApplicationForm, 
+                kubeflowJobName: string, sequenceId: number, url: string, k8sPluralName: string) {
+        this.id = id;
+        this.status = status;
+        this.submitTime = submitTime;
+        this.workingDirectory = workingDirectory;
+        this.form = form;
+        this.kubeflowJobName = kubeflowJobName;
+        this.sequenceId = sequenceId;
+        this.tags = [];
+        this.queryJobFailedCount = 0;
+        this.url = url;
+        this.k8sPluralName = k8sPluralName;
+    }
+}
+
+export const KUBEFLOW_RUN_SHELL_FORMAT: string = 
+`#!/bin/bash
+export NNI_PLATFORM=kubeflow
+export NNI_SYS_DIR={0} 
+export NNI_OUTPUT_DIR={1} 
+export MULTI_PHASE=false
+export NNI_TRIAL_JOB_ID={2} 
+export NNI_EXP_ID={3} 
+export NNI_CODE_DIR={4} 
+export NNI_TRIAL_SEQ_ID={5}
+mkdir -p $NNI_SYS_DIR
+mkdir -p $NNI_OUTPUT_DIR
+cp -rT $NNI_CODE_DIR $NNI_SYS_DIR
+cd $NNI_SYS_DIR
+sh install_nni.sh # Check and install NNI pkg
+python3 -m nni_trial_tool.trial_keeper --trial_command '{6}' --nnimanager_ip '{7}' --nnimanager_port '{8}' 1>$NNI_OUTPUT_DIR/trialkeeper_stdout 2>$NNI_OUTPUT_DIR//trialkeeper_stderr
+`
+
+export type KubeflowTFJobType = 'Created' | 'Running' | 'Failed' | 'Succeeded';