nnictl_utils.py 23.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
import datetime
26
import time
Deshui Yu's avatar
Deshui Yu committed
27
from subprocess import call, check_output
28
from .rest_utils import rest_get, rest_delete, check_rest_server_quick, check_response
SparkSnail's avatar
SparkSnail committed
29
30
from pyhdfs import HdfsClient, HdfsFileNotFoundException
from .config_utils import Config, Experiments, HDFSConfig
Deshui Yu's avatar
Deshui Yu committed
31
from .url_utils import trial_jobs_url, experiment_url, trial_job_id_url
SparkSnail's avatar
SparkSnail committed
32
from .constants import NNICTL_HOME_DIR, EXPERIMENT_INFORMATION_FORMAT, EXPERIMENT_DETAIL_FORMAT, \
33
     EXPERIMENT_MONITOR_INFO, TRIAL_MONITOR_HEAD, TRIAL_MONITOR_CONTENT, TRIAL_MONITOR_TAIL, REST_TIME_OUT
34
from .common_utils import print_normal, print_error, print_warning, detect_process
35
from .command_utils import check_output_command, kill_command
Deshui Yu's avatar
Deshui Yu committed
36

37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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
53
54
55
56
57
58
59
    '''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):
60
            if experiment_dict[key].get('status') != 'STOPPED':
SparkSnail's avatar
SparkSnail committed
61
62
63
                nni_config = Config(experiment_dict[key]['fileName'])
                rest_pid = nni_config.get_config('restServerPid')
                if not detect_process(rest_pid):
64
65
66
67
68
69
70
71
72
73
74
                    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
75

76
77
78
def check_experiment_id(args):
    '''check if the id is valid
    '''
79
    update_experiment()
80
81
82
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
83
        print_normal('There is no experiment running...')
chicm-ms's avatar
chicm-ms committed
84
        return None
85
    if not args.id:
86
87
        running_experiment_list = []
        for key in experiment_dict.keys():
SparkSnail's avatar
SparkSnail committed
88
            if isinstance(experiment_dict[key], dict):
89
                if experiment_dict[key].get('status') != 'STOPPED':
SparkSnail's avatar
SparkSnail committed
90
91
92
93
                    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)
94
        if len(running_experiment_list) > 1:
95
            print_error('There are multiple experiments, please set the experiment id...')
96
97
98
            experiment_information = ""
            for key in running_experiment_list:
                experiment_information += (EXPERIMENT_DETAIL_FORMAT % (key, experiment_dict[key]['status'], \
SparkSnail's avatar
SparkSnail committed
99
                experiment_dict[key]['port'], experiment_dict[key].get('platform'), experiment_dict[key]['startTime'], experiment_dict[key]['endTime']))
100
101
102
103
            print(EXPERIMENT_INFORMATION_FORMAT % experiment_information)
            exit(1)
        elif not running_experiment_list:
            print_error('There is no experiment running!')
chicm-ms's avatar
chicm-ms committed
104
            return None
105
106
        else:
            return running_experiment_list[0]
107
108
109
110
111
    if experiment_dict.get(args.id):
        return args.id
    else:
        print_error('Id not correct!')
        return None
Deshui Yu's avatar
Deshui Yu committed
112

113
def parse_ids(args):
114
115
116
117
118
119
120
121
    '''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
    '''
122
    update_experiment()
123
124
125
126
127
128
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
        print_normal('Experiment is not running...')
        return None
    result_list = []
129
130
    running_experiment_list = []
    for key in experiment_dict.keys():
SparkSnail's avatar
SparkSnail committed
131
        if isinstance(experiment_dict[key], dict):
132
            if experiment_dict[key].get('status') != 'STOPPED':
SparkSnail's avatar
SparkSnail committed
133
134
135
136
                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)
137
    if not args.id:
138
        if len(running_experiment_list) > 1:
139
            print_error('There are multiple experiments, please set the experiment id...')
140
            experiment_information = ""
141
142
            for key in running_experiment_list:
                experiment_information += (EXPERIMENT_DETAIL_FORMAT % (key, experiment_dict[key]['status'], \
SparkSnail's avatar
SparkSnail committed
143
                experiment_dict[key]['port'], experiment_dict[key].get('platform'), experiment_dict[key]['startTime'], experiment_dict[key]['endTime']))
144
145
146
147
            print(EXPERIMENT_INFORMATION_FORMAT % experiment_information)
            exit(1)
        else:
            result_list = running_experiment_list
148
    elif args.id == 'all':
149
        result_list = running_experiment_list
150
    elif args.id.endswith('*'):
151
        for id in running_experiment_list:
152
153
            if id.startswith(args.id[:-1]):
                result_list.append(id)
154
    elif args.id in running_experiment_list:
155
156
        result_list.append(args.id)
    else:
157
        for id in running_experiment_list:
158
159
160
161
162
            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
chicm-ms's avatar
chicm-ms committed
163
164
165
166
    if not result_list and args.id:
        print_error('There are no experiments matched, please set correct experiment id...')
    elif not result_list:
        print_error('There is no experiment running...')
167
168
    return result_list

169
170
171
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
172
173
174
    if experiment_id is None:
        print_error('Please set the experiment id!')
        exit(1)
175
176
177
178
179
180
181
    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
182
183
184
    if experiment_id is None:
        print_error('Please set the experiment id!')
        exit(1)
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
    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:
        start_time = datetime.datetime.utcfromtimestamp(start_time_stamp // 1000).strftime("%Y/%m/%d %H:%M:%S")
        content['startTime'] = str(start_time)
    if end_time_stamp:
        end_time = datetime.datetime.utcfromtimestamp(end_time_stamp // 1000).strftime("%Y/%m/%d %H:%M:%S")
        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
211
212
def stop_experiment(args):
    '''Stop the experiment which is running'''
213
214
215
216
217
218
    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)
219
            nni_config = Config(experiment_dict[experiment_id]['fileName'])
220
221
222
            rest_port = nni_config.get_config('restServerPort')
            rest_pid = nni_config.get_config('restServerPid')
            if rest_pid:
223
                kill_command(rest_pid)
SparkSnail's avatar
SparkSnail committed
224
225
226
227
                tensorboard_pid_list = nni_config.get_config('tensorboardPidList')
                if tensorboard_pid_list:
                    for tensorboard_pid in tensorboard_pid_list:
                        try:
228
                            kill_command(tensorboard_pid)
SparkSnail's avatar
SparkSnail committed
229
230
231
                        except Exception as exception:
                            print_error(exception)
                    nni_config.set_config('tensorboardPidList', [])
SparkSnail's avatar
SparkSnail committed
232
            print_normal('Stop experiment success!')
233
            experiment_config.update_experiment(experiment_id, 'status', 'STOPPED')
234
235
            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
236
237
238

def trial_ls(args):
    '''List trial'''
239
    nni_config = Config(get_config_filename(args))
Deshui Yu's avatar
Deshui Yu committed
240
241
242
243
244
    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
245
246
    running, response = check_rest_server_quick(rest_port)
    if running:
247
        response = rest_get(trial_jobs_url(rest_port), REST_TIME_OUT)
248
        if response and check_response(response):
249
            content = json.loads(response.text)
250
            for index, value in enumerate(content):
251
252
                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
253
254
255
256
257
258
259
        else:
            print_error('List trial failed...')
    else:
        print_error('Restful server is not running...')

def trial_kill(args):
    '''List trial'''
260
    nni_config = Config(get_config_filename(args))
Deshui Yu's avatar
Deshui Yu committed
261
262
263
264
265
    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
266
267
    running, _ = check_rest_server_quick(rest_port)
    if running:
268
        response = rest_delete(trial_job_id_url(rest_port, args.id), REST_TIME_OUT)
269
        if response and check_response(response):
Deshui Yu's avatar
Deshui Yu committed
270
271
272
273
274
275
276
277
            print(response.text)
        else:
            print_error('Kill trial job failed...')
    else:
        print_error('Restful server is not running...')

def list_experiment(args):
    '''Get experiment information'''
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:
286
        response = rest_get(experiment_url(rest_port), REST_TIME_OUT)
287
        if response and check_response(response):
288
289
            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
290
291
292
293
294
        else:
            print_error('List experiment failed...')
    else:
        print_error('Restful server is not running...')

295
296
def experiment_status(args):
    '''Show the status of experiment'''
297
    nni_config = Config(get_config_filename(args))
298
299
300
301
302
303
304
    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
305
306
def log_internal(args, filetype):
    '''internal function to call get_log_content'''
307
    file_name = get_config_filename(args)
Deshui Yu's avatar
Deshui Yu committed
308
    if filetype == 'stdout':
309
        file_full_path = os.path.join(NNICTL_HOME_DIR, file_name, 'stdout')
Deshui Yu's avatar
Deshui Yu committed
310
    else:
311
        file_full_path = os.path.join(NNICTL_HOME_DIR, file_name, 'stderr')
312
313
    print(check_output_command(file_full_path, head=args.head, tail=args.tail))
    
Deshui Yu's avatar
Deshui Yu committed
314
315
316
317
318
319
320
321
def log_stdout(args):
    '''get stdout log'''
    log_internal(args, 'stdout')

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

322
323
324
def log_trial(args):
    ''''get trial log path'''
    trial_id_path_dict = {}
325
    nni_config = Config(get_config_filename(args))
326
327
328
329
330
331
332
    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:
333
        response = rest_get(trial_jobs_url(rest_port), REST_TIME_OUT)
334
335
336
337
338
339
        if response and check_response(response):
            content = json.loads(response.text)
            for trial in content:
                trial_id_path_dict[trial['id']] = trial['logPath']
    else:
        print_error('Restful server is not running...')
goooxu's avatar
goooxu committed
340
        exit(1)
341
342
343
344
    if args.id:
        if args.trial_id:
            if trial_id_path_dict.get(args.trial_id):
                print_normal('id:' + args.trial_id + ' path:' + trial_id_path_dict[args.trial_id])
345
346
347
            else:
                print_error('trial id is not valid!')
                exit(1)
348
        else:
349
            print_error('please specific the trial id!')
goooxu's avatar
goooxu committed
350
            exit(1)
351
    else:
352
        for key in trial_id_path_dict:
353
354
            print('id:' + key + ' path:' + trial_id_path_dict[key])

Deshui Yu's avatar
Deshui Yu committed
355
356
def get_config(args):
    '''get config info'''
357
    nni_config = Config(get_config_filename(args))
Deshui Yu's avatar
Deshui Yu committed
358
    print(nni_config.get_all_config())
359
360
361

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

365
366
def experiment_list(args):
    '''get the information of all experiments'''
367
368
369
370
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
        print('There is no experiment running...')
371
        exit(1)
372
    update_experiment()
373
374
375
376
    experiment_id_list = []
    if args.all and args.all == 'all':
        for key in experiment_dict.keys():
            experiment_id_list.append(key)
377
378
    else:
        for key in experiment_dict.keys():
379
            if experiment_dict[key]['status'] != 'STOPPED':
380
381
382
383
384
                experiment_id_list.append(key)
        if not experiment_id_list:
            print_warning('There is no experiment running...\nYou can use \'nnictl experiment list all\' to list all stopped experiments!')
    experiment_information = ""
    for key in experiment_id_list:
385
        
386
        experiment_information += (EXPERIMENT_DETAIL_FORMAT % (key, experiment_dict[key]['status'], experiment_dict[key]['port'],\
SparkSnail's avatar
SparkSnail committed
387
        experiment_dict[key].get('platform'), experiment_dict[key]['startTime'], experiment_dict[key]['endTime']))
388
    print(EXPERIMENT_INFORMATION_FORMAT % experiment_information)
SparkSnail's avatar
SparkSnail committed
389

SparkSnail's avatar
SparkSnail committed
390
391
392
393
def get_time_interval(time1, time2):
    '''get the interval of two times'''
    try:
        #convert time to timestamp
394
395
        time1 = time.mktime(time.strptime(time1, '%Y/%m/%d %H:%M:%S'))
        time2 = time.mktime(time.strptime(time2, '%Y/%m/%d %H:%M:%S'))
SparkSnail's avatar
SparkSnail committed
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
        seconds = (datetime.datetime.fromtimestamp(time2) - datetime.datetime.fromtimestamp(time1)).seconds
        #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'''
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    if not experiment_dict:
        print('There is no experiment running...')
        exit(1)
415
    update_experiment()
SparkSnail's avatar
SparkSnail committed
416
417
    experiment_id_list = []
    for key in experiment_dict.keys():
418
        if experiment_dict[key]['status'] != 'STOPPED':
SparkSnail's avatar
SparkSnail committed
419
420
421
422
423
424
            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'], \
425
             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
426
427
428
        print(TRIAL_MONITOR_HEAD)
        running, response = check_rest_server_quick(experiment_dict[key]['port'])
        if running:
429
            response = rest_get(trial_jobs_url(experiment_dict[key]['port']), REST_TIME_OUT)
SparkSnail's avatar
SparkSnail committed
430
431
            if response and check_response(response):
                content = json.loads(response.text)
432
                for index, value in enumerate(content):
SparkSnail's avatar
SparkSnail committed
433
434
435
436
437
438
439
440
441
442
443
444
                    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')
445
            update_experiment()
SparkSnail's avatar
SparkSnail committed
446
447
448
449
450
451
452
            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
453
454
455
456
457
458
459
460


def parse_trial_data(content):
    """output: List[Dict]"""
    trial_records = []
    for trial_data in content:
        for phase_i in range(len(trial_data['hyperParameters'])):
            hparam = json.loads(trial_data['hyperParameters'][phase_i])['parameters']
Yan Ni's avatar
Yan Ni committed
461
            hparam['id'] = trial_data['id']
Yan Ni's avatar
Yan Ni committed
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
            if 'finalMetricData' in trial_data.keys() and phase_i < len(trial_data['finalMetricData']):
                reward = json.loads(trial_data['finalMetricData'][phase_i]['data'])
                if isinstance(reward, (float, int)):
                    dict_tmp = {**hparam, **{'reward': reward}}
                elif isinstance(reward, dict):
                    dict_tmp = {**hparam, **reward}
                else:
                    raise ValueError("Invalid finalMetricsData format: {}/{}".format(type(reward), reward))
            else:
                dict_tmp = hparam
            trial_records.append(dict_tmp)
    return trial_records

def export_trials_data(args):
    """export experiment metadata to csv
    """
    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:
        response = rest_get(trial_jobs_url(rest_port), 20)
        if response is not None and check_response(response):
            content = json.loads(response.text)
            # dframe = pd.DataFrame.from_records([parse_trial_data(t_data) for t_data in content])
            # dframe.to_csv(args.csv_path, sep='\t')
            records = parse_trial_data(content)
492
493
494
495
496
497
498
499
500
501
502
503
504
            if args.type == 'json':
                json_records = []
                for trial in records:
                    value = trial.pop('reward', None)
                    trial_id =  trial.pop('id', None)
                    json_records.append({'parameter': trial, 'value': value, 'id': trial_id})
            with open(args.path, 'w') as file:
                if args.type == 'csv':
                    writer = csv.DictWriter(file, set.union(*[set(r.keys()) for r in records]))
                    writer.writeheader()
                    writer.writerows(records)
                else:
                    json.dump(json_records, file)
Yan Ni's avatar
Yan Ni committed
505
506
507
508
        else:
            print_error('Export failed...')
    else:
        print_error('Restful server is not Running')
SparkSnail's avatar
SparkSnail committed
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540

def hdfs_set(args):
    hdfsConfig = HDFSConfig()
    hdfsConfig.set_config(args.host, args.user_name)
    print_normal('HDFS account update success!')

def hdfs_clean(args):
    hdfsConfig = HDFSConfig()
    if not hdfsConfig.get_config():
        print_error('Please use \'nnictl hdfs set\' command to set hdfs account first!')
        exit(1)
    host = hdfsConfig.get_config().get('host')
    user_name = hdfsConfig.get_config().get('userName')
    hdfs_client = HdfsClient(hosts='{0}:80'.format(host), user_name=user_name, webhdfs_path='/webhdfs/api/v1', timeout=5)
    root_path = os.path.join('/', user_name, 'nni', 'experiments')
    while True:
        inputs = input('INFO: clean up all files in {0}, do you want to continue?[Y/N]:'.format(root_path))
        if inputs.lower() not in ['y', 'n', 'yes', 'no']:
            print_warning('please input Y or N!')
        elif inputs.lower() in ['n', 'no']:
            exit(0)
        else:
            break
    path_list = hdfs_client.listdir(root_path)
    for path in path_list:
        full_path = os.path.join(root_path, path)
        print_normal('deleting {0}'.format(full_path))
        if hdfs_client.delete(full_path, recursive=True):
            print_normal('delete success!')
        else:
            print_normal('delete failed!')
    print_normal('DONE')