Support pai and paiYarn trainingservice (#1853)

docs/en_US/TrainingService/PaiMode.md

@@ -34,10 +34,13 @@ trial:
   cpuNum: 1
   memoryMB: 8196
   image: msranni/nni:latest
+  virtualCluster: default
+  nniManagerNFSMountPath: /home/user/mnt
+  containerNFSMountPath: /mnt/data/user
 # Configuration to access OpenPAI Cluster
 paiConfig:
   userName: your_pai_nni_user
-  passWord: your_pai_password
+  token: your_pai_token
   host: 10.1.1.1
 ```
 
@@ -53,52 +56,13 @@ Compared with [LocalMode](LocalMode.md) and [RemoteMachineMode](RemoteMachineMod
     * We already build a docker image [nnimsra/nni](https://hub.docker.com/r/msranni/nni/) on [Docker Hub](https://hub.docker.com/). It contains NNI python packages, Node modules and javascript artifact files required to start experiment, and all of NNI dependencies. The docker file used to build this image can be found at [here](https://github.com/Microsoft/nni/tree/master/deployment/docker/Dockerfile). You can either use this image directly in your config file, or build your own image based on it.
 * virtualCluster
     * Optional key. Set the virtualCluster of OpenPAI. If omitted, the job will run on default virtual cluster.
-* shmMB
-    * Optional key. Set the shmMB configuration of OpenPAI, it set the shared memory for one task in the task role.
-* authFile
-    * Optional key, Set the auth file path for private registry while using PAI mode, [Refer](https://github.com/microsoft/pai/blob/2ea69b45faa018662bc164ed7733f6fdbb4c42b3/docs/faq.md#q-how-to-use-private-docker-registry-job-image-when-submitting-an-openpai-job), you can prepare the authFile and simply provide the local path of this file, NNI will upload this file to HDFS for you.
-* portList  
-    * Optional key. Set the portList configuration of OpenPAI, it specifies a list of port used in container, [Refer](https://github.com/microsoft/pai/blob/b2324866d0280a2d22958717ea6025740f71b9f0/docs/job_tutorial.md#specification).  
-    The config schema in NNI is shown below:
-    ```
-    portList:
-      - label: test
-        beginAt: 8080
-        portNumber: 2
-    ``` 
-    Let's say you want to launch a tensorboard in the mnist example using the port. So the first step is to write a wrapper script `launch_pai.sh` of `mnist.py`.
-
-    ```bash
-    export TENSORBOARD_PORT=PAI_PORT_LIST_${PAI_CURRENT_TASK_ROLE_NAME}_0_tensorboard
-    tensorboard --logdir . --port ${!TENSORBOARD_PORT} &
-    python3 mnist.py
-    ```
-    The config file of portList should be filled as following:
+* nniManagerNFSMountPath
+    * Required key. Set the mount path in your nniManager machine.
+* containerNFSMountPath
+    * Required key. Set the mount path in your container used in PAI.
+* paiStoragePlugin
+    * Required key. Set the storage plugin name used in PAI.
 
-    ```yaml
-  trial:
-    command: bash launch_pai.sh
-    portList:
-      - label: tensorboard
-        beginAt: 0
-        portNumber: 1
-    ```
-
-NNI support two kind of authorization method in PAI, including password and PAI token, [refer](https://github.com/microsoft/pai/blob/b6bd2ab1c8890f91b7ac5859743274d2aa923c22/docs/rest-server/API.md#2-authentication). The authorization is configured in `paiConfig` field.  
-For password authorization, the `paiConfig` schema is:
-```
-paiConfig:
-  userName: your_pai_nni_user
-  passWord: your_pai_password
-  host: 10.1.1.1
-```  
-For pai token authorization, the `paiConfig` schema is:
-```
-paiConfig:
-  userName: your_pai_nni_user
-  token: your_pai_token
-  host: 10.1.1.1
-```
 
 Once complete to fill NNI experiment config file and save (for example, save as exp_pai.yml), then run the following command
 ```
@@ -121,9 +85,7 @@ And you will be redirected to HDFS web portal to browse the output files of that
 You can see there're three fils in output folder: stderr, stdout, and trial.log
 
 ## data management
-If your training data is not too large, it could be put into codeDir, and nni will upload the data to hdfs, or you could build your own docker image with the data. If you have large dataset, it's not appropriate to put the data in codeDir, and you could follow the [guidance](https://github.com/microsoft/pai/blob/master/docs/user/storage.md) to mount the data folder in container.
-
-If you also want to save trial's other output into HDFS, like model files, you can use environment variable `NNI_OUTPUT_DIR` in your trial code to save your own output files, and NNI SDK will copy all the files in `NNI_OUTPUT_DIR` from trial's container to HDFS, the target path is `hdfs://host:port/{username}/nni/{experiments}/{experimentId}/trials/{trialId}/nnioutput`
+Befour using NNI to start your experiment, users should set the corresponding mount data path in your nniManager machine. PAI has their own storage(NFS, AzureBlob ...), and the storage will used in PAI will be mounted to the container when it start a job. Users should set the PAI storage type by `paiStoragePlugin` field to choose a storage in PAI. Then users should mount the storage to their nniManager machine, and set the `nniManagerNFSMountPath` field in configuration file, NNI will generate bash files and copy data in `codeDir` to the `nniManagerNFSMountPath` folder, then NNI will start a trial job. The data in `nniManagerNFSMountPath` will be sync to PAI storage, and will be mounted to PAI's container. The data path in container is set in `containerNFSMountPath`, NNI will enter this folder first, and then run scripts to start a trial job. 
 
 ## version check
 NNI support version check feature in since version 0.6. It is a policy to insure the version of NNIManager is consistent with trialKeeper, and avoid errors caused by version incompatibility.

docs/en_US/TrainingService/PaiYarnMode.md

+**Run an Experiment on OpenpaiYarn**
+===
+The original `pai` mode is modificated to `paiYarn` mode, which is a distributed training platform based on Yarn.
+
+## Setup environment
+Install NNI, follow the install guide [here](../Tutorial/QuickStart.md).
+
+## Run an experiment
+Use `examples/trials/mnist-annotation` as an example. The NNI config YAML file's content is like:
+
+```yaml
+authorName: your_name
+experimentName: auto_mnist
+# how many trials could be concurrently running
+trialConcurrency: 2
+# maximum experiment running duration
+maxExecDuration: 3h
+# empty means never stop
+maxTrialNum: 100
+# choice: local, remote, pai, paiYarn
+trainingServicePlatform: paiYarn
+# search space file
+searchSpacePath: search_space.json
+# choice: true, false
+useAnnotation: true
+tuner:
+  builtinTunerName: TPE
+  classArgs:
+    optimize_mode: maximize
+trial:
+  command: python3 mnist.py
+  codeDir: ~/nni/examples/trials/mnist-annotation
+  gpuNum: 0
+  cpuNum: 1
+  memoryMB: 8196
+  image: msranni/nni:latest
+# Configuration to access OpenpaiYarn Cluster
+paiYarnConfig:
+  userName: your_paiYarn_nni_user
+  passWord: your_paiYarn_password
+  host: 10.1.1.1
+```
+
+Note: You should set `trainingServicePlatform: paiYarn` in NNI config YAML file if you want to start experiment in paiYarn mode.
+
+Compared with [LocalMode](LocalMode.md) and [RemoteMachineMode](RemoteMachineMode.md), trial configuration in paiYarn mode have these additional keys:
+* cpuNum
+    * Required key. Should be positive number based on your trial program's CPU  requirement
+* memoryMB
+    * Required key. Should be positive number based on your trial program's memory requirement
+* image
+    * Required key. In paiYarn mode, your trial program will be scheduled by OpenpaiYarn to run in [Docker container](https://www.docker.com/). This key is used to specify the Docker image used to create the container in which your trial will run.
+    * We already build a docker image [nnimsra/nni](https://hub.docker.com/r/msranni/nni/) on [Docker Hub](https://hub.docker.com/). It contains NNI python packages, Node modules and javascript artifact files required to start experiment, and all of NNI dependencies. The docker file used to build this image can be found at [here](https://github.com/Microsoft/nni/tree/master/deployment/docker/Dockerfile). You can either use this image directly in your config file, or build your own image based on it.
+* virtualCluster
+    * Optional key. Set the virtualCluster of OpenpaiYarn. If omitted, the job will run on default virtual cluster.
+* shmMB
+    * Optional key. Set the shmMB configuration of OpenpaiYarn, it set the shared memory for one task in the task role.
+* authFile
+    * Optional key, Set the auth file path for private registry while using paiYarn mode, [Refer](https://github.com/microsoft/paiYarn/blob/2ea69b45faa018662bc164ed7733f6fdbb4c42b3/docs/faq.md#q-how-to-use-private-docker-registry-job-image-when-submitting-an-openpaiYarn-job), you can prepare the authFile and simply provide the local path of this file, NNI will upload this file to HDFS for you.
+* portList  
+    * Optional key. Set the portList configuration of OpenpaiYarn, it specifies a list of port used in container, [Refer](https://github.com/microsoft/paiYarn/blob/b2324866d0280a2d22958717ea6025740f71b9f0/docs/job_tutorial.md#specification).  
+    The config schema in NNI is shown below:
+    ```
+    portList:
+      - label: test
+        beginAt: 8080
+        portNumber: 2
+    ``` 
+    Let's say you want to launch a tensorboard in the mnist example using the port. So the first step is to write a wrapper script `launch_paiYarn.sh` of `mnist.py`.
+
+    ```bash
+    export TENSORBOARD_PORT=paiYarn_PORT_LIST_${paiYarn_CURRENT_TASK_ROLE_NAME}_0_tensorboard
+    tensorboard --logdir . --port ${!TENSORBOARD_PORT} &
+    python3 mnist.py
+    ```
+    The config file of portList should be filled as following:
+
+    ```yaml
+  trial:
+    command: bash launch_paiYarn.sh
+    portList:
+      - label: tensorboard
+        beginAt: 0
+        portNumber: 1
+    ```
+
+NNI support two kind of authorization method in paiYarn, including password and paiYarn token, [refer](https://github.com/microsoft/paiYarn/blob/b6bd2ab1c8890f91b7ac5859743274d2aa923c22/docs/rest-server/API.md#2-authentication). The authorization is configured in `paiYarnConfig` field.  
+For password authorization, the `paiYarnConfig` schema is:
+```
+paiYarnConfig:
+  userName: your_paiYarn_nni_user
+  passWord: your_paiYarn_password
+  host: 10.1.1.1
+```  
+For paiYarn token authorization, the `paiYarnConfig` schema is:
+```
+paiYarnConfig:
+  userName: your_paiYarn_nni_user
+  token: your_paiYarn_token
+  host: 10.1.1.1
+```
+
+Once complete to fill NNI experiment config file and save (for example, save as exp_paiYarn.yml), then run the following command
+```
+nnictl create --config exp_paiYarn.yml
+```
+to start the experiment in paiYarn mode. NNI will create OpenpaiYarn job for each trial, and the job name format is something like `nni_exp_{experiment_id}_trial_{trial_id}`.
+You can see jobs created by NNI in the OpenpaiYarn cluster's web portal, like:
+![](../../img/nni_paiYarn_joblist.jpg)
+
+Notice: In paiYarn mode, NNIManager will start a rest server and listen on a port which is your NNI WebUI's port plus 1. For example, if your WebUI port is `8080`, the rest server will listen on `8081`, to receive metrics from trial job running in Kubernetes. So you should `enable 8081` TCP port in your firewall rule to allow incoming traffic.
+
+Once a trial job is completed, you can goto NNI WebUI's overview page (like http://localhost:8080/oview) to check trial's information.
+
+Expand a trial information in trial list view, click the logPath link like:
+![](../../img/nni_webui_joblist.jpg)
+
+And you will be redirected to HDFS web portal to browse the output files of that trial in HDFS:
+![](../../img/nni_trial_hdfs_output.jpg)
+
+You can see there're three fils in output folder: stderr, stdout, and trial.log
+
+## data management
+If your training data is not too large, it could be put into codeDir, and nni will upload the data to hdfs, or you could build your own docker image with the data. If you have large dataset, it's not appropriate to put the data in codeDir, and you could follow the [guidance](https://github.com/microsoft/paiYarn/blob/master/docs/user/storage.md) to mount the data folder in container.
+
+If you also want to save trial's other output into HDFS, like model files, you can use environment variable `NNI_OUTPUT_DIR` in your trial code to save your own output files, and NNI SDK will copy all the files in `NNI_OUTPUT_DIR` from trial's container to HDFS, the target path is `hdfs://host:port/{username}/nni/{experiments}/{experimentId}/trials/{trialId}/nnioutput`
+
+## version check
+NNI support version check feature in since version 0.6. It is a policy to insure the version of NNIManager is consistent with trialKeeper, and avoid errors caused by version incompatibility.
+Check policy:
+1. NNIManager before v0.6 could run any version of trialKeeper, trialKeeper support backward compatibility.
+2. Since version 0.6, NNIManager version should keep same with triakKeeper version. For example, if NNIManager version is 0.6, trialKeeper version should be 0.6 too.
+3. Note that the version check feature only check first two digits of version.For example, NNIManager v0.6.1 could use trialKeeper v0.6 or trialKeeper v0.6.2, but could not use trialKeeper v0.5.1 or trialKeeper v0.7.
+
+If you could not run your experiment and want to know if it is caused by version check, you could check your webUI, and there will be an error message about version check.
+![](../../img/version_check.png)

src/nni_manager/main.ts

@@ -20,7 +20,8 @@ import { NNIRestServer } from './rest_server/nniRestServer';
 import { FrameworkControllerTrainingService } from './training_service/kubernetes/frameworkcontroller/frameworkcontrollerTrainingService';
 import { KubeflowTrainingService } from './training_service/kubernetes/kubeflow/kubeflowTrainingService';
 import { LocalTrainingService } from './training_service/local/localTrainingService';
-import { PAITrainingService } from './training_service/pai/paiTrainingService';
+import { PAIK8STrainingService } from './training_service/pai/paiK8S/paiK8STrainingService';
+import { PAIYarnTrainingService } from './training_service/pai/paiYarn/paiYarnTrainingService';
 import {
     RemoteMachineTrainingService
 } from './training_service/remote_machine/remoteMachineTrainingService';
@@ -44,7 +45,11 @@ async function initContainer(platformMode: string, logFileName?: string): Promis
             .scope(Scope.Singleton);
     } else if (platformMode === 'pai') {
         Container.bind(TrainingService)
-            .to(PAITrainingService)
+            .to(PAIK8STrainingService)
+            .scope(Scope.Singleton);
+    } else if (platformMode === 'paiYarn') {
+            Container.bind(TrainingService)
+            .to(PAIYarnTrainingService)
             .scope(Scope.Singleton);
     } else if (platformMode === 'kubeflow') {
         Container.bind(TrainingService)
@@ -76,7 +81,7 @@ async function initContainer(platformMode: string, logFileName?: string): Promis
 
 function usage(): void {
     console.info('usage: node main.js --port <port> --mode \
-    <local/remote/pai/kubeflow/frameworkcontroller> --start_mode <new/resume> --experiment_id <id>');
+    <local/remote/pai/kubeflow/frameworkcontroller/paiYarn> --start_mode <new/resume> --experiment_id <id>');
 }
 
 const strPort: string = parseArg(['--port', '-p']);
@@ -88,7 +93,7 @@ if (!strPort || strPort.length === 0) {
 const port: number = parseInt(strPort, 10);
 
 const mode: string = parseArg(['--mode', '-m']);
-if (!['local', 'remote', 'pai', 'kubeflow', 'frameworkcontroller'].includes(mode)) {
+if (!['local', 'remote', 'pai', 'kubeflow', 'frameworkcontroller', 'paiYarn'].includes(mode)) {
     console.log(`FATAL: unknown mode: ${mode}`);
     usage();
     process.exit(1);

src/nni_manager/rest_server/restValidationSchemas.ts

@@ -36,6 +36,9 @@ export namespace ValidationSchemas {
                 virtualCluster: joi.string(),
                 shmMB: joi.number(),
                 authFile: joi.string(),
+                nniManagerNFSMountPath: joi.string().min(1),
+                containerNFSMountPath: joi.string().min(1),
+                paiStoragePlugin: joi.string().min(1),
                 nasMode: joi.string().valid('classic_mode', 'enas_mode', 'oneshot_mode', 'darts_mode'),
                 portList: joi.array().items(joi.object({
                     label: joi.string().required(),
@@ -89,12 +92,18 @@ export namespace ValidationSchemas {
                     })
                 })
             }),
-            pai_config: joi.object({ // eslint-disable-line @typescript-eslint/camelcase
+            pai_yarn_config: joi.object({ // eslint-disable-line @typescript-eslint/camelcase
                 userName: joi.string().min(1).required(),
                 passWord: joi.string().min(1),
                 token: joi.string().min(1),
                 host: joi.string().min(1).required()
             }),
+            pai_config: joi.object({ // eslint-disable-line @typescript-eslint/camelcase
+                userName: joi.string().min(1).required(),
+                passWord: joi.string().min(1),
+                token: joi.string().min(1),
+                host: joi.string().min(1).required(),
+            }),
             kubeflow_config: joi.object({ // eslint-disable-line @typescript-eslint/camelcase
                 operator: joi.string().min(1).required(),
                 storage: joi.string().min(1),

src/nni_manager/training_service/common/trialConfigMetadataKey.ts

@@ -13,6 +13,7 @@ export enum TrialConfigMetadataKey {
     EXPERIMENT_ID = 'experimentId',
     MULTI_PHASE = 'multiPhase',
     RANDOM_SCHEDULER = 'random_scheduler',
+    PAI_YARN_CLUSTER_CONFIG = 'pai_yarn_config',
     PAI_CLUSTER_CONFIG = 'pai_config',
     KUBEFLOW_CLUSTER_CONFIG = 'kubeflow_config',
     NNI_MANAGER_IP = 'nni_manager_ip',

src/nni_manager/training_service/pai/paiConfig.ts

@@ -4,91 +4,8 @@
 'use strict';
 
 import {TrialConfig} from '../common/trialConfig';
+import { TrialJobApplicationForm, TrialJobDetail, TrialJobStatus  } from '../../common/trainingService';
 
-/**
- * Task role for PAI
- */
-export class PAITaskRole {
-    // Name for the task role
-    public readonly name: string;
-    // Number of tasks for the task role, no less than 1
-    public readonly taskNumber: number;
-    // CPU number for one task in the task role, no less than 1
-    public readonly cpuNumber: number;
-    // Memory for one task in the task role, no less than 100
-    public readonly memoryMB: number;
-    // GPU number for one task in the task role, no less than 0
-    public readonly gpuNumber: number;
-    // Executable command for tasks in the task role, can not be empty
-    public readonly command: string;
-    //Shared memory for one task in the task role
-    public readonly shmMB?: number;
-    //portList to specify the port used in container
-    public portList?: PortListMetaData[];
-
-    /**
-     * Constructor
-     * @param name Name for the task role
-     * @param taskNumber Number of tasks for the task role, no less than 1
-     * @param cpuNumber CPU number for one task in the task role, no less than 1
-     * @param memoryMB Memory for one task in the task role, no less than 100
-     * @param gpuNumber GPU number for one task in the task role, no less than 0
-     * @param command Executable command for tasks in the task role, can not be empty
-     */
-    constructor(name: string, taskNumber: number, cpuNumber: number, memoryMB: number, gpuNumber: number,
-                command: string, shmMB?: number, portList?: PortListMetaData[]) {
-        this.name = name;
-        this.taskNumber = taskNumber;
-        this.cpuNumber = cpuNumber;
-        this.memoryMB = memoryMB;
-        this.gpuNumber = gpuNumber;
-        this.command = command;
-        this.shmMB = shmMB;
-        this.portList = portList;
-    }
-}
-
-/**
- * Trial job configuration submitted to PAI
- */
-export class PAIJobConfig {
-    // Name for the job, need to be unique
-    public readonly jobName: string;
-    // URL pointing to the Docker image for all tasks in the job
-    public readonly image: string;
-    // Code directory on HDFS
-    public readonly codeDir: string;
-    //authentication file used for private Docker registry 
-    public readonly authFile?: string;
-
-    // List of taskRole, one task role at least
-    public taskRoles: PAITaskRole[];
-
-    //The virtual cluster job runs on.
-    public readonly virtualCluster: string;
-
-    /**
-     * Constructor
-     * @param jobName Name for the job, need to be unique
-     * @param image URL pointing to the Docker image for all tasks in the job
-     * @param dataDir Data directory existing on HDFS
-     * @param outputDir Output directory on HDFS
-     * @param taskRoles List of taskRole, one task role at least
-     */
-    constructor(jobName: string, image: string, codeDir: string,
-                taskRoles: PAITaskRole[], virtualCluster: string, authFile?: string) {
-        this.jobName = jobName;
-        this.image = image;
-        this.codeDir = codeDir;
-        this.taskRoles = taskRoles;
-        this.virtualCluster = virtualCluster;
-        this.authFile = authFile;
-    }
-}
-
-/**
- * PAI cluster configuration
- */
 export class PAIClusterConfig {
     public readonly userName: string;
     public readonly passWord?: string;
@@ -111,41 +28,31 @@ export class PAIClusterConfig {
 }
 
 /**
- * portList data structure used in PAI taskRole
+ * PAI trial job detail
  */
-export class PortListMetaData {
-    public readonly label: string = '';
-    public readonly beginAt: number = 0;
-    public readonly portNumber: number = 0;
-}
-  
-
-/**
- * PAI trial configuration
- */
-export class NNIPAITrialConfig extends TrialConfig {
-    public readonly cpuNum: number;
-    public readonly memoryMB: number;
-    public readonly image: string;
-
-    //The virtual cluster job runs on. If omitted, the job will run on default virtual cluster
-    public virtualCluster?: string;
-    //Shared memory for one task in the task role
-    public shmMB?: number;
-    //authentication file used for private Docker registry 
-    public authFile?: string;
-    //portList to specify the port used in container
-    public portList?: PortListMetaData[];
-
-    constructor(command: string, codeDir: string, gpuNum: number, cpuNum: number, memoryMB: number,
-                image: string, virtualCluster?: string, shmMB?: number, authFile?: string, portList?: PortListMetaData[]) {
-        super(command, codeDir, gpuNum);
-        this.cpuNum = cpuNum;
-        this.memoryMB = memoryMB;
-        this.image = image;
-        this.virtualCluster = virtualCluster;
-        this.shmMB = shmMB;
-        this.authFile = authFile;
-        this.portList = portList;
+export class PAITrialJobDetail implements TrialJobDetail {
+    public id: string;
+    public status: TrialJobStatus;
+    public paiJobName: string;
+    public submitTime: number;
+    public startTime?: number;
+    public endTime?: number;
+    public tags?: string[];
+    public url?: string;
+    public workingDirectory: string;
+    public form: TrialJobApplicationForm;
+    public logPath: string;
+    public isEarlyStopped?: boolean;
+
+    constructor(id: string, status: TrialJobStatus, paiJobName: string,
+                submitTime: number, workingDirectory: string, form: TrialJobApplicationForm, logPath: string) {
+        this.id = id;
+        this.status = status;
+        this.paiJobName = paiJobName;
+        this.submitTime = submitTime;
+        this.workingDirectory = workingDirectory;
+        this.form = form;
+        this.tags = [];
+        this.logPath = logPath;
     }
 }

src/nni_manager/training_service/pai/paiJobInfoCollector.ts

@@ -8,8 +8,7 @@ import { Deferred } from 'ts-deferred';
 import { NNIError, NNIErrorNames } from '../../common/errors';
 import { getLogger, Logger } from '../../common/log';
 import { TrialJobStatus } from '../../common/trainingService';
-import { PAIClusterConfig } from './paiConfig';
-import { PAITrialJobDetail } from './paiData';
+import { PAIClusterConfig, PAITrialJobDetail } from './paiConfig';
 
 /**
  * Collector PAI jobs info from PAI cluster, and update pai job status locally
@@ -26,8 +25,8 @@ export class PAIJobInfoCollector {
         this.finalStatuses = ['SUCCEEDED', 'FAILED', 'USER_CANCELED', 'SYS_CANCELED', 'EARLY_STOPPED'];
     }
 
-    public async retrieveTrialStatus(paiToken? : string, paiClusterConfig?: PAIClusterConfig): Promise<void> {
-        if (paiClusterConfig === undefined || paiToken === undefined) {
+    public async retrieveTrialStatus(token? : string, paiBaseClusterConfig?: PAIClusterConfig): Promise<void> {
+        if (paiBaseClusterConfig === undefined || token === undefined) {
             return Promise.resolve();
         }
 
@@ -36,7 +35,7 @@ export class PAIJobInfoCollector {
             if (paiTrialJob === undefined) {
                 throw new NNIError(NNIErrorNames.NOT_FOUND, `trial job id ${trialJobId} not found`);
             }
-            updatePaiTrialJobs.push(this.getSinglePAITrialJobInfo(paiTrialJob, paiToken, paiClusterConfig));
+            updatePaiTrialJobs.push(this.getSinglePAITrialJobInfo(paiTrialJob, token, paiBaseClusterConfig));
         }
 
         await Promise.all(updatePaiTrialJobs);
@@ -56,7 +55,7 @@ export class PAIJobInfoCollector {
             uri: `http://${paiClusterConfig.host}/rest-server/api/v1/user/${paiClusterConfig.userName}/jobs/${paiTrialJob.paiJobName}`,
             method: 'GET',
             json: true,
-            headers: {
+               headers: {
                 'Content-Type': 'application/json',
                 Authorization: `Bearer ${paiToken}`
             }
@@ -114,8 +113,12 @@ export class PAIJobInfoCollector {
                             paiTrialJob.endTime = response.body.jobStatus.completedTime;
                         }
                         // Set pai trial job's url to WebHDFS output path
-                        if (paiTrialJob.hdfsLogPath !== undefined) {
-                            paiTrialJob.url += `,${paiTrialJob.hdfsLogPath}`;
+                        if (paiTrialJob.logPath !== undefined) {
+                            if (paiTrialJob.url) {
+                                paiTrialJob.url += `,${paiTrialJob.logPath}`;
+                            } else {
+                                paiTrialJob.url = `${paiTrialJob.logPath}`;
+                            }
                         }
                     }
                 }

src/nni_manager/training_service/pai/paiJobRestServer.ts

@@ -19,19 +19,17 @@ export interface ParameterFileMeta {
  * PAI Training service Rest server, provides rest API to support pai job metrics update
  *
  */
-@component.Singleton
 export class PAIJobRestServer extends ClusterJobRestServer {
-    private parameterFileMetaList: ParameterFileMeta[] = [];
+    protected parameterFileMetaList: ParameterFileMeta[] = [];
 
-    @Inject
-    private readonly paiTrainingService: PAITrainingService;
+    protected readonly paiTrainingService: PAITrainingService;
 
     /**
      * constructor to provide NNIRestServer's own rest property, e.g. port
      */
-    constructor() {
+    constructor (paiTrainingService: PAITrainingService) {
         super();
-        this.paiTrainingService = component.get(PAITrainingService);
+        this.paiTrainingService = paiTrainingService;
     }
 
     protected handleTrialMetrics(jobId: string, metrics: any[]): void {

src/nni_manager/training_service/pai/paiK8S/paiK8SConfig.ts

+/**
+ * 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 {TrialConfig} from '../../common/trialConfig';
+
+/**
+ * PAI trial configuration
+ */
+export class NNIPAIK8STrialConfig extends TrialConfig {
+    public readonly cpuNum: number;
+    public readonly memoryMB: number;
+    public readonly image: string;
+    public virtualCluster?: string;
+    public readonly nniManagerNFSMountPath: string;
+    public readonly containerNFSMountPath: string;
+    public readonly paiStoragePlugin: string;
+
+    constructor(command: string, codeDir: string, gpuNum: number, cpuNum: number, memoryMB: number,
+                image: string, nniManagerNFSMountPath: string, containerNFSMountPath: string,
+                paiStoragePlugin: string, virtualCluster?: string) {
+        super(command, codeDir, gpuNum);
+        this.cpuNum = cpuNum;
+        this.memoryMB = memoryMB;
+        this.image = image;
+        this.virtualCluster = virtualCluster;
+        this.nniManagerNFSMountPath = nniManagerNFSMountPath;
+        this.containerNFSMountPath = containerNFSMountPath;
+        this.paiStoragePlugin = paiStoragePlugin;
+    }
+}

src/nni_manager/training_service/pai/paiK8S/paiK8SData.ts

+/**
+ * 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 PAI_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 nni
+fi`;
+
+export const PAI_K8S_TRIAL_COMMAND_FORMAT: string =
+`export NNI_PLATFORM=pai NNI_SYS_DIR={0} NNI_OUTPUT_DIR={1} NNI_TRIAL_JOB_ID={2} NNI_EXP_ID={3} NNI_TRIAL_SEQ_ID={4} MULTI_PHASE={5} \
+&& ls $NNI_SYS_DIR \
+&& cd $NNI_SYS_DIR && sh install_nni.sh \
+&& python3 -m nni_trial_tool.trial_keeper --trial_command '{6}' --nnimanager_ip '{7}' --nnimanager_port '{8}' \
+--nni_manager_version '{9}' --log_collection '{10}'`;

src/nni_manager/training_service/pai/paiK8S/paiK8STrainingService.ts

+/**
+ * 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 cpp from 'child-process-promise';
+import * as fs from 'fs';
+import * as path from 'path';
+// tslint:disable-next-line:no-implicit-dependencies
+import * as request from 'request';
+import * as component from '../../../common/component';
+
+import { Deferred } from 'ts-deferred';
+import { String } from 'typescript-string-operations';
+import {
+    HyperParameters, NNIManagerIpConfig, TrainingService,
+    TrialJobApplicationForm, TrialJobDetail, TrialJobMetric
+} from '../../../common/trainingService';
+import { delay, generateParamFileName,
+    getExperimentRootDir, getIPV4Address, getVersion, uniqueString, unixPathJoin } from '../../../common/utils';
+import { CONTAINER_INSTALL_NNI_SHELL_FORMAT } from '../../common/containerJobData';
+import { TrialConfigMetadataKey } from '../../common/trialConfigMetadataKey';
+import { execMkdir, validateCodeDir, execCopydir } from '../../common/util';
+import { PAI_K8S_TRIAL_COMMAND_FORMAT } from './paiK8SData';
+import { NNIPAIK8STrialConfig } from './paiK8SConfig';
+import { PAITrainingService } from '../paiTrainingService';
+import { PAIClusterConfig, PAITrialJobDetail } from '../paiConfig';
+import { PAIJobRestServer } from '../paiJobRestServer';
+
+const yaml = require('js-yaml');
+
+/**
+ * Training Service implementation for OpenPAI (Open Platform for AI)
+ * Refer https://github.com/Microsoft/pai for more info about OpenPAI
+ */
+@component.Singleton
+class PAIK8STrainingService extends PAITrainingService {
+    protected paiTrialConfig: NNIPAIK8STrialConfig | undefined;
+
+    constructor() {
+        super();
+    }
+
+    public async setClusterMetadata(key: string, value: string): Promise<void> {
+        switch (key) {
+            case TrialConfigMetadataKey.PAI_CLUSTER_CONFIG:
+                this.paiJobRestServer = new PAIJobRestServer(component.get(PAIK8STrainingService));
+                this.paiClusterConfig = <PAIClusterConfig>JSON.parse(value);
+                if(this.paiClusterConfig.passWord) {
+                    // Get PAI authentication token
+                    await this.updatePaiToken();
+                } else if(this.paiClusterConfig.token) {
+                    this.paiToken = this.paiClusterConfig.token;
+                }
+                break;
+
+            case TrialConfigMetadataKey.TRIAL_CONFIG:
+                if (this.paiClusterConfig === undefined) {
+                    this.log.error('pai cluster config is not initialized');
+                    break;
+                }
+                this.paiTrialConfig = <NNIPAIK8STrialConfig>JSON.parse(value);
+                // Validate to make sure codeDir doesn't have too many files
+                await validateCodeDir(this.paiTrialConfig.codeDir);
+                break;
+            case TrialConfigMetadataKey.VERSION_CHECK:
+                this.versionCheck = (value === 'true' || value === 'True');
+                break;
+            case TrialConfigMetadataKey.LOG_COLLECTION:
+                this.logCollection = value;
+                break;
+            case TrialConfigMetadataKey.MULTI_PHASE:
+                this.isMultiPhase = (value === 'true' || value === 'True');
+                break;
+            default:
+                //Reject for unknown keys
+                this.log.error(`Uknown key: ${key}`);
+        }
+    }
+    
+    //TODO: update trial parameters
+    public async updateTrialJob(trialJobId: string, form: TrialJobApplicationForm): Promise<TrialJobDetail> {
+        const trialJobDetail: undefined | TrialJobDetail = this.trialJobsMap.get(trialJobId);
+        if (trialJobDetail === undefined) {
+            throw new Error(`updateTrialJob failed: ${trialJobId} not found`);
+        }
+        return trialJobDetail;
+    }
+
+    public async submitTrialJob(form: TrialJobApplicationForm): Promise<TrialJobDetail> {
+        if (this.paiClusterConfig === undefined) {
+            throw new Error(`paiClusterConfig not initialized!`);
+        }
+        if (this.paiTrialConfig === undefined) {
+            throw new Error(`paiTrialConfig not initialized!`);
+        }
+
+        this.log.info(`submitTrialJob: form: ${JSON.stringify(form)}`);
+
+        const trialJobId: string = uniqueString(5);
+        //TODO: use HDFS working folder instead
+        const trialWorkingFolder: string = path.join(this.expRootDir, 'trials', trialJobId);
+        const paiJobName: string = `nni_exp_${this.experimentId}_trial_${trialJobId}`;
+        const logPath: string = path.join(this.paiTrialConfig.nniManagerNFSMountPath, this.experimentId, trialJobId);
+        const trialJobDetail: PAITrialJobDetail = new PAITrialJobDetail(
+            trialJobId,
+            'WAITING',
+            paiJobName,
+            Date.now(),
+            trialWorkingFolder,
+            form,
+            logPath);
+
+        this.trialJobsMap.set(trialJobId, trialJobDetail);
+        this.jobQueue.push(trialJobId);
+
+        return trialJobDetail;
+    }
+
+    public generateJobConfigInYamlFormat(trialJobId: string, command: string) {
+        if (this.paiTrialConfig === undefined) {
+            throw new Error('trial config is not initialized');
+        }
+        const jobName = `nni_exp_${this.experimentId}_trial_${trialJobId}`
+        const paiJobConfig: any = {
+            protocolVersion: 2, 
+            name: jobName,
+            type: 'job',
+            jobRetryCount: 0,
+            prerequisites: [
+              {
+                type: 'dockerimage',
+                uri: this.paiTrialConfig.image,
+                name: 'docker_image_0'
+              }
+            ],
+            taskRoles: {
+                taskrole: {
+                    instances: 1,
+                    completion: {
+                       minFailedInstances: 1,
+                       minSucceededInstances: -1
+                    },
+                    taskRetryCount: 0,
+                    dockerImage: 'docker_image_0',
+                    resourcePerInstance: {
+                        gpu: this.paiTrialConfig.gpuNum,
+                        cpu: this.paiTrialConfig.cpuNum,
+                        memoryMB: this.paiTrialConfig.memoryMB
+                    },
+                    commands: [
+                        command
+                    ]
+                }
+            },
+            extras: {
+                'com.microsoft.pai.runtimeplugin': [
+                    {
+                        plugin: this.paiTrialConfig.paiStoragePlugin
+                    }
+                ],
+                submitFrom: 'submit-job-v2'
+            }
+        }
+        if (this.paiTrialConfig.virtualCluster) {
+            paiJobConfig.defaults=  {
+                virtualCluster: this.paiTrialConfig.virtualCluster
+            }
+        }
+
+        return yaml.safeDump(paiJobConfig);
+      }
+
+    protected async submitTrialJobToPAI(trialJobId: string): Promise<boolean> {
+        const deferred: Deferred<boolean> = new Deferred<boolean>();
+        const trialJobDetail: PAITrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);
+
+        if (trialJobDetail === undefined) {
+            throw new Error(`Failed to find PAITrialJobDetail for job ${trialJobId}`);
+        }
+
+        if (this.paiClusterConfig === undefined) {
+            throw new Error('PAI Cluster config is not initialized');
+        }
+        if (this.paiTrialConfig === undefined) {
+            throw new Error('trial config is not initialized');
+        }
+        if (this.paiToken === undefined) {
+            throw new Error('PAI token is not initialized');
+        }
+
+        if (this.paiJobRestServer === undefined) {
+            throw new Error('paiJobRestServer is not initialized');
+        }
+
+        this.paiRestServerPort = this.paiJobRestServer.clusterRestServerPort;
+
+        // Step 1. Prepare PAI job configuration
+        const trialLocalFolder: string = path.join(this.paiTrialConfig.nniManagerNFSMountPath, this.experimentId, trialJobId);
+        //create trial local working folder locally.
+        await execMkdir(trialLocalFolder);
+
+        const runScriptContent: string = CONTAINER_INSTALL_NNI_SHELL_FORMAT;
+        // Write NNI installation file to local files
+        await fs.promises.writeFile(path.join(trialLocalFolder, 'install_nni.sh'), runScriptContent, { encoding: 'utf8' });
+
+        // Write file content ( parameter.cfg ) to local working folders
+        if (trialJobDetail.form !== undefined) {
+            await fs.promises.writeFile(
+                path.join(trialLocalFolder, generateParamFileName(trialJobDetail.form.hyperParameters)),
+                trialJobDetail.form.hyperParameters.value, { encoding: 'utf8' }
+            );
+        }
+
+        //Copy codeDir files to local working folder
+        await execCopydir(this.paiTrialConfig.codeDir, trialLocalFolder);
+        
+        const nniManagerIp: string = this.nniManagerIpConfig ? this.nniManagerIpConfig.nniManagerIp : getIPV4Address();
+        const version: string = this.versionCheck ? await getVersion() : '';
+        const containerWorkingDir: string = `${this.paiTrialConfig.containerNFSMountPath}/${this.experimentId}/${trialJobId}`;
+        const nniPaiTrialCommand: string = String.Format(
+            PAI_K8S_TRIAL_COMMAND_FORMAT,
+            `${containerWorkingDir}`,
+            `${containerWorkingDir}/nnioutput`,
+            trialJobId,
+            this.experimentId,
+            trialJobDetail.form.sequenceId,
+            this.isMultiPhase,
+            this.paiTrialConfig.command,
+            nniManagerIp,
+            this.paiRestServerPort,
+            version,
+            this.logCollection
+        )
+        .replace(/\r\n|\n|\r/gm, '');
+
+        this.log.info(`nniPAItrial command is ${nniPaiTrialCommand.trim()}`);
+        
+        const paiJobConfig = this.generateJobConfigInYamlFormat(trialJobId, nniPaiTrialCommand);
+
+        // Step 3. Submit PAI job via Rest call
+        // Refer https://github.com/Microsoft/pai/blob/master/docs/rest-server/API.md for more detail about PAI Rest API
+        const submitJobRequest: request.Options = {
+            uri: `http://${this.paiClusterConfig.host}/rest-server/api/v2/jobs`,
+            method: 'POST',
+            body: paiJobConfig,
+            headers: {
+                'Content-Type': 'text/yaml',
+                Authorization: `Bearer ${this.paiToken}`
+            }
+        };
+        request(submitJobRequest, (error: Error, response: request.Response, body: any) => {
+            if ((error !== undefined && error !== null) || response.statusCode >= 400) {
+                const errorMessage: string = (error !== undefined && error !== null) ? error.message :
+                    `Submit trial ${trialJobId} failed, http code:${response.statusCode}, http body: ${body}`;
+
+                this.log.error(errorMessage);
+                trialJobDetail.status = 'FAILED';
+            } else {
+                trialJobDetail.submitTime = Date.now();
+            }
+            deferred.resolve(true);
+        });
+
+        return deferred.promise;
+    }
+}
+
+export { PAIK8STrainingService };

src/nni_manager/training_service/pai/hdfsClientUtility.ts → src/nni_manager/training_service/pai/paiYarn/hdfsClientUtility.ts

@@ -4,9 +4,9 @@
 import * as fs from 'fs';
 import * as path from 'path';
 import { Deferred } from 'ts-deferred';
-import { getExperimentId } from '../../common/experimentStartupInfo';
-import { getLogger } from '../../common/log';
-import { unixPathJoin } from '../../common/utils';
+import { getExperimentId } from '../../../common/experimentStartupInfo';
+import { getLogger } from '../../../common/log';
+import { unixPathJoin } from '../../../common/utils';
 
 /**
  * HDFS client utility, including copy file/directory

src/nni_manager/training_service/pai/paiYarn/paiYarnConfig.ts

+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+'use strict';
+
+import {TrialConfig} from '../../common/trialConfig';
+
+/**
+ * Task role for PAI
+ */
+export class PAITaskRole {
+    // Name for the task role
+    public readonly name: string;
+    // Number of tasks for the task role, no less than 1
+    public readonly taskNumber: number;
+    // CPU number for one task in the task role, no less than 1
+    public readonly cpuNumber: number;
+    // Memory for one task in the task role, no less than 100
+    public readonly memoryMB: number;
+    // GPU number for one task in the task role, no less than 0
+    public readonly gpuNumber: number;
+    // Executable command for tasks in the task role, can not be empty
+    public readonly command: string;
+    //Shared memory for one task in the task role
+    public readonly shmMB?: number;
+    //portList to specify the port used in container
+    public portList?: PortListMetaData[];
+
+    /**
+     * Constructor
+     * @param name Name for the task role
+     * @param taskNumber Number of tasks for the task role, no less than 1
+     * @param cpuNumber CPU number for one task in the task role, no less than 1
+     * @param memoryMB Memory for one task in the task role, no less than 100
+     * @param gpuNumber GPU number for one task in the task role, no less than 0
+     * @param command Executable command for tasks in the task role, can not be empty
+     */
+    constructor(name: string, taskNumber: number, cpuNumber: number, memoryMB: number, gpuNumber: number,
+                command: string, shmMB?: number, portList?: PortListMetaData[]) {
+        this.name = name;
+        this.taskNumber = taskNumber;
+        this.cpuNumber = cpuNumber;
+        this.memoryMB = memoryMB;
+        this.gpuNumber = gpuNumber;
+        this.command = command;
+        this.shmMB = shmMB;
+        this.portList = portList;
+    }
+}
+
+/**
+ * Trial job configuration submitted to PAI
+ */
+export class PAIJobConfig {
+    // Name for the job, need to be unique
+    public readonly jobName: string;
+    // URL pointing to the Docker image for all tasks in the job
+    public readonly image: string;
+    // Code directory on HDFS
+    public readonly codeDir: string;
+    //authentication file used for private Docker registry 
+    public readonly authFile?: string;
+
+    // List of taskRole, one task role at least
+    public taskRoles: PAITaskRole[];
+
+    //The virtual cluster job runs on.
+    public readonly virtualCluster: string;
+
+    /**
+     * Constructor
+     * @param jobName Name for the job, need to be unique
+     * @param image URL pointing to the Docker image for all tasks in the job
+     * @param dataDir Data directory existing on HDFS
+     * @param outputDir Output directory on HDFS
+     * @param taskRoles List of taskRole, one task role at least
+     */
+    constructor(jobName: string, image: string, codeDir: string,
+                taskRoles: PAITaskRole[], virtualCluster: string, authFile?: string) {
+        this.jobName = jobName;
+        this.image = image;
+        this.codeDir = codeDir;
+        this.taskRoles = taskRoles;
+        this.virtualCluster = virtualCluster;
+        this.authFile = authFile;
+    }
+}
+
+/**
+ * portList data structure used in PAI taskRole
+ */
+export class PortListMetaData {
+    public readonly label: string = '';
+    public readonly beginAt: number = 0;
+    public readonly portNumber: number = 0;
+}
+  
+
+/**
+ * PAI trial configuration
+ */
+export class NNIPAITrialConfig extends TrialConfig {
+    public readonly cpuNum: number;
+    public readonly memoryMB: number;
+    public readonly image: string;
+
+    //The virtual cluster job runs on. If omitted, the job will run on default virtual cluster
+    public virtualCluster?: string;
+    //Shared memory for one task in the task role
+    public shmMB?: number;
+    //authentication file used for private Docker registry 
+    public authFile?: string;
+    //portList to specify the port used in container
+    public portList?: PortListMetaData[];
+
+    constructor(command: string, codeDir: string, gpuNum: number, cpuNum: number, memoryMB: number,
+                image: string, virtualCluster?: string, shmMB?: number, authFile?: string, portList?: PortListMetaData[]) {
+        super(command, codeDir, gpuNum);
+        this.cpuNum = cpuNum;
+        this.memoryMB = memoryMB;
+        this.image = image;
+        this.virtualCluster = virtualCluster;
+        this.shmMB = shmMB;
+        this.authFile = authFile;
+        this.portList = portList;
+    }
+}

src/nni_manager/training_service/pai/paiData.ts → src/nni_manager/training_service/pai/paiYarn/paiYarnData.ts

@@ -3,37 +3,7 @@
 
 'use strict';
 
-import { TrialJobApplicationForm, TrialJobDetail, TrialJobStatus  } from '../../common/trainingService';
-
-/**
- * PAI trial job detail
- */
-export class PAITrialJobDetail implements TrialJobDetail {
-    public id: string;
-    public status: TrialJobStatus;
-    public paiJobName: string;
-    public submitTime: number;
-    public startTime?: number;
-    public endTime?: number;
-    public tags?: string[];
-    public url?: string;
-    public workingDirectory: string;
-    public form: TrialJobApplicationForm;
-    public hdfsLogPath: string;
-    public isEarlyStopped?: boolean;
-
-    constructor(id: string, status: TrialJobStatus, paiJobName: string,
-                submitTime: number, workingDirectory: string, form: TrialJobApplicationForm, hdfsLogPath: string) {
-        this.id = id;
-        this.status = status;
-        this.paiJobName = paiJobName;
-        this.submitTime = submitTime;
-        this.workingDirectory = workingDirectory;
-        this.form = form;
-        this.tags = [];
-        this.hdfsLogPath = hdfsLogPath;
-    }
-}
+import { TrialJobApplicationForm, TrialJobDetail, TrialJobStatus  } from '../../../common/trainingService';
 
 export const PAI_INSTALL_NNI_SHELL_FORMAT: string =
 `#!/bin/bash
@@ -46,7 +16,7 @@ else
 fi`;
 
 export const PAI_TRIAL_COMMAND_FORMAT: string =
-`export NNI_PLATFORM=pai NNI_SYS_DIR={0} NNI_OUTPUT_DIR={1} NNI_TRIAL_JOB_ID={2} NNI_EXP_ID={3} NNI_TRIAL_SEQ_ID={4} MULTI_PHASE={5} \
+`export NNI_PLATFORM=paiYarn NNI_SYS_DIR={0} NNI_OUTPUT_DIR={1} NNI_TRIAL_JOB_ID={2} NNI_EXP_ID={3} NNI_TRIAL_SEQ_ID={4} MULTI_PHASE={5} \
 && cd $NNI_SYS_DIR && sh install_nni.sh \
 && python3 -m nni_trial_tool.trial_keeper --trial_command '{6}' --nnimanager_ip '{7}' --nnimanager_port '{8}' \
 --pai_hdfs_output_dir '{9}' --pai_hdfs_host '{10}' --pai_user_name {11} --nni_hdfs_exp_dir '{12}' --webhdfs_path '/webhdfs/api/v1' \

src/nni_manager/training_service/pai/paiYarn/paiYarnTrainingService.ts

+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+'use strict';
+
+import * as fs from 'fs';
+import * as path from 'path';
+import * as request from 'request';
+import * as component from '../../../common/component';
+
+import { EventEmitter } from 'events';
+import { Deferred } from 'ts-deferred';
+import { String } from 'typescript-string-operations';
+import { getExperimentId } from '../../../common/experimentStartupInfo';
+import { getLogger, Logger } from '../../../common/log';
+import {
+    HyperParameters, NNIManagerIpConfig, TrainingService,
+    TrialJobApplicationForm, TrialJobDetail, TrialJobMetric
+} from '../../../common/trainingService';
+import { delay, generateParamFileName,
+    getExperimentRootDir, getIPV4Address, getVersion, uniqueString, unixPathJoin } from '../../../common/utils';
+import { CONTAINER_INSTALL_NNI_SHELL_FORMAT } from '../../common/containerJobData';
+import { TrialConfigMetadataKey } from '../../common/trialConfigMetadataKey';
+import { execMkdir, validateCodeDir } from '../../common/util';
+import { HDFSClientUtility } from './hdfsClientUtility';
+import { NNIPAITrialConfig, PAIJobConfig, PAITaskRole } from './paiYarnConfig';
+import { PAI_LOG_PATH_FORMAT, PAI_TRIAL_COMMAND_FORMAT } from './paiYarnData';
+import { PAIJobInfoCollector } from '../paiJobInfoCollector';
+import { PAITrainingService } from '../paiTrainingService';
+import { PAIClusterConfig, PAITrialJobDetail } from '../paiConfig';
+
+import * as WebHDFS from 'webhdfs';
+import { PAIJobRestServer, ParameterFileMeta } from '../paiJobRestServer';
+
+/**
+ * Training Service implementation for OpenPAI (Open Platform for AI)
+ * Refer https://github.com/Microsoft/pai for more info about OpenPAI
+ */
+@component.Singleton
+class PAIYarnTrainingService extends PAITrainingService {
+    private hdfsClient: any;
+    private copyExpCodeDirPromise?: Promise<void>;
+    private copyAuthFilePromise?: Promise<void>;
+    private paiTrialConfig?: NNIPAITrialConfig;
+
+    constructor() {
+        super();
+    }
+
+    public async submitTrialJob(form: TrialJobApplicationForm): Promise<TrialJobDetail> {
+        if (this.paiClusterConfig === undefined) {
+            throw new Error(`paiBaseClusterConfig not initialized!`);
+        }
+
+        this.log.info(`submitTrialJob: form: ${JSON.stringify(form)}`);
+
+        const trialJobId: string = uniqueString(5);
+        //TODO: use HDFS working folder instead
+        const trialWorkingFolder: string = path.join(this.expRootDir, 'trials', trialJobId);
+        const paiJobName: string = `nni_exp_${this.experimentId}_trial_${trialJobId}`;
+        const hdfsCodeDir: string = HDFSClientUtility.getHdfsTrialWorkDir(this.paiClusterConfig.userName, trialJobId);
+        const hdfsOutputDir: string = unixPathJoin(hdfsCodeDir, 'nnioutput');
+
+        const hdfsLogPath: string = String.Format(
+            PAI_LOG_PATH_FORMAT,
+            this.paiClusterConfig.host,
+            hdfsOutputDir
+            );
+
+        const trialJobDetail: PAITrialJobDetail = new PAITrialJobDetail(
+            trialJobId,
+            'WAITING',
+            paiJobName,
+            Date.now(),
+            trialWorkingFolder,
+            form,
+            hdfsLogPath);
+
+        this.trialJobsMap.set(trialJobId, trialJobDetail);
+        this.jobQueue.push(trialJobId);
+
+        return trialJobDetail;
+    }
+
+    public async setClusterMetadata(key: string, value: string): Promise<void> {
+        switch (key) {
+            case TrialConfigMetadataKey.NNI_MANAGER_IP:
+                this.nniManagerIpConfig = <NNIManagerIpConfig>JSON.parse(value);
+                break;
+
+            case TrialConfigMetadataKey.PAI_YARN_CLUSTER_CONFIG:
+                this.paiJobRestServer = new PAIJobRestServer(component.get(PAIYarnTrainingService));
+                this.paiClusterConfig = <PAIClusterConfig>JSON.parse(value);
+
+                this.hdfsClient = WebHDFS.createClient({
+                    user: this.paiClusterConfig.userName,
+                    // Refer PAI document for Pylon mapping https://github.com/Microsoft/pai/tree/master/docs/pylon
+                    port: 80,
+                    path: '/webhdfs/api/v1',
+                    host: this.paiClusterConfig.host
+                });
+                if(this.paiClusterConfig.passWord) {
+                    // Get PAI authentication token
+                    await this.updatePaiToken();
+                } else if(this.paiClusterConfig.token) {
+                    this.paiToken = this.paiClusterConfig.token;
+                } else {
+                    throw new Error('pai cluster config format error, please set password or token!');
+                }
+
+                break;
+
+            case TrialConfigMetadataKey.TRIAL_CONFIG:
+                if (this.paiClusterConfig === undefined) {
+                    this.log.error('pai cluster config is not initialized');
+                    break;
+                }
+                this.paiTrialConfig = <NNIPAITrialConfig>JSON.parse(value);
+
+                // Validate to make sure codeDir doesn't have too many files
+                await validateCodeDir(this.paiTrialConfig.codeDir);
+           
+                // Copy experiment files from local folder to HDFS
+                this.copyExpCodeDirPromise = HDFSClientUtility.copyDirectoryToHdfs(
+                    this.paiTrialConfig.codeDir,
+                    HDFSClientUtility.getHdfsExpCodeDir(this.paiClusterConfig.userName),
+                    this.hdfsClient
+                );
+                
+                // Upload authFile to hdfs
+                if (this.paiTrialConfig.authFile) {
+                    this.authFileHdfsPath = unixPathJoin(HDFSClientUtility.hdfsExpRootDir(this.paiClusterConfig.userName), 'authFile');
+                    this.copyAuthFilePromise = HDFSClientUtility.copyFileToHdfs(this.paiTrialConfig.authFile, this.authFileHdfsPath, this.hdfsClient);
+                }
+
+                break;
+            case TrialConfigMetadataKey.VERSION_CHECK:
+                this.versionCheck = (value === 'true' || value === 'True');
+                break;
+            case TrialConfigMetadataKey.LOG_COLLECTION:
+                this.logCollection = value;
+                break;
+            case TrialConfigMetadataKey.MULTI_PHASE:
+                this.isMultiPhase = (value === 'true' || value === 'True');
+                break;
+            default:
+                //Reject for unknown keys
+                throw new Error(`Uknown key: ${key}`);
+        }
+    }
+
+    protected async submitTrialJobToPAI(trialJobId: string): Promise<boolean> {
+        const deferred: Deferred<boolean> = new Deferred<boolean>();
+        const trialJobDetail: PAITrialJobDetail | undefined = this.trialJobsMap.get(trialJobId);
+
+        if (trialJobDetail === undefined) {
+            throw new Error(`Failed to find PAITrialJobDetail for job ${trialJobId}`);
+        }
+
+        if (this.paiClusterConfig === undefined) {
+            throw new Error('PAI Cluster config is not initialized');
+        }
+        if (this.paiTrialConfig === undefined) {
+            throw new Error('trial config is not initialized');
+        }
+        if (this.paiToken === undefined) {
+            throw new Error('PAI token is not initialized');
+        }
+
+        if (this.paiJobRestServer === undefined) {
+            throw new Error('paiJobRestServer is not initialized');
+        }
+
+        this.paiRestServerPort = this.paiJobRestServer.clusterRestServerPort;
+
+        // Make sure experiment code files is copied from local to HDFS
+        if (this.copyExpCodeDirPromise !== undefined) {
+            await this.copyExpCodeDirPromise;
+        }
+
+        //Make sure authFile is copied from local to HDFS
+        if (this.paiTrialConfig.authFile) {
+            await this.copyAuthFilePromise;
+        }
+        // Step 1. Prepare PAI job configuration
+
+        const trialLocalTempFolder: string = path.join(getExperimentRootDir(), 'trials-local', trialJobId);
+        //create tmp trial working folder locally.
+        await execMkdir(trialLocalTempFolder);
+
+        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' });
+
+        // Write file content ( parameter.cfg ) to local tmp folders
+        if (trialJobDetail.form !== undefined) {
+            await fs.promises.writeFile(
+                path.join(trialLocalTempFolder, generateParamFileName(trialJobDetail.form.hyperParameters)),
+                trialJobDetail.form.hyperParameters.value, { encoding: 'utf8' }
+            );
+        }
+        const hdfsCodeDir: string = HDFSClientUtility.getHdfsTrialWorkDir(this.paiClusterConfig.userName, trialJobId);
+        const hdfsOutputDir: string = unixPathJoin(hdfsCodeDir, 'nnioutput');
+        const nniManagerIp: string = this.nniManagerIpConfig ? this.nniManagerIpConfig.nniManagerIp : getIPV4Address();
+        const version: string = this.versionCheck ? await getVersion() : '';
+        const nniPaiTrialCommand: string = String.Format(
+            PAI_TRIAL_COMMAND_FORMAT,
+            // PAI will copy job's codeDir into /root directory
+            `$PWD/${trialJobId}`,
+            `$PWD/${trialJobId}/nnioutput`,
+            trialJobId,
+            this.experimentId,
+            trialJobDetail.form.sequenceId,
+            this.isMultiPhase,
+            this.paiTrialConfig.command,
+            nniManagerIp,
+            this.paiRestServerPort,
+            hdfsOutputDir,
+            this.paiClusterConfig.host,
+            this.paiClusterConfig.userName,
+            HDFSClientUtility.getHdfsExpCodeDir(this.paiClusterConfig.userName),
+            version,
+            this.logCollection
+        )
+        .replace(/\r\n|\n|\r/gm, '');
+
+        this.log.info(`nniPAItrial command is ${nniPaiTrialCommand.trim()}`);
+        const paiTaskRoles: PAITaskRole[] = [
+            new PAITaskRole(
+                `nni_trail_${trialJobId}`,
+                // Task role number
+                1,
+                // Task CPU number
+                this.paiTrialConfig.cpuNum,
+                // Task memory
+                this.paiTrialConfig.memoryMB,
+                // Task GPU number
+                this.paiTrialConfig.gpuNum,
+                // Task command
+                nniPaiTrialCommand,
+                // Task shared memory
+                this.paiTrialConfig.shmMB,
+                // Task portList
+                this.paiTrialConfig.portList
+            )
+        ];
+
+        const paiJobConfig: PAIJobConfig = new PAIJobConfig(
+            // Job name
+            trialJobDetail.paiJobName,
+            // Docker image
+            this.paiTrialConfig.image,
+            // codeDir
+            `$PAI_DEFAULT_FS_URI${hdfsCodeDir}`,
+            // PAI Task roles
+            paiTaskRoles,
+            // Add Virutal Cluster
+            this.paiTrialConfig.virtualCluster === undefined ? 'default' : this.paiTrialConfig.virtualCluster.toString(),
+            //Task auth File
+            this.authFileHdfsPath
+        );
+
+        // Step 2. Upload code files in codeDir onto HDFS
+        try {
+            await HDFSClientUtility.copyDirectoryToHdfs(trialLocalTempFolder, hdfsCodeDir, this.hdfsClient);
+        } catch (error) {
+            this.log.error(`PAI Training service: copy ${this.paiTrialConfig.codeDir} to HDFS ${hdfsCodeDir} failed, error is ${error}`);
+            trialJobDetail.status = 'FAILED'; // eslint-disable-line require-atomic-updates
+            return true;
+        }
+
+        // Step 3. Submit PAI job via Rest call
+        // Refer https://github.com/Microsoft/pai/blob/master/docs/rest-server/API.md for more detail about PAI Rest API
+        const submitJobRequest: request.Options = {
+            uri: `http://${this.paiClusterConfig.host}/rest-server/api/v1/user/${this.paiClusterConfig.userName}/jobs`,
+            method: 'POST',
+            json: true,
+            body: paiJobConfig,
+            headers: {
+                'Content-Type': 'application/json',
+                Authorization: `Bearer ${this.paiToken}`
+            }
+        };
+        request(submitJobRequest, (error: Error, response: request.Response, body: any) => {
+            if ((error !== undefined && error !== null) || response.statusCode >= 400) {
+                const errorMessage: string = (error !== undefined && error !== null) ? error.message :
+                    `Submit trial ${trialJobId} failed, http code:${response.statusCode}, http body: ${response.body.message}`;
+                trialJobDetail.status = 'FAILED';
+                deferred.resolve(true);
+            } else {
+                trialJobDetail.submitTime = Date.now();
+                deferred.resolve(true);
+            }
+        });
+
+        return deferred.promise;
+    }
+
+    public async updateTrialJob(trialJobId: string, form: TrialJobApplicationForm): Promise<TrialJobDetail> {
+        const trialJobDetail: undefined | TrialJobDetail = this.trialJobsMap.get(trialJobId);
+        if (trialJobDetail === undefined) {
+            throw new Error(`updateTrialJob failed: ${trialJobId} not found`);
+        }
+        await this.writeParameterFile(trialJobId, form.hyperParameters);
+
+        return trialJobDetail;
+    }
+
+    protected async writeParameterFile(trialJobId: string, hyperParameters: HyperParameters): Promise<void> {
+        if (this.paiClusterConfig === undefined) {
+            throw new Error('PAI Cluster config is not initialized');
+        }
+        if (this.paiTrialConfig === undefined) {
+            throw new Error('PAI trial config is not initialized');
+        }
+
+        const trialLocalTempFolder: string = path.join(getExperimentRootDir(), 'trials-local', trialJobId);
+        const hpFileName: string = generateParamFileName(hyperParameters);
+        const localFilepath: string = path.join(trialLocalTempFolder, hpFileName);
+        await fs.promises.writeFile(localFilepath, hyperParameters.value, { encoding: 'utf8' });
+        const hdfsCodeDir: string = HDFSClientUtility.getHdfsTrialWorkDir(this.paiClusterConfig.userName, trialJobId);
+        const hdfsHpFilePath: string = path.join(hdfsCodeDir, hpFileName);
+
+        await HDFSClientUtility.copyFileToHdfs(localFilepath, hdfsHpFilePath, this.hdfsClient);
+
+        await this.postParameterFileMeta({
+            experimentId: this.experimentId,
+            trialId: trialJobId,
+            filePath: hdfsHpFilePath
+        });
+    }
+
+    protected postParameterFileMeta(parameterFileMeta: ParameterFileMeta): Promise<void> {
+        const deferred: Deferred<void> = new Deferred<void>();
+        if (this.paiJobRestServer === undefined) {
+            throw new Error('paiJobRestServer not implemented!');
+        }
+        const req: request.Options = {
+            uri: `${this.paiJobRestServer.endPoint}${this.paiJobRestServer.apiRootUrl}/parameter-file-meta`,
+            method: 'POST',
+            json: true,
+            body: parameterFileMeta
+        };
+        request(req, (err: Error, res: request.Response) => {
+            if (err) {
+                deferred.reject(err);
+            } else {
+                deferred.resolve();
+            }
+        });
+
+        return deferred.promise;
+    }
+}
+
+export { PAIYarnTrainingService };

src/nni_manager/training_service/pai/paiTrialConfig.ts → src/nni_manager/training_service/pai/paiYarn/paiYarnTrialConfig.ts

@@ -3,12 +3,12 @@
 
 'use strict';
 
-import {TrialConfig} from '../common/trialConfig';
+import {TrialConfig} from '../../common/trialConfig';
 
 /**
  * PAI configuration to run trials
  */
-export class PAITrialConfig extends TrialConfig {
+export class PAIYarnTrialConfig extends TrialConfig {
     public readonly cpuNum: number;
     public readonly memoryMB: number;
     public readonly image: string;

src/nni_manager/training_service/test/hdfsClientUtility.test.ts

@@ -9,7 +9,7 @@ import * as os from 'os';
 import * as path from 'path';
 import * as tmp from 'tmp';
 import { cleanupUnitTest, prepareUnitTest, uniqueString } from '../../common/utils';
-import { HDFSClientUtility } from '../pai/hdfsClientUtility';
+import { HDFSClientUtility } from '../pai/paiYarn/hdfsClientUtility';
 
 var WebHDFS = require('webhdfs');
 var rmdir = require('rmdir');

src/nni_manager/training_service/test/paiTrainingService.test.ts → src/nni_manager/training_service/test/paiYarnTrainingService.test.ts

@@ -11,14 +11,14 @@ import * as component from '../../common/component';
 import { TrialJobApplicationForm } from '../../common/trainingService';
 import { cleanupUnitTest, prepareUnitTest } from '../../common/utils';
 import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
-import { PAITrainingService } from '../pai/paiTrainingService';
+import { PAIYarnTrainingService } from '../pai/paiYarn/paiYarnTrainingService';
 
 // TODO: copy mockedTrail.py to local folder
 const localCodeDir: string = tmp.dirSync().name
 const mockedTrialPath: string = './training_service/test/mockedTrial.py'
 fs.copyFileSync(mockedTrialPath, localCodeDir + '/mockedTrial.py')
 
-describe('Unit Test for PAITrainingService', () => {
+describe('Unit Test for PAIYarnTrainingService', () => {
     let skip: boolean = false;
     let testPaiClusterInfo: any;
     let paiCluster: any;
@@ -33,7 +33,7 @@ describe('Unit Test for PAITrainingService', () => {
         skip = true;
     }
 
-    let paiTrainingService: PAITrainingService;
+    let paiYarnTrainingService: PAIYarnTrainingService;
 
     console.log(tmp.dirSync().name);
 
@@ -51,15 +51,15 @@ describe('Unit Test for PAITrainingService', () => {
         if (skip) {
             return;
         }
-        paiTrainingService = component.get(PAITrainingService);
-        paiTrainingService.run();
+        paiYarnTrainingService = component.get(PAIYarnTrainingService);
+        paiYarnTrainingService.run();
     });
 
     afterEach(() => {
         if (skip) {
             return;
         }
-        paiTrainingService.cleanUp();
+        paiYarnTrainingService.cleanUp();
     });
 
     it('Get PAI token', async () => {
@@ -67,14 +67,14 @@ describe('Unit Test for PAITrainingService', () => {
             return;
         }
         console.log(`paiCluster is ${paiCluster}`)
-        await paiTrainingService.setClusterMetadata(TrialConfigMetadataKey.PAI_CLUSTER_CONFIG, paiCluster);
-        await paiTrainingService.setClusterMetadata(TrialConfigMetadataKey.TRIAL_CONFIG, paiTrialConfig);
+        await paiYarnTrainingService.setClusterMetadata(TrialConfigMetadataKey.PAI_YARN_CLUSTER_CONFIG, paiCluster);
+        await paiYarnTrainingService.setClusterMetadata(TrialConfigMetadataKey.TRIAL_CONFIG, paiTrialConfig);
         const form: TrialJobApplicationForm = {
             sequenceId: 0,
             hyperParameters: { value: '', index: 0 }
         };
         try {
-            const trialDetail = await paiTrainingService.submitTrialJob(form);
+            const trialDetail = await paiYarnTrainingService.submitTrialJob(form);
             chai.expect(trialDetail.status).to.be.equals('WAITING');
         } catch(error) {
             console.log('Submit job failed:' + error);

src/sdk/pynni/nni/platform/__init__.py

@@ -7,7 +7,7 @@ if trial_env_vars.NNI_PLATFORM is None:
     from .standalone import *
 elif trial_env_vars.NNI_PLATFORM == 'unittest':
     from .test import *
-elif trial_env_vars.NNI_PLATFORM in ('local', 'remote', 'pai', 'kubeflow', 'frameworkcontroller'):
+elif trial_env_vars.NNI_PLATFORM in ('local', 'remote', 'pai', 'kubeflow', 'frameworkcontroller', 'paiYarn'):
     from .local import *
 else:
     raise RuntimeError('Unknown platform %s' % trial_env_vars.NNI_PLATFORM)