nnictl_utils.py 30.3 KB
Newer Older
Deshui Yu's avatar
Deshui Yu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Copyright (c) Microsoft Corporation
# All rights reserved.
#
# MIT License
#
# Permission is hereby granted, free of charge,
# to any person obtaining a copy of this software and associated
# documentation files (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and
# to permit persons to whom the Software is furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
# BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

Yan Ni's avatar
Yan Ni committed
21
import csv
Deshui Yu's avatar
Deshui Yu committed
22
23
24
import os
import psutil
import json
25
from datetime import datetime, timezone
26
import time
SparkSnail's avatar
SparkSnail committed
27
28
29
30
import re
from pathlib import Path
from pyhdfs import HdfsClient, HdfsFileNotFoundException
import shutil
Deshui Yu's avatar
Deshui Yu committed
31
from subprocess import call, check_output
32
from nni_annotation import expand_annotations
33
from .rest_utils import rest_get, rest_delete, check_rest_server_quick, check_response
34
from .url_utils import trial_jobs_url, experiment_url, trial_job_id_url, export_data_url
35
from .config_utils import Config, Experiments
SparkSnail's avatar
SparkSnail committed
36
from .constants import NNICTL_HOME_DIR, EXPERIMENT_INFORMATION_FORMAT, EXPERIMENT_DETAIL_FORMAT, \
37
     EXPERIMENT_MONITOR_INFO, TRIAL_MONITOR_HEAD, TRIAL_MONITOR_CONTENT, TRIAL_MONITOR_TAIL, REST_TIME_OUT
SparkSnail's avatar
SparkSnail committed
38
from .common_utils import print_normal, print_error, print_warning, detect_process, get_yml_content
39
from .command_utils import check_output_command, kill_command
SparkSnail's avatar
SparkSnail committed
40
from .ssh_utils import create_ssh_sftp_client, remove_remote_directory
Deshui Yu's avatar
Deshui Yu committed
41

42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def get_experiment_time(port):
    '''get the startTime and endTime of an experiment'''
    response = rest_get(experiment_url(port), REST_TIME_OUT)
    if response and check_response(response):
        content = convert_time_stamp_to_date(json.loads(response.text))
        return content.get('startTime'), content.get('endTime')
    return None, None

def get_experiment_status(port):
    '''get the status of an experiment'''
    result, response = check_rest_server_quick(port)
    if result:
        return json.loads(response.text).get('status')
    return None

def update_experiment():
SparkSnail's avatar
SparkSnail committed
58
59
60
61
62
63
64
    '''Update the experiment status in config file'''
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
        return None
    for key in experiment_dict.keys():
        if isinstance(experiment_dict[key], dict):
65
            if experiment_dict[key].get('status') != 'STOPPED':
SparkSnail's avatar
SparkSnail committed
66
67
68
                nni_config = Config(experiment_dict[key]['fileName'])
                rest_pid = nni_config.get_config('restServerPid')
                if not detect_process(rest_pid):
69
70
71
72
73
74
75
76
77
78
79
                    experiment_config.update_experiment(key, 'status', 'STOPPED')
                    continue
                rest_port = nni_config.get_config('restServerPort')
                startTime, endTime = get_experiment_time(rest_port)
                if startTime:
                    experiment_config.update_experiment(key, 'startTime', startTime)
                if endTime:
                    experiment_config.update_experiment(key, 'endTime', endTime)
                status = get_experiment_status(rest_port)
                if status:
                    experiment_config.update_experiment(key, 'status', status)
SparkSnail's avatar
SparkSnail committed
80

SparkSnail's avatar
SparkSnail committed
81
def check_experiment_id(args, update=True):
82
83
    '''check if the id is valid
    '''
SparkSnail's avatar
SparkSnail committed
84
85
    if update:
        update_experiment()
86
87
88
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
89
        print_normal('There is no experiment running...')
chicm-ms's avatar
chicm-ms committed
90
        return None
91
    if not args.id:
92
93
        running_experiment_list = []
        for key in experiment_dict.keys():
SparkSnail's avatar
SparkSnail committed
94
            if isinstance(experiment_dict[key], dict):
95
                if experiment_dict[key].get('status') != 'STOPPED':
SparkSnail's avatar
SparkSnail committed
96
97
98
99
                    running_experiment_list.append(key)
            elif isinstance(experiment_dict[key], list):
                # if the config file is old version, remove the configuration from file
                experiment_config.remove_experiment(key)
100
        if len(running_experiment_list) > 1:
101
            print_error('There are multiple experiments, please set the experiment id...')
102
103
104
            experiment_information = ""
            for key in running_experiment_list:
                experiment_information += (EXPERIMENT_DETAIL_FORMAT % (key, experiment_dict[key]['status'], \
SparkSnail's avatar
SparkSnail committed
105
                experiment_dict[key]['port'], experiment_dict[key].get('platform'), experiment_dict[key]['startTime'], experiment_dict[key]['endTime']))
106
107
108
            print(EXPERIMENT_INFORMATION_FORMAT % experiment_information)
            exit(1)
        elif not running_experiment_list:
SparkSnail's avatar
SparkSnail committed
109
            print_error('There is no experiment running.')
chicm-ms's avatar
chicm-ms committed
110
            return None
111
112
        else:
            return running_experiment_list[0]
113
114
115
    if experiment_dict.get(args.id):
        return args.id
    else:
SparkSnail's avatar
SparkSnail committed
116
        print_error('Id not correct.')
117
        return None
Deshui Yu's avatar
Deshui Yu committed
118

119
def parse_ids(args):
120
    '''Parse the arguments for nnictl stop
121
122
123
124
125
126
127
128
    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
129
    '''
130
    update_experiment()
131
132
133
134
135
136
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
        print_normal('Experiment is not running...')
        return None
    result_list = []
137
138
    running_experiment_list = []
    for key in experiment_dict.keys():
SparkSnail's avatar
SparkSnail committed
139
        if isinstance(experiment_dict[key], dict):
140
            if experiment_dict[key].get('status') != 'STOPPED':
SparkSnail's avatar
SparkSnail committed
141
142
143
144
                running_experiment_list.append(key)
        elif isinstance(experiment_dict[key], list):
            # if the config file is old version, remove the configuration from file
            experiment_config.remove_experiment(key)
145
146
    if args.all:
        return running_experiment_list
147
148
149
150
151
152
153
154
    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:
155
        if len(running_experiment_list) > 1:
156
            print_error('There are multiple experiments, please set the experiment id...')
157
            experiment_information = ""
158
159
            for key in running_experiment_list:
                experiment_information += (EXPERIMENT_DETAIL_FORMAT % (key, experiment_dict[key]['status'], \
SparkSnail's avatar
SparkSnail committed
160
                experiment_dict[key]['port'], experiment_dict[key].get('platform'), experiment_dict[key]['startTime'], experiment_dict[key]['endTime']))
161
162
163
164
            print(EXPERIMENT_INFORMATION_FORMAT % experiment_information)
            exit(1)
        else:
            result_list = running_experiment_list
165
    elif args.id.endswith('*'):
166
        for id in running_experiment_list:
167
168
            if id.startswith(args.id[:-1]):
                result_list.append(id)
169
    elif args.id in running_experiment_list:
170
171
        result_list.append(args.id)
    else:
172
        for id in running_experiment_list:
173
174
175
176
177
            if id.startswith(args.id):
                result_list.append(id)
        if len(result_list) > 1:
            print_error(args.id + ' is ambiguous, please choose ' + ' '.join(result_list) )
            return None
178
    if not result_list and (args.id  or args.port):
179
        print_error('There are no experiments matched, please set correct experiment id or restful server port')
SparkSnail's avatar
SparkSnail committed
180
    elif not result_list:
SparkSnail's avatar
SparkSnail committed
181
        print_error('There is no experiment running...')
182
183
    return result_list

184
185
186
def get_config_filename(args):
    '''get the file name of config file'''
    experiment_id = check_experiment_id(args)
chicm-ms's avatar
chicm-ms committed
187
    if experiment_id is None:
SparkSnail's avatar
SparkSnail committed
188
        print_error('Please set correct experiment id.')
chicm-ms's avatar
chicm-ms committed
189
        exit(1)
190
191
192
193
194
195
196
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    return experiment_dict[experiment_id]['fileName']

def get_experiment_port(args):
    '''get the port of experiment'''
    experiment_id = check_experiment_id(args)
chicm-ms's avatar
chicm-ms committed
197
    if experiment_id is None:
SparkSnail's avatar
SparkSnail committed
198
        print_error('Please set correct experiment id.')
chicm-ms's avatar
chicm-ms committed
199
        exit(1)
200
201
202
203
204
205
206
207
208
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    return experiment_dict[experiment_id]['port']

def convert_time_stamp_to_date(content):
    '''Convert time stamp to date time format'''
    start_time_stamp = content.get('startTime')
    end_time_stamp = content.get('endTime')
    if start_time_stamp:
209
        start_time = datetime.fromtimestamp(start_time_stamp // 1000, timezone.utc).astimezone().strftime("%Y/%m/%d %H:%M:%S")
210
211
        content['startTime'] = str(start_time)
    if end_time_stamp:
212
        end_time = datetime.fromtimestamp(end_time_stamp // 1000, timezone.utc).astimezone().strftime("%Y/%m/%d %H:%M:%S")
213
214
215
216
217
218
219
220
221
222
223
224
225
        content['endTime'] = str(end_time)
    return content

def check_rest(args):
    '''check if restful server is running'''
    nni_config = Config(get_config_filename(args))
    rest_port = nni_config.get_config('restServerPort')
    running, _ = check_rest_server_quick(rest_port)
    if not running:
        print_normal('Restful server is running...')
    else:
        print_normal('Restful server is not running...')

Deshui Yu's avatar
Deshui Yu committed
226
227
def stop_experiment(args):
    '''Stop the experiment which is running'''
228
229
230
    if args.id and args.id == 'all':
        print_warning('\'nnictl stop all\' is abolished, please use \'nnictl stop --all\' to stop all of experiments!')
        exit(1)
231
232
233
234
235
236
    experiment_id_list = parse_ids(args)
    if experiment_id_list:
        experiment_config = Experiments()
        experiment_dict = experiment_config.get_all_experiments()
        for experiment_id in experiment_id_list:
            print_normal('Stoping experiment %s' % experiment_id)
237
            nni_config = Config(experiment_dict[experiment_id]['fileName'])
238
239
240
            rest_port = nni_config.get_config('restServerPort')
            rest_pid = nni_config.get_config('restServerPid')
            if rest_pid:
241
                kill_command(rest_pid)
SparkSnail's avatar
SparkSnail committed
242
243
244
245
                tensorboard_pid_list = nni_config.get_config('tensorboardPidList')
                if tensorboard_pid_list:
                    for tensorboard_pid in tensorboard_pid_list:
                        try:
246
                            kill_command(tensorboard_pid)
SparkSnail's avatar
SparkSnail committed
247
248
249
                        except Exception as exception:
                            print_error(exception)
                    nni_config.set_config('tensorboardPidList', [])
SparkSnail's avatar
SparkSnail committed
250
            print_normal('Stop experiment success.')
251
            experiment_config.update_experiment(experiment_id, 'status', 'STOPPED')
252
253
            time_now = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
            experiment_config.update_experiment(experiment_id, 'endTime', str(time_now))
Deshui Yu's avatar
Deshui Yu committed
254
255
256

def trial_ls(args):
    '''List trial'''
257
    nni_config = Config(get_config_filename(args))
Deshui Yu's avatar
Deshui Yu committed
258
259
260
261
262
    rest_port = nni_config.get_config('restServerPort')
    rest_pid = nni_config.get_config('restServerPid')
    if not detect_process(rest_pid):
        print_error('Experiment is not running...')
        return
263
264
    running, response = check_rest_server_quick(rest_port)
    if running:
265
        response = rest_get(trial_jobs_url(rest_port), REST_TIME_OUT)
266
        if response and check_response(response):
267
            content = json.loads(response.text)
268
            for index, value in enumerate(content):
269
270
                content[index] = convert_time_stamp_to_date(value)
            print(json.dumps(content, indent=4, sort_keys=True, separators=(',', ':')))
Deshui Yu's avatar
Deshui Yu committed
271
272
273
274
275
276
277
        else:
            print_error('List trial failed...')
    else:
        print_error('Restful server is not running...')

def trial_kill(args):
    '''List trial'''
278
    nni_config = Config(get_config_filename(args))
Deshui Yu's avatar
Deshui Yu committed
279
280
281
282
283
    rest_port = nni_config.get_config('restServerPort')
    rest_pid = nni_config.get_config('restServerPid')
    if not detect_process(rest_pid):
        print_error('Experiment is not running...')
        return
284
285
    running, _ = check_rest_server_quick(rest_port)
    if running:
SparkSnail's avatar
SparkSnail committed
286
        response = rest_delete(trial_job_id_url(rest_port, args.trial_id), REST_TIME_OUT)
287
        if response and check_response(response):
Deshui Yu's avatar
Deshui Yu committed
288
289
290
291
292
293
            print(response.text)
        else:
            print_error('Kill trial job failed...')
    else:
        print_error('Restful server is not running...')

294
295
296
297
298
299
300
301
302
303
304
def trial_codegen(args):
    '''Generate code for a specific trial'''
    print_warning('Currently, this command is only for nni nas programming interface.')
    exp_id = check_experiment_id(args)
    nni_config = Config(get_config_filename(args))
    if not nni_config.get_config('experimentConfig')['useAnnotation']:
        print_error('The experiment is not using annotation')
        exit(1)
    code_dir = nni_config.get_config('experimentConfig')['trial']['codeDir']
    expand_annotations(code_dir, './exp_%s_trial_%s_code'%(exp_id, args.trial_id), exp_id, args.trial_id)

Deshui Yu's avatar
Deshui Yu committed
305
306
def list_experiment(args):
    '''Get experiment information'''
307
    nni_config = Config(get_config_filename(args))
Deshui Yu's avatar
Deshui Yu committed
308
309
310
311
312
    rest_port = nni_config.get_config('restServerPort')
    rest_pid = nni_config.get_config('restServerPid')
    if not detect_process(rest_pid):
        print_error('Experiment is not running...')
        return
313
314
    running, _ = check_rest_server_quick(rest_port)
    if running:
315
        response = rest_get(experiment_url(rest_port), REST_TIME_OUT)
316
        if response and check_response(response):
317
318
            content = convert_time_stamp_to_date(json.loads(response.text))
            print(json.dumps(content, indent=4, sort_keys=True, separators=(',', ':')))
Deshui Yu's avatar
Deshui Yu committed
319
320
321
322
323
        else:
            print_error('List experiment failed...')
    else:
        print_error('Restful server is not running...')

324
325
def experiment_status(args):
    '''Show the status of experiment'''
326
    nni_config = Config(get_config_filename(args))
327
328
329
330
331
332
333
    rest_port = nni_config.get_config('restServerPort')
    result, response = check_rest_server_quick(rest_port)
    if not result:
        print_normal('Restful server is not running...')
    else:
        print(json.dumps(json.loads(response.text), indent=4, sort_keys=True, separators=(',', ':')))

Deshui Yu's avatar
Deshui Yu committed
334
335
def log_internal(args, filetype):
    '''internal function to call get_log_content'''
336
    file_name = get_config_filename(args)
Deshui Yu's avatar
Deshui Yu committed
337
    if filetype == 'stdout':
338
        file_full_path = os.path.join(NNICTL_HOME_DIR, file_name, 'stdout')
Deshui Yu's avatar
Deshui Yu committed
339
    else:
340
        file_full_path = os.path.join(NNICTL_HOME_DIR, file_name, 'stderr')
341
    print(check_output_command(file_full_path, head=args.head, tail=args.tail))
342

Deshui Yu's avatar
Deshui Yu committed
343
344
345
346
347
348
349
350
def log_stdout(args):
    '''get stdout log'''
    log_internal(args, 'stdout')

def log_stderr(args):
    '''get stderr log'''
    log_internal(args, 'stderr')

351
352
353
def log_trial(args):
    ''''get trial log path'''
    trial_id_path_dict = {}
SparkSnail's avatar
SparkSnail committed
354
    trial_id_list = []
355
    nni_config = Config(get_config_filename(args))
356
357
358
359
360
361
362
    rest_port = nni_config.get_config('restServerPort')
    rest_pid = nni_config.get_config('restServerPid')
    if not detect_process(rest_pid):
        print_error('Experiment is not running...')
        return
    running, response = check_rest_server_quick(rest_port)
    if running:
363
        response = rest_get(trial_jobs_url(rest_port), REST_TIME_OUT)
364
365
366
        if response and check_response(response):
            content = json.loads(response.text)
            for trial in content:
SparkSnail's avatar
SparkSnail committed
367
368
369
                trial_id_list.append(trial.get('id'))
                if trial.get('logPath'):
                    trial_id_path_dict[trial.get('id')] = trial['logPath']
370
371
    else:
        print_error('Restful server is not running...')
goooxu's avatar
goooxu committed
372
        exit(1)
SparkSnail's avatar
SparkSnail committed
373
374
375
376
377
378
    if args.trial_id:
        if args.trial_id not in trial_id_list:
            print_error('Trial id {0} not correct, please check your command!'.format(args.trial_id))
            exit(1)
        if trial_id_path_dict.get(args.trial_id):
            print_normal('id:' + args.trial_id + ' path:' + trial_id_path_dict[args.trial_id])
379
        else:
SparkSnail's avatar
SparkSnail committed
380
            print_error('Log path is not available yet, please wait...')
goooxu's avatar
goooxu committed
381
            exit(1)
382
    else:
SparkSnail's avatar
SparkSnail committed
383
        print_normal('All of trial log info:')
384
        for key in trial_id_path_dict:
SparkSnail's avatar
SparkSnail committed
385
386
387
            print_normal('id:' + key + ' path:' + trial_id_path_dict[key])
        if not trial_id_path_dict:
            print_normal('None')
388

Deshui Yu's avatar
Deshui Yu committed
389
390
def get_config(args):
    '''get config info'''
391
    nni_config = Config(get_config_filename(args))
Deshui Yu's avatar
Deshui Yu committed
392
    print(nni_config.get_all_config())
393
394
395

def webui_url(args):
    '''show the url of web ui'''
396
    nni_config = Config(get_config_filename(args))
397
398
    print_normal('{0} {1}'.format('Web UI url:', ' '.join(nni_config.get_config('webuiUrl'))))

SparkSnail's avatar
SparkSnail committed
399
400
401
402
403
404
def local_clean(directory):
    '''clean up local data'''
    print_normal('removing folder {0}'.format(directory))
    try:
        shutil.rmtree(directory)
    except FileNotFoundError as err:
SparkSnail's avatar
SparkSnail committed
405
        print_error('{0} does not exist.'.format(directory))
SparkSnail's avatar
SparkSnail committed
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
    
def remote_clean(machine_list, experiment_id=None):
    '''clean up remote data'''
    for machine in machine_list:
        passwd = machine.get('passwd')
        userName = machine.get('username')
        host = machine.get('ip')
        port = machine.get('port')
        if experiment_id:
            remote_dir = '/' + '/'.join(['tmp', 'nni', 'experiments', experiment_id])
        else:
            remote_dir = '/' + '/'.join(['tmp', 'nni', 'experiments'])
        sftp = create_ssh_sftp_client(host, port, userName, passwd)
        print_normal('removing folder {0}'.format(host + ':' + str(port) + remote_dir))
        remove_remote_directory(sftp, remote_dir)
    
def hdfs_clean(host, user_name, output_dir, experiment_id=None):
    '''clean up hdfs data'''
    hdfs_client = HdfsClient(hosts='{0}:80'.format(host), user_name=user_name, webhdfs_path='/webhdfs/api/v1', timeout=5)
    if experiment_id:
        full_path = '/' + '/'.join([user_name, 'nni', 'experiments', experiment_id])
    else:
        full_path = '/' + '/'.join([user_name, 'nni', 'experiments'])
    print_normal('removing folder {0} in hdfs'.format(full_path))
    hdfs_client.delete(full_path, recursive=True)
    if output_dir:
        pattern = re.compile('hdfs://(?P<host>([0-9]{1,3}.){3}[0-9]{1,3})(:[0-9]{2,5})?(?P<baseDir>/.*)?')
        match_result = pattern.match(output_dir)
        if match_result:
            output_host = match_result.group('host')
            output_dir = match_result.group('baseDir')
            #check if the host is valid
            if output_host != host:
                print_warning('The host in {0} is not consistent with {1}'.format(output_dir, host))
            else:
                if experiment_id:
                    output_dir = output_dir + '/' + experiment_id
                print_normal('removing folder {0} in hdfs'.format(output_dir))
                hdfs_client.delete(output_dir, recursive=True)

def experiment_clean(args):
    '''clean up the experiment data'''
    experiment_id_list = []
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if args.all:
        experiment_id_list = list(experiment_dict.keys())
    else:
        if args.id is None:
SparkSnail's avatar
SparkSnail committed
455
            print_error('please set experiment id.')
SparkSnail's avatar
SparkSnail committed
456
457
            exit(1)
        if args.id not in experiment_dict:
SparkSnail's avatar
SparkSnail committed
458
            print_error('Cannot find experiment {0}.'.format(args.id))
SparkSnail's avatar
SparkSnail committed
459
460
461
            exit(1)
        experiment_id_list.append(args.id)
    while True:
SparkSnail's avatar
SparkSnail committed
462
        print('INFO: This action will delete experiment {0}, and it\'s not recoverable.'.format(' '.join(experiment_id_list)))
SparkSnail's avatar
SparkSnail committed
463
464
465
466
        inputs = input('INFO: do you want to continue?[y/N]:')
        if not inputs.lower() or inputs.lower() in ['n', 'no']:
            exit(0)
        elif inputs.lower() not in ['y', 'n', 'yes', 'no']:
SparkSnail's avatar
SparkSnail committed
467
            print_warning('please input Y or N.')
SparkSnail's avatar
SparkSnail committed
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
        else:
            break
    for experiment_id in experiment_id_list:
        nni_config = Config(experiment_dict[experiment_id]['fileName'])
        platform = nni_config.get_config('experimentConfig').get('trainingServicePlatform')
        experiment_id = nni_config.get_config('experimentId')
        if platform == 'remote':
            machine_list = nni_config.get_config('experimentConfig').get('machineList')
            remote_clean(machine_list, experiment_id)
        elif platform == 'pai':
            host = nni_config.get_config('experimentConfig').get('paiConfig').get('host')	
            user_name = nni_config.get_config('experimentConfig').get('paiConfig').get('userName')
            output_dir = nni_config.get_config('experimentConfig').get('trial').get('outputDir')
            hdfs_clean(host, user_name, output_dir, experiment_id)
        elif platform != 'local':
            #TODO: support all platforms
SparkSnail's avatar
SparkSnail committed
484
            print_warning('platform {0} clean up not supported yet.'.format(platform))
SparkSnail's avatar
SparkSnail committed
485
486
487
488
489
490
491
492
493
494
            exit(0)
        #clean local data
        home = str(Path.home())
        local_dir = nni_config.get_config('experimentConfig').get('logDir')
        if not local_dir:
            local_dir = os.path.join(home, 'nni', 'experiments', experiment_id)
        local_clean(local_dir)
        experiment_config = Experiments()
        print_normal('removing metadata of experiment {0}'.format(experiment_id))
        experiment_config.remove_experiment(experiment_id)
SparkSnail's avatar
SparkSnail committed
495
        print_normal('Done.') 
SparkSnail's avatar
SparkSnail committed
496
497
498
499
500
501
502
503
504
505

def get_platform_dir(config_content):
    '''get the dir list to be deleted'''
    platform = config_content.get('trainingServicePlatform')
    dir_list = []
    if platform == 'remote':
        machine_list = config_content.get('machineList')
        for machine in machine_list:
            host = machine.get('ip')
            port = machine.get('port')
SparkSnail's avatar
SparkSnail committed
506
            dir_list.append(host + ':' + str(port) + '/tmp/nni')
SparkSnail's avatar
SparkSnail committed
507
508
509
510
511
    elif platform == 'pai':
        pai_config = config_content.get('paiConfig')
        host = config_content.get('paiConfig').get('host')	
        user_name = config_content.get('paiConfig').get('userName')
        output_dir = config_content.get('trial').get('outputDir')
SparkSnail's avatar
SparkSnail committed
512
        dir_list.append('server: {0}, path: {1}/nni'.format(host, user_name))
SparkSnail's avatar
SparkSnail committed
513
514
515
516
517
518
519
520
        if output_dir:
            dir_list.append(output_dir)
    return dir_list

def platform_clean(args):
    '''clean up the experiment data'''
    config_path = os.path.abspath(args.config)
    if not os.path.exists(config_path):
SparkSnail's avatar
SparkSnail committed
521
        print_error('Please set correct config path.')
SparkSnail's avatar
SparkSnail committed
522
523
524
        exit(1)
    config_content = get_yml_content(config_path)
    platform = config_content.get('trainingServicePlatform')
SparkSnail's avatar
SparkSnail committed
525
526
527
    if platform == 'local':
        print_normal('it doesn’t need to clean local platform.')
        exit(0)
SparkSnail's avatar
SparkSnail committed
528
    if platform not in ['remote', 'pai']:
SparkSnail's avatar
SparkSnail committed
529
        print_normal('platform {0} not supported.'.format(platform))
SparkSnail's avatar
SparkSnail committed
530
        exit(0)
Shinai Yang's avatar
Shinai Yang committed
531
    update_experiment()
SparkSnail's avatar
SparkSnail committed
532
533
534
535
536
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    id_list = list(experiment_dict.keys())
    dir_list = get_platform_dir(config_content)
    if not dir_list:
SparkSnail's avatar
SparkSnail committed
537
        print_normal('No folder of NNI caches is found.')
SparkSnail's avatar
SparkSnail committed
538
539
540
541
542
543
544
545
546
        exit(1)
    while True:
        print_normal('This command will remove below folders of NNI caches. If other users are using experiments on below hosts, it will be broken.')
        for dir in dir_list:
            print('       ' + dir)
        inputs = input('INFO: do you want to continue?[y/N]:')
        if not inputs.lower() or inputs.lower() in ['n', 'no']:
            exit(0)
        elif inputs.lower() not in ['y', 'n', 'yes', 'no']:
SparkSnail's avatar
SparkSnail committed
547
            print_warning('please input Y or N.')
SparkSnail's avatar
SparkSnail committed
548
549
550
551
552
553
554
555
556
557
558
559
        else:
            break
    if platform == 'remote':
        machine_list = config_content.get('machineList')
        for machine in machine_list:
            remote_clean(machine_list, None)
    elif platform == 'pai':
        pai_config = config_content.get('paiConfig')
        host = config_content.get('paiConfig').get('host')	
        user_name = config_content.get('paiConfig').get('userName')
        output_dir = config_content.get('trial').get('outputDir')
        hdfs_clean(host, user_name, output_dir, None)
SparkSnail's avatar
SparkSnail committed
560
    print_normal('Done.')
SparkSnail's avatar
SparkSnail committed
561

562
563
def experiment_list(args):
    '''get the information of all experiments'''
Shinai Yang's avatar
Shinai Yang committed
564
    update_experiment()
565
566
567
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
SparkSnail's avatar
SparkSnail committed
568
        print_normal('Cannot find experiments.')
569
570
        exit(1)
    experiment_id_list = []
SparkSnail's avatar
SparkSnail committed
571
    if args.all:
572
573
        for key in experiment_dict.keys():
            experiment_id_list.append(key)
574
575
    else:
        for key in experiment_dict.keys():
576
            if experiment_dict[key]['status'] != 'STOPPED':
577
578
                experiment_id_list.append(key)
        if not experiment_id_list:
579
            print_warning('There is no experiment running...\nYou can use \'nnictl experiment list --all\' to list all experiments.')
580
581
    experiment_information = ""
    for key in experiment_id_list:
582
        experiment_information += (EXPERIMENT_DETAIL_FORMAT % (key, experiment_dict[key]['status'], experiment_dict[key]['port'],\
SparkSnail's avatar
SparkSnail committed
583
        experiment_dict[key].get('platform'), experiment_dict[key]['startTime'], experiment_dict[key]['endTime']))
584
    print(EXPERIMENT_INFORMATION_FORMAT % experiment_information)
SparkSnail's avatar
SparkSnail committed
585

SparkSnail's avatar
SparkSnail committed
586
587
588
589
def get_time_interval(time1, time2):
    '''get the interval of two times'''
    try:
        #convert time to timestamp
590
591
        time1 = time.mktime(time.strptime(time1, '%Y/%m/%d %H:%M:%S'))
        time2 = time.mktime(time.strptime(time2, '%Y/%m/%d %H:%M:%S'))
592
        seconds = (datetime.fromtimestamp(time2) - datetime.fromtimestamp(time1)).seconds
SparkSnail's avatar
SparkSnail committed
593
594
595
596
597
598
599
600
601
602
603
604
605
        #convert seconds to day:hour:minute:second
        days = seconds / 86400
        seconds %= 86400
        hours = seconds / 3600
        seconds %= 3600
        minutes = seconds / 60
        seconds %= 60
        return '%dd %dh %dm %ds' % (days, hours, minutes, seconds)
    except:
        return 'N/A'

def show_experiment_info():
    '''show experiment information in monitor'''
Shinai Yang's avatar
Shinai Yang committed
606
    update_experiment()
SparkSnail's avatar
SparkSnail committed
607
608
609
610
611
612
613
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
        print('There is no experiment running...')
        exit(1)
    experiment_id_list = []
    for key in experiment_dict.keys():
614
        if experiment_dict[key]['status'] != 'STOPPED':
SparkSnail's avatar
SparkSnail committed
615
616
617
618
619
620
            experiment_id_list.append(key)
    if not experiment_id_list:
        print_warning('There is no experiment running...')
        return
    for key in experiment_id_list:
        print(EXPERIMENT_MONITOR_INFO % (key, experiment_dict[key]['status'], experiment_dict[key]['port'], \
621
             experiment_dict[key].get('platform'), experiment_dict[key]['startTime'], get_time_interval(experiment_dict[key]['startTime'], experiment_dict[key]['endTime'])))
SparkSnail's avatar
SparkSnail committed
622
623
624
        print(TRIAL_MONITOR_HEAD)
        running, response = check_rest_server_quick(experiment_dict[key]['port'])
        if running:
625
            response = rest_get(trial_jobs_url(experiment_dict[key]['port']), REST_TIME_OUT)
SparkSnail's avatar
SparkSnail committed
626
627
            if response and check_response(response):
                content = json.loads(response.text)
628
                for index, value in enumerate(content):
SparkSnail's avatar
SparkSnail committed
629
630
631
632
633
634
635
636
637
638
639
640
                    content[index] = convert_time_stamp_to_date(value)
                    print(TRIAL_MONITOR_CONTENT % (content[index].get('id'), content[index].get('startTime'), content[index].get('endTime'), content[index].get('status')))
        print(TRIAL_MONITOR_TAIL)

def monitor_experiment(args):
    '''monitor the experiment'''
    if args.time <= 0:
        print_error('please input a positive integer as time interval, the unit is second.')
        exit(1)
    while True:
        try:
            os.system('clear')
641
            update_experiment()
SparkSnail's avatar
SparkSnail committed
642
643
644
645
646
647
648
            show_experiment_info()
            time.sleep(args.time)
        except KeyboardInterrupt:
            exit(0)
        except Exception as exception:
            print_error(exception)
            exit(1)
Yan Ni's avatar
Yan Ni committed
649
650

def export_trials_data(args):
651
652
    '''export experiment metadata to csv
    '''
Yan Ni's avatar
Yan Ni committed
653
654
655
656
657
658
659
660
    nni_config = Config(get_config_filename(args))
    rest_port = nni_config.get_config('restServerPort')
    rest_pid = nni_config.get_config('restServerPid')
    if not detect_process(rest_pid):
        print_error('Experiment is not running...')
        return
    running, response = check_rest_server_quick(rest_port)
    if running:
661
        response = rest_get(export_data_url(rest_port), 20)
Yan Ni's avatar
Yan Ni committed
662
        if response is not None and check_response(response):
663
            if args.type == 'json':
664
665
666
667
668
669
670
671
672
673
674
675
676
                with open(args.path, 'w') as file:
                    file.write(response.text)
            elif args.type == 'csv':
                content = json.loads(response.text)
                trial_records = []
                for record in content:
                    if not isinstance(record['value'], (float, int)):
                        formated_record = {**record['parameter'], **record['value'], **{'id': record['id']}}
                    else:
                        formated_record = {**record['parameter'], **{'reward': record['value'], 'id': record['id']}}
                    trial_records.append(formated_record)
                with open(args.path, 'w') as file:
                    writer = csv.DictWriter(file, set.union(*[set(r.keys()) for r in trial_records]))
677
                    writer.writeheader()
678
679
680
681
                    writer.writerows(trial_records)
            else:
                print_error('Unknown type: %s' % args.type)
                exit(1)
Yan Ni's avatar
Yan Ni committed
682
683
684
        else:
            print_error('Export failed...')
    else:
685
        print_error('Restful server is not Running')