Fix version check

tools/nni_trial_tool/trial_keeper.py

@@ -109,22 +109,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 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'''