Support kubeflow reuse mode

docs/en_US/TrainingService/KubeflowMode.rst

@@ -253,3 +253,41 @@ version check
 NNI support version check feature in since version 0.6, `refer <PaiMode.rst>`__
 
 Any problems when using NNI in Kubeflow mode, please create issues on `NNI Github repo <https://github.com/Microsoft/nni>`__.
+
+
+Kubeflow reuse mode
+----------------------
+NNI support setting reuse mode for trial jobs. In reuse mode, NNI will submit a long-running trial runner process to occupy the container, and start trial jobs as the subprocess of the trial runner process, it means k8s do not need to schedule new container again, it just reuse old container.
+Currently, kubeflow reuse mode only support V2 config.
+Here is the example:
+
+.. code-block:: yaml
+
+   searchSpaceFile: search_space.json
+   trialCommand: python3 mnist.py
+   trialGpuNumber: 0
+   trialConcurrency: 4
+   maxTrialNumber: 20
+   tuner:
+     name: TPE
+     classArgs:
+       optimize_mode: maximize
+   trainingService:
+     reuseMode: true
+     platform: kubeflow
+     worker:
+       command: python3 mnist.py
+       code_directory: .
+       dockerImage: msranni/nni
+       cpuNumber: 1
+       gpuNumber: 0
+       memorySize: 8192
+       replicas: 1
+     operator: tf-operator
+     storage:
+       storageType: azureStorage
+       azureAccount: {your_account}
+       azureShare: {your_share}
+       keyVaultName: {your_valut_name}
+       keyVaultKey: {your_valut_key}
+     apiVersion: v1

nni/experiment/config/kubeflow.py

@@ -8,21 +8,21 @@
 from .common import TrainingServiceConfig
 from . import util
 
-__all__ = ['KubeflowConfig', 'KubeflowRoleConfig', 'KubeflowNfsConfig', 'KubeflowAzureStorageConfig']
+__all__ = ['KubeflowConfig', 'KubeflowRoleConfig', 'KubeflowStorageConfig', 'KubeflowNfsConfig', 'KubeflowAzureStorageConfig']
 
 
 @dataclass(init=False)
-class _KubeflowStorageConfig(ConfigBase):
-    storage: str
+class KubeflowStorageConfig(ConfigBase):
+    storage_type: str
     server: Optional[str] = None
     path: Optional[str] = None
     azure_account: Optional[str] = None
     azure_share: Optional[str] = None
-    key_vault: Optional[str] = None
-    key_vault_secret: Optional[str] = None
+    key_vault_name: Optional[str] = None
+    key_vault_key: Optional[str] = None
 
 @dataclass(init=False)
-class KubeflowNfsConfig(_KubeflowStorageConfig):
+class KubeflowNfsConfig(KubeflowStorageConfig):
     storage: str = 'nfs'
     server: str
     path: str
@@ -32,37 +32,41 @@ class KubeflowAzureStorageConfig(ConfigBase):
     storage: str = 'azureStorage'
     azure_account: str
     azure_share: str
-    key_vault: str
-    key_vault_secret: str
+    key_vault_name: str
+    key_vault_key: str
 
 
 @dataclass(init=False)
 class KubeflowRoleConfig(ConfigBase):
     replicas: int
     command: str
-    gpu_number: int
+    gpu_number: Optional[int] = 0
     cpu_number: int
     memory_size: str
     docker_image: str = 'msranni/nni:latest'
+    code_directory: str
 
 
 @dataclass(init=False)
 class KubeflowConfig(TrainingServiceConfig):
     platform: str = 'kubeflow'
     operator: str
     api_version: str
-    storage: _KubeflowStorageConfig
-    worker: KubeflowRoleConfig
-    parameter_server: Optional[KubeflowRoleConfig] = None
+    storage: KubeflowStorageConfig
+    worker: Optional[KubeflowRoleConfig] = None
+    ps: Optional[KubeflowRoleConfig] = None
+    master: Optional[KubeflowRoleConfig] = None
+    reuse_mode: Optional[bool] = True #set reuse mode as true for v2 config
 
     def __init__(self, **kwargs):
         kwargs = util.case_insensitive(kwargs)
-        kwargs['storage'] = util.load_config(_KubeflowStorageConfig, kwargs.get('storage'))
+        kwargs['storage'] = util.load_config(KubeflowStorageConfig, kwargs.get('storage'))
         kwargs['worker'] = util.load_config(KubeflowRoleConfig, kwargs.get('worker'))
-        kwargs['parameterserver'] = util.load_config(KubeflowRoleConfig, kwargs.get('parameterserver'))
+        kwargs['ps'] = util.load_config(KubeflowRoleConfig, kwargs.get('ps'))
+        kwargs['master'] = util.load_config(KubeflowRoleConfig, kwargs.get('master'))
         super().__init__(**kwargs)
 
     _validation_rules = {
         'platform': lambda value: (value == 'kubeflow', 'cannot be modified'),
         'operator': lambda value: value in ['tf-operator', 'pytorch-operator']
-    }
+    }

ts/nni_manager/common/experimentConfig.ts

@@ -5,6 +5,9 @@
 
 import * as assert from 'assert';
 
+import { KubeflowOperator, OperatorApiVersion } from '../training_service/kubernetes/kubeflow/kubeflowConfig'
+import { KubernetesStorageKind } from '../training_service/kubernetes/kubernetesConfig';
+
 export interface TrainingServiceConfig {
     platform: string;
 }
@@ -68,35 +71,38 @@ export interface AmlConfig extends TrainingServiceConfig {
     maxTrialNumberPerGpu: number;
 }
 
-/* Kubeflow */
-
-// FIXME: merge with shared storage config
 export interface KubeflowStorageConfig {
-    storage: string;
+    storageType: string;
+    maxTrialNumberPerGpu?: number;
     server?: string;
     path?: string;
     azureAccount?: string;
     azureShare?: string;
-    keyVault?: string;
-    keyVaultSecret?: string;
+    keyVaultName?: string;
+    keyVaultKey?: string;
 }
 
 export interface KubeflowRoleConfig {
     replicas: number;
+    codeDirectory: string;
     command: string;
     gpuNumber: number;
     cpuNumber: number;
-    memorySize: string;
+    memorySize: number;
     dockerImage: string;
+    privateRegistryAuthPath?: string;
 }
 
 export interface KubeflowConfig extends TrainingServiceConfig {
     platform: 'kubeflow';
-    operator: string;
-    apiVersion: string;
+    ps?: KubeflowRoleConfig;
+    master?: KubeflowRoleConfig;
+    worker?: KubeflowRoleConfig;
+    maxTrialNumberPerGpu: number;
+    operator: KubeflowOperator;
+    apiVersion: OperatorApiVersion;
     storage: KubeflowStorageConfig;
-    worker: KubeflowRoleConfig;
-    parameterServer?: KubeflowRoleConfig;
+    reuseMode: boolean;
 }
 
 /* FrameworkController */
@@ -221,4 +227,4 @@ export function flattenConfig<T>(config: ExperimentConfig, platform: string): T
         Object.assign(flattened, config.trainingService);
     }
     return <T>flattened;
-}
+}

ts/nni_manager/core/nnimanager.ts

@@ -453,9 +453,6 @@ class NNIManager implements Manager {
         if (platform === 'local') {
             const module_ = await import('../training_service/local/localTrainingService');
             return new module_.LocalTrainingService(config);
-        } else if (platform === 'kubeflow') {
-            const module_ = await import('../training_service/kubernetes/kubeflow/kubeflowTrainingService');
-            return new module_.KubeflowTrainingService();
         } else if (platform === 'frameworkcontroller') {
             const module_ = await import('../training_service/kubernetes/frameworkcontroller/frameworkcontrollerTrainingService');
             return new module_.FrameworkControllerTrainingService();

ts/nni_manager/training_service/reusable/environments/environmentServiceFactory.ts

@@ -2,6 +2,7 @@ import { AMLEnvironmentService } from './amlEnvironmentService';
 import { OpenPaiEnvironmentService } from './openPaiEnvironmentService';
 import { LocalEnvironmentService } from './localEnvironmentService';
 import { RemoteEnvironmentService } from './remoteEnvironmentService';
+import { KubeflowEnvironmentService } from './kubernetes/kubeflowEnvironmentService';
 import { EnvironmentService } from '../environment';
 import { ExperimentConfig } from '../../../common/experimentConfig';
 import { ExperimentStartupInfo } from '../../../common/experimentStartupInfo';
@@ -20,6 +21,8 @@ export async function createEnvironmentService(name: string, config: ExperimentC
             return new AMLEnvironmentService(config, info);
         case 'openpai':
             return new OpenPaiEnvironmentService(config, info);
+        case 'kubeflow':
+            return new KubeflowEnvironmentService(config, info);
     }
 
     const esConfig = await getCustomEnvironmentServiceConfig(name);
@@ -29,4 +32,4 @@ export async function createEnvironmentService(name: string, config: ExperimentC
     const esModule = importModule(esConfig.nodeModulePath);
     const esClass = esModule[esConfig.nodeClassName] as any;
     return new esClass(config, info);
-}
+}