Merge pull request #147 from Microsoft/master

merge master

tools/nni_trial_tool/trial_keeper.py

@@ -44,10 +44,9 @@ def main_loop(args):
     
     stdout_file = open(STDOUT_FULL_PATH, 'a+')
     stderr_file = open(STDERR_FULL_PATH, 'a+')
-    
-    trial_keeper_syslogger = RemoteLogger(args.nnimanager_ip, args.nnimanager_port, 'trial_keeper', StdOutputType.Stdout)
+    trial_keeper_syslogger = RemoteLogger(args.nnimanager_ip, args.nnimanager_port, 'trial_keeper', StdOutputType.Stdout, args.log_collection)
     # redirect trial keeper's stdout and stderr to syslog
-    trial_syslogger_stdout = RemoteLogger(args.nnimanager_ip, args.nnimanager_port, 'trial', StdOutputType.Stdout)
+    trial_syslogger_stdout = RemoteLogger(args.nnimanager_ip, args.nnimanager_port, 'trial', StdOutputType.Stdout, args.log_collection)
     sys.stdout = sys.stderr = trial_keeper_syslogger
     # backward compatibility
     hdfs_host = None
@@ -109,22 +108,25 @@ def check_version(args):
         trial_keeper_version = pkg_resources.get_distribution('nni').version
     except pkg_resources.ResolutionError as err:
         #package nni does not exist, try nni-tool package
-        nni_log(LogType.Warning, 'Package nni does not exist!')
-        try:
-            trial_keeper_version = pkg_resources.get_distribution('nni-tool').version
-        except pkg_resources.ResolutionError as err:
-            #package nni-tool does not exist
-            nni_log(LogType.Error, 'Package nni-tool does not exist!')
-            os._exit(1)
+        nni_log(LogType.Error, 'Package nni does not exist!')
+        os._exit(1)
     if not args.version:
         # skip version check
         nni_log(LogType.Warning, 'Skipping version check!')
-    elif trial_keeper_version != args.version:
-        nni_log(LogType.Error, 'Exit trial keeper, trial keeper version is {}, and trainingService version is {}, \
-        versions does not match, please check your code and image versions!'.format(trial_keeper_version, args.version))
-        os._exit(1)
     else:
-        nni_log(LogType.Info,  'NNI version is {}'.format(args.version))
+        regular = re.compile('v?(?P<version>[0-9](\.[0-9]){0,2}).*')
+        try:
+            trial_keeper_version = regular.search(trial_keeper_version).group('version')
+            nni_log(LogType.Info, 'trial_keeper_version is {0}'.format(trial_keeper_version))
+            training_service_version = regular.search(args.version).group('version')
+            nni_log(LogType.Info, 'training_service_version is {0}'.format(training_service_version))
+            if trial_keeper_version != training_service_version:
+                nni_log(LogType.Error, 'Version does not match!')
+                os._exit(1)
+            else:
+                nni_log(LogType.Info, 'Version match!')
+        except AttributeError as err:
+            nni_log(LogType.Error, err)
 
 if __name__ == '__main__':
     '''NNI Trial Keeper main function'''
@@ -141,6 +143,7 @@ if __name__ == '__main__':
     PARSER.add_argument('--nni_hdfs_exp_dir', type=str, help='nni experiment directory in hdfs')
     PARSER.add_argument('--webhdfs_path', type=str, help='the webhdfs path used in webhdfs URL')
     PARSER.add_argument('--version', type=str, help='the nni version transmitted from trainingService')
+    PARSER.add_argument('--log_collection', type=str, help='set the way to collect log in trialkeeper')
     args, unknown = PARSER.parse_known_args()
     if args.trial_command is None:
         exit(1)