launcher.py 19.4 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
21
22
23
24
# 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.


import json
import os
import shutil
25
import string
26
from subprocess import Popen, PIPE, call, check_output
Deshui Yu's avatar
Deshui Yu committed
27
import tempfile
28
from nni_annotation import *
Deshui Yu's avatar
Deshui Yu committed
29
from .launcher_utils import validate_all_content
30
from .rest_utils import rest_put, rest_post, check_rest_server, check_rest_server_quick, check_response
SparkSnail's avatar
SparkSnail committed
31
from .url_utils import cluster_metadata_url, experiment_url, get_local_urls
32
33
from .config_utils import Config, Experiments
from .common_utils import get_yml_content, get_json_content, print_error, print_normal, print_warning, detect_process, detect_port
34
from .constants import *
35
import time
36
import random
QuanluZhang's avatar
QuanluZhang committed
37
import site
Gems Guo's avatar
Gems Guo committed
38
39
from pathlib import Path

40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def get_log_path(config_file_name):
    '''generate stdout and stderr log path'''
    stdout_full_path = os.path.join(NNICTL_HOME_DIR, config_file_name, 'stdout')
    stderr_full_path = os.path.join(NNICTL_HOME_DIR, config_file_name, 'stderr')
    return stdout_full_path, stderr_full_path

def print_log_content(config_file_name):
    '''print log information'''
    stdout_full_path, stderr_full_path = get_log_path(config_file_name)
    print_normal(' Stdout:')
    stdout_cmds = ['cat', stdout_full_path]
    stdout_content = check_output(stdout_cmds)
    print(stdout_content.decode('utf-8'))
    print('\n\n')
    print_normal(' Stderr:')
    stderr_cmds = ['cat', stderr_full_path]
    stderr_content = check_output(stderr_cmds)
    print(stderr_content.decode('utf-8'))


60
def start_rest_server(port, platform, mode, config_file_name, experiment_id=None):
Deshui Yu's avatar
Deshui Yu committed
61
    '''Run nni manager process'''
62
    nni_config = Config(config_file_name)
63
    if detect_port(port):
64
65
66
67
        print_error('Port %s is used by another process, please reset the port!\n' \
        'You could use \'nnictl create --help\' to get help information' % port)
        exit(1)
    
68
    if (platform == 'pai' or platform == 'kubeflow') and detect_port(int(port) + 1):
69
70
71
        print_error('PAI mode need an additional adjacent port %d, and the port %d is used by another process!\n' \
        'You could set another port to start experiment!\n' \
        'You could use \'nnictl create --help\' to get help information' % ((int(port) + 1), (int(port) + 1)))
72
        exit(1)
Deshui Yu's avatar
Deshui Yu committed
73
74

    print_normal('Starting restful server...')
QuanluZhang's avatar
QuanluZhang committed
75
76
77
78
79
80
81
82
83
84
85
    python_dir = str(Path(site.getusersitepackages()).parents[2])
    entry_file = os.path.join(python_dir, 'nni', 'main.js')
    entry_dir = os.path.join(python_dir, 'nni')
    local_entry_dir = entry_dir
    if not os.path.isfile(entry_file):
        python_dir = str(Path(site.getsitepackages()[0]).parents[2])
        entry_file = os.path.join(python_dir, 'nni', 'main.js')
        entry_dir = os.path.join(python_dir, 'nni')
        if not os.path.isfile(entry_file):
            raise Exception('Fail to find main.js under both %s and %s!' % (local_entry_dir, entry_dir))
    cmds = ['node', entry_file, '--port', str(port), '--mode', platform, '--start_mode', mode]
Deshui Yu's avatar
Deshui Yu committed
86
87
    if mode == 'resume':
        cmds += ['--experiment_id', experiment_id]
88
    stdout_full_path, stderr_full_path = get_log_path(config_file_name)
goooxu's avatar
goooxu committed
89
90
    stdout_file = open(stdout_full_path, 'a+')
    stderr_file = open(stderr_full_path, 'a+')
91
92
93
94
95
    time_now = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
    #add time information in the header of log files
    log_header = LOG_HEADER % str(time_now)
    stdout_file.write(log_header)
    stderr_file.write(log_header)
QuanluZhang's avatar
QuanluZhang committed
96
    process = Popen(cmds, cwd=entry_dir, stdout=stdout_file, stderr=stderr_file)
97
    return process, str(time_now)
Deshui Yu's avatar
Deshui Yu committed
98

99
def set_trial_config(experiment_config, port, config_file_name):
100
    '''set trial configuration'''
Deshui Yu's avatar
Deshui Yu committed
101
    request_data = dict()
102
    value_dict = dict()
103
104
105
    value_dict['command'] = experiment_config['trial']['command']
    value_dict['codeDir'] = experiment_config['trial']['codeDir']
    value_dict['gpuNum'] = experiment_config['trial']['gpuNum']
106
107
108
109
110
111
112
113
114
115
    if experiment_config['trial'].get('cpuNum'):
        value_dict['cpuNum'] = experiment_config['trial']['cpuNum']
    if experiment_config['trial'].get('memoryMB'):
        value_dict['memoryMB'] = experiment_config['trial']['memoryMB']
    if experiment_config['trial'].get('image'):
        value_dict['image'] = experiment_config['trial']['image']
    if experiment_config['trial'].get('dataDir'):
        value_dict['dataDir'] = experiment_config['trial']['dataDir']
    if experiment_config['trial'].get('outputDir'):
        value_dict['outputDir'] = experiment_config['trial']['outputDir']
116
    request_data['trial_config'] = value_dict
Deshui Yu's avatar
Deshui Yu committed
117
    response = rest_put(cluster_metadata_url(port), json.dumps(request_data), 20)
118
119
120
    if check_response(response):
        return True
    else:
121
        print('Error message is {}'.format(response.text))
122
        _, stderr_full_path = get_log_path(config_file_name)
SparkSnail's avatar
SparkSnail committed
123
124
125
        if response:
            with open(stderr_full_path, 'a+') as fout:
                fout.write(json.dumps(json.loads(response.text), indent=4, sort_keys=True, separators=(',', ':')))
126
        return False
127

128
def set_local_config(experiment_config, port, config_file_name):
129
    '''set local configuration'''
130
    return set_trial_config(experiment_config, port, config_file_name)
Deshui Yu's avatar
Deshui Yu committed
131

132
def set_remote_config(experiment_config, port, config_file_name):
Deshui Yu's avatar
Deshui Yu committed
133
134
135
136
137
    '''Call setClusterMetadata to pass trial'''
    #set machine_list
    request_data = dict()
    request_data['machine_list'] = experiment_config['machineList']
    response = rest_put(cluster_metadata_url(port), json.dumps(request_data), 20)
138
    err_message = ''
139
    if not response or not check_response(response):
140
141
        if response is not None:
            err_message = response.text
142
            _, stderr_full_path = get_log_path(config_file_name)
goooxu's avatar
goooxu committed
143
            with open(stderr_full_path, 'a+') as fout:
144
                fout.write(json.dumps(json.loads(err_message), indent=4, sort_keys=True, separators=(',', ':')))
145
        return False, err_message
Deshui Yu's avatar
Deshui Yu committed
146
147

    #set trial_config
148
    return set_trial_config(experiment_config, port, config_file_name), err_message
Deshui Yu's avatar
Deshui Yu committed
149

150
def set_pai_config(experiment_config, port, config_file_name):
151
152
153
154
    '''set pai configuration''' 
    pai_config_data = dict()
    pai_config_data['pai_config'] = experiment_config['paiConfig']
    response = rest_put(cluster_metadata_url(port), json.dumps(pai_config_data), 20)
155
    err_message = None
156
157
158
    if not response or not response.status_code == 200:
        if response is not None:
            err_message = response.text
159
            _, stderr_full_path = get_log_path(config_file_name)
160
            with open(stderr_full_path, 'a+') as fout:
chicm-ms's avatar
chicm-ms committed
161
                fout.write(json.dumps(json.loads(err_message), indent=4, sort_keys=True, separators=(',', ':')))
162
163
164
        return False, err_message

    #set trial_config
165
    return set_trial_config(experiment_config, port, config_file_name), err_message
166

167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def set_kubeflow_config(experiment_config, port, config_file_name):
    '''set kubeflow configuration''' 
    kubeflow_config_data = dict()
    kubeflow_config_data['kubeflow_config'] = experiment_config['kubeflowConfig']
    response = rest_put(cluster_metadata_url(port), json.dumps(kubeflow_config_data), 20)
    err_message = None
    if not response or not response.status_code == 200:
        if response is not None:
            err_message = response.text
            _, stderr_full_path = get_log_path(config_file_name)
            with open(stderr_full_path, 'a+') as fout:
                fout.write(json.dumps(json.loads(err_message), indent=4, sort_keys=True, separators=(',', ':')))
        return False, err_message

    #set trial_config
    return set_trial_config(experiment_config, port, config_file_name), err_message

184
def set_experiment(experiment_config, mode, port, config_file_name):
Deshui Yu's avatar
Deshui Yu committed
185
186
187
188
189
190
191
    '''Call startExperiment (rest POST /experiment) with yaml file content'''
    request_data = dict()
    request_data['authorName'] = experiment_config['authorName']
    request_data['experimentName'] = experiment_config['experimentName']
    request_data['trialConcurrency'] = experiment_config['trialConcurrency']
    request_data['maxExecDuration'] = experiment_config['maxExecDuration']
    request_data['maxTrialNum'] = experiment_config['maxTrialNum']
192
    request_data['searchSpace'] = experiment_config.get('searchSpace')
193
194
195
196
    request_data['trainingServicePlatform'] = experiment_config.get('trainingServicePlatform')

    if experiment_config.get('description'):
        request_data['description'] = experiment_config['description']
chicm-ms's avatar
chicm-ms committed
197
198
    if experiment_config.get('multiPhase'):
        request_data['multiPhase'] = experiment_config.get('multiPhase')
chicm-ms's avatar
chicm-ms committed
199
200
    if experiment_config.get('multiThread'):
        request_data['multiThread'] = experiment_config.get('multiThread')
Deshui Yu's avatar
Deshui Yu committed
201
202
203
204
205
206
207
    request_data['tuner'] = experiment_config['tuner']
    if 'assessor' in experiment_config:
        request_data['assessor'] = experiment_config['assessor']

    request_data['clusterMetaData'] = []
    if experiment_config['trainingServicePlatform'] == 'local':
        request_data['clusterMetaData'].append(
208
            {'key':'codeDir', 'value':experiment_config['trial']['codeDir']})
Deshui Yu's avatar
Deshui Yu committed
209
        request_data['clusterMetaData'].append(
210
            {'key': 'command', 'value': experiment_config['trial']['command']})
211
    elif experiment_config['trainingServicePlatform'] == 'remote':
Deshui Yu's avatar
Deshui Yu committed
212
213
214
        request_data['clusterMetaData'].append(
            {'key': 'machine_list', 'value': experiment_config['machineList']})
        value_dict = dict()
215
216
217
        value_dict['command'] = experiment_config['trial']['command']
        value_dict['codeDir'] = experiment_config['trial']['codeDir']
        value_dict['gpuNum'] = experiment_config['trial']['gpuNum']
Deshui Yu's avatar
Deshui Yu committed
218
219
        request_data['clusterMetaData'].append(
            {'key': 'trial_config', 'value': value_dict})
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
    elif experiment_config['trainingServicePlatform'] == 'pai':
        request_data['clusterMetaData'].append(
            {'key': 'pai_config', 'value': experiment_config['paiConfig']})
        value_dict = dict()
        value_dict['command'] = experiment_config['trial']['command']
        value_dict['codeDir'] = experiment_config['trial']['codeDir']
        value_dict['gpuNum'] = experiment_config['trial']['gpuNum']
        if experiment_config['trial'].get('cpuNum'):
            value_dict['cpuNum'] = experiment_config['trial']['cpuNum']
        if experiment_config['trial'].get('memoryMB'):
            value_dict['memoryMB'] = experiment_config['trial']['memoryMB']
        if experiment_config['trial'].get('image'):
            value_dict['image'] = experiment_config['trial']['image']
        if experiment_config['trial'].get('dataDir'):
            value_dict['dataDir'] = experiment_config['trial']['dataDir']
        if experiment_config['trial'].get('outputDir'):
            value_dict['outputDir'] = experiment_config['trial']['outputDir']
        request_data['clusterMetaData'].append(
            {'key': 'trial_config', 'value': value_dict})
Deshui Yu's avatar
Deshui Yu committed
239
240

    response = rest_post(experiment_url(port), json.dumps(request_data), 20)
241
242
243
    if check_response(response):
        return response
    else:
244
        _, stderr_full_path = get_log_path(config_file_name)
SparkSnail's avatar
SparkSnail committed
245
246
247
        if response:
            with open(stderr_full_path, 'a+') as fout:
                fout.write(json.dumps(json.loads(response.text), indent=4, sort_keys=True, separators=(',', ':')))
SparkSnail's avatar
SparkSnail committed
248
            print_error('Setting experiment error, error message is {}'.format(response.text))
249
        return None
Deshui Yu's avatar
Deshui Yu committed
250

251
def launch_experiment(args, experiment_config, mode, config_file_name, experiment_id=None):
Deshui Yu's avatar
Deshui Yu committed
252
    '''follow steps to start rest server and start experiment'''
253
    nni_config = Config(config_file_name)
Deshui Yu's avatar
Deshui Yu committed
254
    # start rest server
255
    rest_process, start_time = start_rest_server(args.port, experiment_config['trainingServicePlatform'], mode, config_file_name, experiment_id)
Deshui Yu's avatar
Deshui Yu committed
256
257
258
    nni_config.set_config('restServerPid', rest_process.pid)
    # Deal with annotation
    if experiment_config.get('useAnnotation'):
liuzhe-lz's avatar
liuzhe-lz committed
259
        path = os.path.join(tempfile.gettempdir(), 'nni', 'annotation')
QuanluZhang's avatar
QuanluZhang committed
260
261
        if not os.path.isdir(path):
            os.makedirs(path)
liuzhe-lz's avatar
liuzhe-lz committed
262
263
264
265
        path = tempfile.mkdtemp(dir=path)
        code_dir = expand_annotations(experiment_config['trial']['codeDir'], path)
        experiment_config['trial']['codeDir'] = code_dir
        search_space = generate_search_space(code_dir)
266
        experiment_config['searchSpace'] = json.dumps(search_space)
Deshui Yu's avatar
Deshui Yu committed
267
        assert search_space, ERROR_INFO % 'Generated search space is empty'
268
269
270
    elif experiment_config.get('searchSpacePath'):
            search_space = get_json_content(experiment_config.get('searchSpacePath'))
            experiment_config['searchSpace'] = json.dumps(search_space)
Deshui Yu's avatar
Deshui Yu committed
271
    else:
272
        experiment_config['searchSpace'] = json.dumps('')
Deshui Yu's avatar
Deshui Yu committed
273
274

    # check rest server
goooxu's avatar
goooxu committed
275
    running, _ = check_rest_server(args.port)
276
    if running:
277
        print_normal('Successfully started Restful server!')
Deshui Yu's avatar
Deshui Yu committed
278
279
    else:
        print_error('Restful server start failed!')
280
        print_log_content(config_file_name)
Deshui Yu's avatar
Deshui Yu committed
281
        try:
Gems Guo's avatar
Gems Guo committed
282
            cmds = ['kill', str(rest_process.pid)]
283
            call(cmds)
Deshui Yu's avatar
Deshui Yu committed
284
285
        except Exception:
            raise Exception(ERROR_INFO % 'Rest server stopped!')
goooxu's avatar
goooxu committed
286
        exit(1)
Deshui Yu's avatar
Deshui Yu committed
287
288
289
290

    # set remote config
    if experiment_config['trainingServicePlatform'] == 'remote':
        print_normal('Setting remote config...')
291
        config_result, err_msg = set_remote_config(experiment_config, args.port, config_file_name)
292
        if config_result:
293
            print_normal('Successfully set remote config!')
Deshui Yu's avatar
Deshui Yu committed
294
        else:
295
            print_error('Failed! Error is: {}'.format(err_msg))
Deshui Yu's avatar
Deshui Yu committed
296
            try:
Gems Guo's avatar
Gems Guo committed
297
                cmds = ['kill', str(rest_process.pid)]
298
                call(cmds)
Deshui Yu's avatar
Deshui Yu committed
299
300
            except Exception:
                raise Exception(ERROR_INFO % 'Rest server stopped!')
goooxu's avatar
goooxu committed
301
            exit(1)
Deshui Yu's avatar
Deshui Yu committed
302
303
304
305

    # set local config
    if experiment_config['trainingServicePlatform'] == 'local':
        print_normal('Setting local config...')
306
        if set_local_config(experiment_config, args.port, config_file_name):
307
            print_normal('Successfully set local config!')
Deshui Yu's avatar
Deshui Yu committed
308
        else:
309
            print_error('Set local config failed!')
Deshui Yu's avatar
Deshui Yu committed
310
            try:
Gems Guo's avatar
Gems Guo committed
311
                cmds = ['kill', str(rest_process.pid)]
312
                call(cmds)
Deshui Yu's avatar
Deshui Yu committed
313
314
            except Exception:
                raise Exception(ERROR_INFO % 'Rest server stopped!')
goooxu's avatar
goooxu committed
315
            exit(1)
316
317
318
319
    
    #set pai config
    if experiment_config['trainingServicePlatform'] == 'pai':
        print_normal('Setting pai config...')
320
        config_result, err_msg = set_pai_config(experiment_config, args.port, config_file_name)
321
        if config_result:
322
            print_normal('Successfully set pai config!')
323
        else:
324
325
            if err_msg:
                print_error('Failed! Error is: {}'.format(err_msg))
326
            try:
Gems Guo's avatar
Gems Guo committed
327
                cmds = ['kill', str(rest_process.pid)]
328
329
                call(cmds)
            except Exception:
330
                raise Exception(ERROR_INFO % 'Restful server stopped!')
goooxu's avatar
goooxu committed
331
            exit(1)
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
    
    #set kubeflow config
    if experiment_config['trainingServicePlatform'] == 'kubeflow':
        print_normal('Setting kubeflow config...')
        config_result, err_msg = set_kubeflow_config(experiment_config, args.port, config_file_name)
        if config_result:
            print_normal('Successfully set kubeflow config!')
        else:
            if err_msg:
                print_error('Failed! Error is: {}'.format(err_msg))
            try:
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
            except Exception:
                raise Exception(ERROR_INFO % 'Restful server stopped!')
            exit(1)
Deshui Yu's avatar
Deshui Yu committed
348
349
350

    # start a new experiment
    print_normal('Starting experiment...')
351
    response = set_experiment(experiment_config, mode, args.port, config_file_name)
Deshui Yu's avatar
Deshui Yu committed
352
353
354
355
356
    if response:
        if experiment_id is None:
            experiment_id = json.loads(response.text).get('experiment_id')
        nni_config.set_config('experimentId', experiment_id)
    else:
357
358
        print_error('Start experiment failed!')
        print_log_content(config_file_name)
Deshui Yu's avatar
Deshui Yu committed
359
        try:
Gems Guo's avatar
Gems Guo committed
360
            cmds = ['kill', str(rest_process.pid)]
361
            call(cmds)
Deshui Yu's avatar
Deshui Yu committed
362
        except Exception:
363
            raise Exception(ERROR_INFO % 'Restful server stopped!')
goooxu's avatar
goooxu committed
364
        exit(1)
SparkSnail's avatar
SparkSnail committed
365
366
    web_ui_url_list = get_local_urls(args.port)
    nni_config.set_config('webuiUrl', web_ui_url_list)
367
368
369
    
    #save experiment information
    experiment_config = Experiments()
370
    experiment_config.add_experiment(experiment_id, args.port, start_time, config_file_name)
371
372

    print_normal(EXPERIMENT_SUCCESS_INFO % (experiment_id, '   '.join(web_ui_url_list)))
Deshui Yu's avatar
Deshui Yu committed
373
374
375

def resume_experiment(args):
    '''resume an experiment'''
376
377
378
379
380
381
    experiment_config = Experiments()
    experiment_dict = experiment_config.get_all_experiments()
    experiment_id = None
    experiment_endTime = None
    #find the latest stopped experiment
    if not args.id:
SparkSnail's avatar
SparkSnail committed
382
383
384
        print_error('Please set experiment id! \nYou could use \'nnictl resume {id}\' to resume a stopped experiment!\n' \
        'You could use \'nnictl experiment list all\' to show all of stopped experiments!')
        exit(1)
385
386
387
388
389
390
391
392
393
394
    else:
        if experiment_dict.get(args.id) is None:
            print_error('Id %s not exist!' % args.id)
            exit(1)
        if experiment_dict[args.id]['status'] == 'running':
            print_error('Experiment %s is running!' % args.id)
            exit(1)
        experiment_id = args.id
    print_normal('Resuming experiment %s...' % experiment_id)
    nni_config = Config(experiment_dict[experiment_id]['fileName'])
Deshui Yu's avatar
Deshui Yu committed
395
396
    experiment_config = nni_config.get_config('experimentConfig')
    experiment_id = nni_config.get_config('experimentId')
SparkSnail's avatar
SparkSnail committed
397
398
399
400
401
    new_config_file_name = ''.join(random.sample(string.ascii_letters + string.digits, 8))
    new_nni_config = Config(new_config_file_name)
    new_nni_config.set_config('experimentConfig', experiment_config)
    launch_experiment(args, experiment_config, 'resume', new_config_file_name, experiment_id)
    new_nni_config.set_config('restServerPort', args.port)
Deshui Yu's avatar
Deshui Yu committed
402
403
404

def create_experiment(args):
    '''start a new experiment'''
405
406
    config_file_name = ''.join(random.sample(string.ascii_letters + string.digits, 8))
    nni_config = Config(config_file_name)
407
    config_path = os.path.abspath(args.config)
goooxu's avatar
goooxu committed
408
409
410
    if not os.path.exists(config_path):
        print_error('Please set correct config path!')
        exit(1)
411
412
    experiment_config = get_yml_content(config_path)
    validate_all_content(experiment_config, config_path)
Deshui Yu's avatar
Deshui Yu committed
413
414

    nni_config.set_config('experimentConfig', experiment_config)
415
    launch_experiment(args, experiment_config, 'new', config_file_name)
goooxu's avatar
goooxu committed
416
    nni_config.set_config('restServerPort', args.port)