Asynchronous dispatcher (#372)

* Asynchronous dispatcher

* updates

* updates

* updates

* updates

src/nni_manager/common/manager.ts

@@ -34,6 +34,7 @@ interface ExperimentParams {
     searchSpace: string;
     trainingServicePlatform: string;
     multiPhase?: boolean;
+    multiThread?: boolean;
     tuner: {
         className: string;
         builtinTunerName?: string;

src/nni_manager/common/utils.ts

@@ -158,12 +158,16 @@ function parseArg(names: string[]): string {
  * @param assessor: similiar as tuner
  *
  */
-function getMsgDispatcherCommand(tuner: any, assessor: any, multiPhase: boolean = false): string {
+function getMsgDispatcherCommand(tuner: any, assessor: any, multiPhase: boolean = false, multiThread: boolean = false): string {
     let command: string = `python3 -m nni --tuner_class_name ${tuner.className}`;
     if (multiPhase) {
         command += ' --multi_phase';
     }
 
+    if (multiThread) {
+        command += ' --multi_thread';
+    }
+
     if (tuner.classArgs !== undefined) {
         command += ` --tuner_args ${JSON.stringify(JSON.stringify(tuner.classArgs))}`;
     }

src/nni_manager/core/commands.ts

@@ -26,6 +26,7 @@ const ADD_CUSTOMIZED_TRIAL_JOB = 'AD';
 const TRIAL_END = 'EN';
 const TERMINATE = 'TE';
 
+const INITIALIZED = 'ID';
 const NEW_TRIAL_JOB = 'TR';
 const SEND_TRIAL_JOB_PARAMETER = 'SP';
 const NO_MORE_TRIAL_JOBS = 'NO';
@@ -39,6 +40,7 @@ const TUNER_COMMANDS: Set<string> = new Set([
     ADD_CUSTOMIZED_TRIAL_JOB,
     TERMINATE,
 
+    INITIALIZED,
     NEW_TRIAL_JOB,
     SEND_TRIAL_JOB_PARAMETER,
     NO_MORE_TRIAL_JOBS
@@ -61,6 +63,7 @@ export {
     ADD_CUSTOMIZED_TRIAL_JOB,
     TRIAL_END,
     TERMINATE,
+    INITIALIZED,
     NEW_TRIAL_JOB,
     NO_MORE_TRIAL_JOBS,
     KILL_TRIAL_JOB,

src/nni_manager/core/nnimanager.ts

@@ -37,8 +37,8 @@ import {
 } from '../common/trainingService';
 import { delay, getLogDir, getMsgDispatcherCommand } from '../common/utils';
 import {
-    ADD_CUSTOMIZED_TRIAL_JOB, KILL_TRIAL_JOB, NEW_TRIAL_JOB, NO_MORE_TRIAL_JOBS, REPORT_METRIC_DATA,
-    REQUEST_TRIAL_JOBS, SEND_TRIAL_JOB_PARAMETER, TERMINATE, TRIAL_END, UPDATE_SEARCH_SPACE
+    ADD_CUSTOMIZED_TRIAL_JOB, INITIALIZE, INITIALIZED, KILL_TRIAL_JOB, NEW_TRIAL_JOB, NO_MORE_TRIAL_JOBS,
+    REPORT_METRIC_DATA, REQUEST_TRIAL_JOBS, SEND_TRIAL_JOB_PARAMETER, TERMINATE, TRIAL_END, UPDATE_SEARCH_SPACE
 } from './commands';
 import { createDispatcherInterface, IpcInterface } from './ipcInterface';
 
@@ -127,7 +127,8 @@ class NNIManager implements Manager {
             this.trainingService.setClusterMetadata('multiPhase', expParams.multiPhase.toString());
         }
 
-        const dispatcherCommand: string = getMsgDispatcherCommand(expParams.tuner, expParams.assessor, expParams.multiPhase);
+        const dispatcherCommand: string = getMsgDispatcherCommand(
+            expParams.tuner, expParams.assessor, expParams.multiPhase, expParams.multiThread);
         this.log.debug(`dispatcher command: ${dispatcherCommand}`);
         this.setupTuner(
             //expParams.tuner.tunerCommand,
@@ -159,7 +160,8 @@ class NNIManager implements Manager {
             this.trainingService.setClusterMetadata('multiPhase', expParams.multiPhase.toString());
         }
 
-        const dispatcherCommand: string = getMsgDispatcherCommand(expParams.tuner, expParams.assessor, expParams.multiPhase);
+        const dispatcherCommand: string = getMsgDispatcherCommand(
+            expParams.tuner, expParams.assessor, expParams.multiPhase, expParams.multiThread);
         this.log.debug(`dispatcher command: ${dispatcherCommand}`);
         this.setupTuner(
             dispatcherCommand,
@@ -419,16 +421,20 @@ class NNIManager implements Manager {
             } else {
                 this.trialConcurrencyChange = requestTrialNum;
             }
-            for (let i: number = 0; i < requestTrialNum; i++) {
+
+            const requestCustomTrialNum: number = Math.min(requestTrialNum, this.customizedTrials.length);
+            for (let i: number = 0; i < requestCustomTrialNum; i++) {
                 // ask tuner for more trials
                 if (this.customizedTrials.length > 0) {
                     const hyperParams: string | undefined = this.customizedTrials.shift();
                     this.dispatcher.sendCommand(ADD_CUSTOMIZED_TRIAL_JOB, hyperParams);
-                } else {
-                    this.dispatcher.sendCommand(REQUEST_TRIAL_JOBS, '1');
                 }
             }
 
+            if (requestTrialNum - requestCustomTrialNum > 0) {
+                this.requestTrialJobs(requestTrialNum - requestCustomTrialNum);
+            }
+
             // check maxtrialnum and maxduration here
             if (this.experimentProfile.execDuration > this.experimentProfile.params.maxExecDuration ||
                 this.currSubmittedTrialNum >= this.experimentProfile.params.maxTrialNum) {
@@ -526,11 +532,9 @@ class NNIManager implements Manager {
         if (this.dispatcher === undefined) {
             throw new Error('Dispatcher error: tuner has not been setup');
         }
-        // TO DO: we should send INITIALIZE command to tuner if user's tuner needs to run init method in tuner
-        this.log.debug(`Send tuner command: update search space: ${this.experimentProfile.params.searchSpace}`);
-        this.dispatcher.sendCommand(UPDATE_SEARCH_SPACE, this.experimentProfile.params.searchSpace);
-        this.log.debug(`Send tuner command: ${this.experimentProfile.params.trialConcurrency}`);
-        this.dispatcher.sendCommand(REQUEST_TRIAL_JOBS, String(this.experimentProfile.params.trialConcurrency));
+        this.log.debug(`Send tuner command: INITIALIZE: ${this.experimentProfile.params.searchSpace}`);
+        // Tuner need to be initialized with search space before generating any hyper parameters
+        this.dispatcher.sendCommand(INITIALIZE, this.experimentProfile.params.searchSpace);
     }
 
     private async onTrialJobMetrics(metric: TrialJobMetric): Promise<void> {
@@ -541,9 +545,32 @@ class NNIManager implements Manager {
         this.dispatcher.sendCommand(REPORT_METRIC_DATA, metric.data);
     }
 
+    private requestTrialJobs(jobNum: number): void {
+        if (jobNum < 1) {
+            return;
+        }
+        if (this.dispatcher === undefined) {
+            throw new Error('Dispatcher error: tuner has not been setup');
+        }
+        if (this.experimentProfile.params.multiThread) {
+            // Send multiple requests to ensure multiple hyper parameters are generated in non-blocking way.
+            // For a single REQUEST_TRIAL_JOBS request, hyper parameters are generated one by one
+            // sequentially.
+            for (let i: number = 0; i < jobNum; i++) {
+                this.dispatcher.sendCommand(REQUEST_TRIAL_JOBS, '1');
+            }
+        } else {
+            this.dispatcher.sendCommand(REQUEST_TRIAL_JOBS, String(jobNum));
+        }
+    }
+
     private async onTunerCommand(commandType: string, content: string): Promise<void> {
         this.log.info(`Command from tuner: ${commandType}, ${content}`);
         switch (commandType) {
+            case INITIALIZED:
+                // Tuner is intialized, search space is set, request tuner to generate hyper parameters
+                this.requestTrialJobs(this.experimentProfile.params.trialConcurrency);
+                break;
             case NEW_TRIAL_JOB:
                 this.waitingTrials.push(content);
                 break;

src/nni_manager/rest_server/restValidationSchemas.ts

@@ -68,6 +68,7 @@ export namespace ValidationSchemas {
             searchSpace: joi.string().required(),
             maxExecDuration: joi.number().min(0).required(),
             multiPhase: joi.boolean(),
+            multiThread: joi.boolean(),
             tuner: joi.object({
                 builtinTunerName: joi.string().valid('TPE', 'Random', 'Anneal', 'Evolution', 'SMAC', 'BatchTuner', 'GridSearch'),
                 codeDir: joi.string(),

src/sdk/pynni/nni/__main__.py

@@ -28,6 +28,7 @@ import json
 import importlib
 
 from .constants import ModuleName, ClassName, ClassArgs
+from nni.common import enable_multi_thread
 from nni.msg_dispatcher import MsgDispatcher
 from nni.multi_phase.multi_phase_dispatcher import MultiPhaseMsgDispatcher
 logger = logging.getLogger('nni.main')
@@ -91,6 +92,7 @@ def parse_args():
     parser.add_argument('--assessor_class_filename', type=str, required=False,
                         help='Assessor class file path')
     parser.add_argument('--multi_phase', action='store_true')
+    parser.add_argument('--multi_thread', action='store_true')
 
     flags, _ = parser.parse_known_args()
     return flags
@@ -101,6 +103,8 @@ def main():
     '''
 
     args = parse_args()
+    if args.multi_thread:
+        enable_multi_thread()
 
     tuner = None
     assessor = None

src/sdk/pynni/nni/common.py

@@ -78,3 +78,12 @@ def init_logger(logger_file_path):
     logging.getLogger('matplotlib').setLevel(logging.INFO)
 
     sys.stdout = _LoggerFileWrapper(logger_file)
+
+_multi_thread = False
+
+def enable_multi_thread():
+    global _multi_thread
+    _multi_thread = True
+
+def multi_thread_enabled():
+    return _multi_thread

src/sdk/pynni/nni/msg_dispatcher.py

@@ -21,6 +21,7 @@
 import logging
 from collections import defaultdict
 import json_tricks
+import threading
 
 from .protocol import CommandType, send
 from .msg_dispatcher_base import MsgDispatcherBase
@@ -69,7 +70,7 @@ def _pack_parameter(parameter_id, params, customized=False):
 
 class MsgDispatcher(MsgDispatcherBase):
     def __init__(self, tuner, assessor=None):
-        super()
+        super().__init__()
         self.tuner = tuner
         self.assessor = assessor
         if assessor is None:
@@ -85,6 +86,14 @@ class MsgDispatcher(MsgDispatcherBase):
         if self.assessor is not None:
             self.assessor.save_checkpoint()
 
+    def handle_initialize(self, data):
+        '''
+        data is search space
+        '''
+        self.tuner.update_search_space(data)
+        send(CommandType.Initialized, '')
+        return True
+
     def handle_request_trial_jobs(self, data):
         # data: number or trial jobs
         ids = [_create_parameter_id() for _ in range(data)]

src/sdk/pynni/nni/msg_dispatcher_base.py

@@ -22,8 +22,8 @@
 import os
 import logging
 import json_tricks
-
-from .common import init_logger
+from multiprocessing.dummy import Pool as ThreadPool
+from .common import init_logger, multi_thread_enabled
 from .recoverable import Recoverable
 from .protocol import CommandType, receive
 
@@ -31,6 +31,10 @@ init_logger('dispatcher.log')
 _logger = logging.getLogger(__name__)
 
 class MsgDispatcherBase(Recoverable):
+    def __init__(self):
+        if multi_thread_enabled():
+            self.pool = ThreadPool()
+
     def run(self):
         """Run the tuner.
         This function will never return unless raise.
@@ -39,17 +43,24 @@ class MsgDispatcherBase(Recoverable):
         if mode == 'resume':
             self.load_checkpoint()
 
-        while self.handle_request():
-            pass
-
-        _logger.info('Terminated by NNI manager')
-
-    def handle_request(self):
+        while True:
             _logger.debug('waiting receive_message')
-
             command, data = receive()
             if command is None:
-            return False
+                break
+            if multi_thread_enabled():
+                self.pool.map_async(self.handle_request, [(command, data)])
+            else:
+                self.handle_request((command, data))
+
+        if multi_thread_enabled():
+            self.pool.close()
+            self.pool.join()
+
+        _logger.info('Terminated by NNI manager')
+
+    def handle_request(self, request):
+        command, data = request
 
         _logger.debug('handle request: command: [{}], data: [{}]'.format(command, data))
 
@@ -60,6 +71,7 @@ class MsgDispatcherBase(Recoverable):
 
         command_handlers = {
             # Tunner commands:
+            CommandType.Initialize: self.handle_initialize,
             CommandType.RequestTrialJobs: self.handle_request_trial_jobs,
             CommandType.UpdateSearchSpace: self.handle_update_search_space,
             CommandType.AddCustomizedTrialJob: self.handle_add_customized_trial,
@@ -74,6 +86,9 @@ class MsgDispatcherBase(Recoverable):
 
         return command_handlers[command](data)
 
+    def handle_initialize(self, data):
+        raise NotImplementedError('handle_initialize not implemented')
+
     def handle_request_trial_jobs(self, data):
         raise NotImplementedError('handle_request_trial_jobs not implemented')
 

src/sdk/pynni/nni/multi_phase/multi_phase_dispatcher.py

@@ -91,6 +91,14 @@ class MultiPhaseMsgDispatcher(MsgDispatcherBase):
         if self.assessor is not None:
             self.assessor.save_checkpoint()
 
+    def handle_initialize(self, data):
+        '''
+        data is search space
+        '''
+        self.tuner.update_search_space(data)
+        send(CommandType.Initialized, '')
+        return True
+
     def handle_request_trial_jobs(self, data):
         # data: number or trial jobs
         ids = [_create_parameter_id() for _ in range(data)]

src/sdk/pynni/nni/protocol.py

@@ -19,7 +19,9 @@
 # ==================================================================================================
 
 import logging
+import threading
 from enum import Enum
+from .common import multi_thread_enabled
 
 
 class CommandType(Enum):
@@ -33,6 +35,7 @@ class CommandType(Enum):
     Terminate = b'TE'
 
     # out
+    Initialized = b'ID'
     NewTrialJob = b'TR'
     SendTrialJobParameter = b'SP'
     NoMoreTrialJobs = b'NO'
@@ -42,6 +45,7 @@ class CommandType(Enum):
 try:
     _in_file = open(3, 'rb')
     _out_file = open(4, 'wb')
+    _lock = threading.Lock()
 except OSError:
     _msg = 'IPC pipeline not exists, maybe you are importing tuner/assessor from trial code?'
     import logging
@@ -53,12 +57,19 @@ def send(command, data):
     command: CommandType object.
     data: string payload.
     """
+    global _lock
+    try:
+        if multi_thread_enabled():
+            _lock.acquire()
         data = data.encode('utf8')
         assert len(data) < 1000000, 'Command too long'
         msg = b'%b%06d%b' % (command.value, len(data), data)
         logging.getLogger(__name__).debug('Sending command, data: [%s]' % msg)
         _out_file.write(msg)
         _out_file.flush()
+    finally:
+        if multi_thread_enabled():
+            _lock.release()
 
 
 def receive():

tools/nni_cmd/config_schema.py

@@ -31,6 +31,7 @@ Optional('maxTrialNum'): And(int, lambda x: 1 <= x <= 99999),
 'trainingServicePlatform': And(str, lambda x: x in ['remote', 'local', 'pai', 'kubeflow']),
 Optional('searchSpacePath'): os.path.exists,
 Optional('multiPhase'): bool,
+Optional('multiThread'): bool,
 'useAnnotation': bool,
 'tuner': Or({
     'builtinTunerName': Or('TPE', 'Random', 'Anneal', 'Evolution', 'SMAC', 'BatchTuner', 'GridSearch'),

tools/nni_cmd/launcher.py

@@ -196,6 +196,8 @@ def set_experiment(experiment_config, mode, port, config_file_name):
         request_data['description'] = experiment_config['description']
     if experiment_config.get('multiPhase'):
         request_data['multiPhase'] = experiment_config.get('multiPhase')
+    if experiment_config.get('multiThread'):
+        request_data['multiThread'] = experiment_config.get('multiThread')
     request_data['tuner'] = experiment_config['tuner']
     if 'assessor' in experiment_config:
         request_data['assessor'] = experiment_config['assessor']