Let nnictl support stopping by port

test/naive_test.py

@@ -88,6 +88,7 @@ def stop_experiment_test():
     subprocess.run(['nnictl', 'create', '--config', 'tuner_test/local.yml', '--port', '8080'], check=True)
     subprocess.run(['nnictl', 'create', '--config', 'tuner_test/local.yml', '--port', '8888'], check=True)
     subprocess.run(['nnictl', 'create', '--config', 'tuner_test/local.yml', '--port', '8989'], check=True)
+    subprocess.run(['nnictl', 'create', '--config', 'tuner_test/local.yml', '--port', '8990'], check=True)
 
     # test cmd 'nnictl stop id`
     experiment_id = get_experiment_id(EXPERIMENT_URL)
@@ -96,6 +97,12 @@ def stop_experiment_test():
     snooze()
     assert not detect_port(8080), '`nnictl stop %s` failed to stop experiments' % experiment_id
 
+    # test cmd `nnictl stop --port`
+    proc = subprocess.run(['nnictl', 'stop', '--port', '8990'])
+    assert proc.returncode == 0, '`nnictl stop %s` failed with code %d' % (experiment_id, proc.returncode)
+    snooze()
+    assert not detect_port(8990), '`nnictl stop %s` failed to stop experiments' % experiment_id
+
     # test cmd `nnictl stop all`
     proc = subprocess.run(['nnictl', 'stop', 'all'])
     assert proc.returncode == 0, '`nnictl stop all` failed with code %d' % proc.returncode

tools/nni_cmd/nnictl.py

@@ -90,6 +90,7 @@ def parse_args():
     #parse stop command
     parser_stop = subparsers.add_parser('stop', help='stop the experiment')
     parser_stop.add_argument('id', nargs='?', help='the id of experiment, use \'all\' to stop all running experiments')
+    parser_stop.add_argument('--port', '-p', dest='port', help='the port of restful server')
     parser_stop.set_defaults(func=stop_experiment)
 
     #parse trial command

tools/nni_cmd/nnictl_utils.py

@@ -118,12 +118,14 @@ def check_experiment_id(args, update=True):
 
 def parse_ids(args):
     '''Parse the arguments for nnictl stop
-    1.If there is an id specified, return the corresponding id
-    2.If there is no id specified, and there is an experiment running, return the id, or return Error
-    3.If the id matches an experiment, nnictl will return the id.
-    4.If the id ends with *, nnictl will match all ids matchs the regular
-    5.If the id does not exist but match the prefix of an experiment id, nnictl will return the matched id
-    6.If the id does not exist but match multiple prefix of the experiment ids, nnictl will give id information
+    1.If port is provided and id is not specified, return the id who owns the port
+    2.If both port and id are provided, return the id if it owns the port, otherwise fail
+    3.If there is an id specified, return the corresponding id
+    4.If there is no id specified, and there is an experiment running, return the id, or return Error
+    5.If the id matches an experiment, nnictl will return the id.
+    6.If the id ends with *, nnictl will match all ids matchs the regular
+    7.If the id does not exist but match the prefix of an experiment id, nnictl will return the matched id
+    8.If the id does not exist but match multiple prefix of the experiment ids, nnictl will give id information
     '''
     update_experiment()
     experiment_config = Experiments()
@@ -140,7 +142,14 @@ def parse_ids(args):
         elif isinstance(experiment_dict[key], list):
             # if the config file is old version, remove the configuration from file
             experiment_config.remove_experiment(key)
-    if not args.id:
+    if args.port is not None:
+        for key in running_experiment_list:
+            if str(experiment_dict[key]['port']) == args.port:
+                result_list.append(key)
+        if args.id and result_list and args.id != result_list[0]:
+            print_error('Experiment id and resful server port not match')
+            exit(1)
+    elif not args.id:
         if len(running_experiment_list) > 1:
             print_error('There are multiple experiments, please set the experiment id...')
             experiment_information = ""
@@ -166,8 +175,8 @@ def parse_ids(args):
         if len(result_list) > 1:
             print_error(args.id + ' is ambiguous, please choose ' + ' '.join(result_list) )
             return None
-    if not result_list and args.id and args.id != 'all':
-        print_error('There are no experiments matched, please set correct experiment id...')
+    if not result_list and ((args.id and args.id != 'all') or args.port):
+        print_error('There are no experiments matched, please set correct experiment id or restful server port')
     elif not result_list:
         print_error('There is no experiment running...')
     return result_list
@@ -665,4 +674,4 @@ def export_trials_data(args):
         else:
             print_error('Export failed...')
     else:
-        print_error('Restful server is not Running')
+        print_error('Restful server is not Running')