Support remote training service use reuse mode

docs/en_US/Tutorial/ExperimentConfig.md

@@ -592,6 +592,14 @@ Specifies the pre-command that will be executed before the remote machine execut
 
 __Note__: Because __preCommand__ will execute before other commands each time, it is strongly not recommended to set __preCommand__ that will make changes to system, i.e. `mkdir` or `touch`.
 
+### remoteConfig
+
+Optional field in remote mode. Set remote machine related configuration.
+
+#### reuse
+
+Optional. Set if use trial_runner to maintan multiple trial.  
+
 ### kubeflowConfig
 
 #### operator

src/nni_manager/main.ts

@@ -23,9 +23,6 @@ import { KubeflowTrainingService } from './training_service/kubernetes/kubeflow/
 import { LocalTrainingService } from './training_service/local/localTrainingService';
 import { RouterTrainingService } from './training_service/reusable/routerTrainingService';
 import { PAIYarnTrainingService } from './training_service/pai/paiYarn/paiYarnTrainingService';
-import {
-    RemoteMachineTrainingService
-} from './training_service/remote_machine/remoteMachineTrainingService';
 import { DLTSTrainingService } from './training_service/dlts/dltsTrainingService';
 
 function initStartupInfo(
@@ -43,7 +40,7 @@ async function initContainer(foreground: boolean, platformMode: string, logFileN
             .scope(Scope.Singleton);
     } else if (platformMode === 'remote') {
         Container.bind(TrainingService)
-            .to(RemoteMachineTrainingService)
+            .to(RouterTrainingService)
             .scope(Scope.Singleton);
     } else if (platformMode === 'pai') {
         Container.bind(TrainingService)

src/nni_manager/rest_server/restValidationSchemas.ts

@@ -166,6 +166,9 @@ export namespace ValidationSchemas {
             }),
             nni_manager_ip: joi.object({ // eslint-disable-line @typescript-eslint/camelcase
                 nniManagerIp: joi.string().min(1)
+            }),
+            remote_config: joi.object({ // eslint-disable-line @typescript-eslint/camelcase
+                reuse: joi.boolean()
             })
         }
     };

src/nni_manager/training_service/common/trialConfigMetadataKey.ts

@@ -10,6 +10,7 @@ export enum TrialConfigMetadataKey {
     MACHINE_LIST = 'machine_list',
     LOCAL_CONFIG = 'local_config',
     TRIAL_CONFIG = 'trial_config',
+    REMOTE_CONFIG = 'remote_config',
     EXPERIMENT_ID = 'experimentId',
     MULTI_PHASE = 'multiPhase',
     RANDOM_SCHEDULER = 'random_scheduler',

src/nni_manager/training_service/reusable/remote/remoteConfig.ts

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { EnvironmentInformation } from '../environment';
+import { RemoteMachineTrialJobDetail } from '../../remote_machine/remoteMachineData';
+import { TrialJobApplicationForm } from '../../../common/trainingService';
+
+
+/**
+ * work around here, need RemoteMachineTrialJobDetail data structure to schedule machines
+ */
+export class RemoteMachineMetaDetail extends RemoteMachineTrialJobDetail {
+    constructor() {
+        // work around, the form data is a placeholder
+        const form: TrialJobApplicationForm = {
+            sequenceId: 0,
+            hyperParameters: {
+                value: '',
+                index: 0
+            }
+        };
+        super('', 'WAITING', 1, '', form);
+    }
+}
+
+/**
+ * RemoteMachineEnvironmentInformation
+ */
+export class RemoteMachineEnvironmentInformation extends EnvironmentInformation {
+    public rmMachineMetaDetail?: RemoteMachineMetaDetail;
+}
+
+export class RemoteConfig {
+    public readonly reuse: boolean;
+
+    /**
+     * Constructor
+     * @param reuse If job is reusable for multiple trials
+     */
+    constructor(reuse: boolean) {
+        this.reuse = reuse;
+    }
+}

src/nni_manager/training_service/reusable/routerTrainingService.ts

@@ -12,12 +12,15 @@ import { delay } from '../../common/utils';
 import { TrialConfigMetadataKey } from '../common/trialConfigMetadataKey';
 import { PAIClusterConfig } from '../pai/paiConfig';
 import { PAIK8STrainingService } from '../pai/paiK8S/paiK8STrainingService';
+import { RemoteMachineTrainingService } from '../remote_machine/remoteMachineTrainingService';
 import { EnvironmentService } from './environment';
 import { OpenPaiEnvironmentService } from './environments/openPaiEnvironmentService';
 import { AMLEnvironmentService } from './environments/amlEnvironmentService';
+import { RemoteEnvironmentService } from './environments/remoteEnvironmentService';
 import { MountedStorageService } from './storages/mountedStorageService';
 import { StorageService } from './storageService';
 import { TrialDispatcher } from './trialDispatcher';
+import { RemoteConfig } from './remote/remoteConfig';
 
 
 /**
@@ -146,6 +149,18 @@ class RouterTrainingService implements TrainingService {
                 await this.internalTrainingService.setClusterMetadata(key, value);
 
                 this.metaDataCache.clear();
+            } else if (key === TrialConfigMetadataKey.REMOTE_CONFIG) {
+                const config = <RemoteConfig>JSON.parse(value);
+                if (config.reuse === true) {
+                    this.log.info(`reuse flag enabled, use EnvironmentManager.`);
+                    this.internalTrainingService = component.get(TrialDispatcher);
+                    Container.bind(EnvironmentService)
+                        .to(RemoteEnvironmentService)
+                        .scope(Scope.Singleton);
+                } else {
+                    this.log.debug(`caching metadata key:{} value:{}, as training service is not determined.`);
+                    this.internalTrainingService = component.get(RemoteMachineTrainingService);
+                }
             } else {
                 this.log.debug(`caching metadata key:{} value:{}, as training service is not determined.`);
                 this.metaDataCache.set(key, value);

tools/nni_cmd/config_schema.py

@@ -372,6 +372,12 @@ def validate(self, data):
     })
 }
 
+remote_config_schema = {
+    Optional('remoteConfig'): {
+        'reuse': setType('reuse', bool)
+    }
+}
+
 machine_list_schema = {
     'machineList': [Or(
         {
@@ -399,7 +405,7 @@ def validate(self, data):
 
 training_service_schema_dict = {
     'local': Schema({**common_schema, **common_trial_schema}),
-    'remote': Schema({**common_schema, **common_trial_schema, **machine_list_schema}),
+    'remote': Schema({**common_schema, **common_trial_schema, **machine_list_schema, **remote_config_schema}),
     'pai': Schema({**common_schema, **pai_trial_schema, **pai_config_schema}),
     'paiYarn': Schema({**common_schema, **pai_yarn_trial_schema, **pai_yarn_config_schema}),
     'kubeflow': Schema({**common_schema, **kubeflow_trial_schema, **kubeflow_config_schema}),

tools/nni_cmd/launcher.py

@@ -138,6 +138,7 @@ def set_remote_config(experiment_config, port, config_file_name):
     '''Call setClusterMetadata to pass trial'''
     #set machine_list
     request_data = dict()
+    request_data['remote_config'] = experiment_config['remoteConfig']
     request_data['machine_list'] = experiment_config['machineList']
     if request_data['machine_list']:
         for i in range(len(request_data['machine_list'])):
@@ -301,7 +302,6 @@ def set_experiment(experiment_config, mode, port, config_file_name):
     request_data['maxTrialNum'] = experiment_config['maxTrialNum']
     request_data['searchSpace'] = experiment_config.get('searchSpace')
     request_data['trainingServicePlatform'] = experiment_config.get('trainingServicePlatform')
-
     if experiment_config.get('description'):
         request_data['description'] = experiment_config['description']
     if experiment_config.get('multiPhase'):
@@ -332,7 +332,6 @@ def set_experiment(experiment_config, mode, port, config_file_name):
         request_data['versionCheck'] = experiment_config.get('versionCheck')
     if experiment_config.get('logCollection'):
         request_data['logCollection'] = experiment_config.get('logCollection')
-
     request_data['clusterMetaData'] = []
     if experiment_config['trainingServicePlatform'] == 'local':
         request_data['clusterMetaData'].append(
@@ -344,6 +343,11 @@ def set_experiment(experiment_config, mode, port, config_file_name):
             {'key': 'machine_list', 'value': experiment_config['machineList']})
         request_data['clusterMetaData'].append(
             {'key': 'trial_config', 'value': experiment_config['trial']})
+        if not experiment_config.get('remoteConfig'):
+            # set default value of reuse in remoteConfig to False
+            experiment_config['remoteConfig'] = {'reuse': False}
+        request_data['clusterMetaData'].append(
+            {'key': 'remote_config', 'value': experiment_config['remoteConfig']})
     elif experiment_config['trainingServicePlatform'] == 'pai':
         request_data['clusterMetaData'].append(
             {'key': 'pai_config', 'value': experiment_config['paiConfig']})

tools/nni_trial_tool/trial_runner.py

@@ -27,6 +27,10 @@ def main_loop(args):
     gpu_refresh_last_time = datetime.now() - timedelta(minutes=1)
 
     try:
+        if args.job_pid_file:
+            with open(args.job_pid_file, 'w') as job_file:
+                job_file.write("%d" % os.getpid())
+
         trials = dict()
 
         command_channel = args.command_channel
@@ -143,6 +147,7 @@ def check_version(args):
     PARSER.add_argument('--nni_manager_version', type=str, help='the nni version transmitted from nniManager')
     PARSER.add_argument('--log_collection', type=str, help='set the way to collect log in trial runner')
     PARSER.add_argument('--node_count', type=int, help='number of nodes, it determines how to consume command and save code file')
+    PARSER.add_argument('--job_pid_file', type=str, help='save trial runner process pid')
     args, unknown = PARSER.parse_known_args()
 
     setting_file = "settings.json"