launcher.py 14.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
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
Deshui Yu's avatar
Deshui Yu committed
31
32
from .url_utils import cluster_metadata_url, experiment_url
from .config_utils import Config
33
34
from .common_utils import get_yml_content, get_json_content, print_error, print_normal, print_warning, detect_process
from .constants import *
Deshui Yu's avatar
Deshui Yu committed
35
36
from .webui_utils import start_web_ui, check_web_ui

37
def start_rest_server(port, platform, mode, experiment_id=None):
Deshui Yu's avatar
Deshui Yu committed
38
    '''Run nni manager process'''
39
    print_normal('Checking environment...')
Deshui Yu's avatar
Deshui Yu committed
40
41
    nni_config = Config()
    rest_port = nni_config.get_config('restServerPort')
42
43
    running, _ = check_rest_server_quick(rest_port)
    if rest_port and running:
Deshui Yu's avatar
Deshui Yu committed
44
45
46
47
48
        print_error('There is an experiment running, please stop it first...')
        print_normal('You can use \'nnictl stop\' command to stop an experiment!')
        exit(0)

    print_normal('Starting restful server...')
49
    manager = os.environ.get('NNI_MANAGER', 'nnimanager')
Deshui Yu's avatar
Deshui Yu committed
50
51
52
53
54
55
56
57
58
59
    cmds = [manager, '--port', str(port), '--mode', platform, '--start_mode', mode]
    if mode == 'resume':
        cmds += ['--experiment_id', experiment_id]
    if not os.path.exists(LOG_DIR):
        os.makedirs(LOG_DIR)
    stdout_file = open(STDOUT_FULL_PATH, 'a+')
    stderr_file = open(STDERR_FULL_PATH, 'a+')
    process = Popen(cmds, stdout=stdout_file, stderr=stderr_file)
    return process

60
61
def set_trial_config(experiment_config, port):
    '''set trial configuration'''
Deshui Yu's avatar
Deshui Yu committed
62
    request_data = dict()
63
    value_dict = dict()
64
65
66
    value_dict['command'] = experiment_config['trial']['command']
    value_dict['codeDir'] = experiment_config['trial']['codeDir']
    value_dict['gpuNum'] = experiment_config['trial']['gpuNum']
67
68
69
70
71
72
73
74
75
76
    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']
77
    request_data['trial_config'] = value_dict
Deshui Yu's avatar
Deshui Yu committed
78
    response = rest_put(cluster_metadata_url(port), json.dumps(request_data), 20)
79
80
81
    if check_response(response):
        return True
    else:
82
        print('Error message is {}'.format(response.text))
83
84
85
        with open(STDERR_FULL_PATH, 'a+') as fout:
            fout.write(json.dumps(json.loads(response.text), indent=4, sort_keys=True, separators=(',', ':')))
        return False
86
87
88
89

def set_local_config(experiment_config, port):
    '''set local configuration'''
    return set_trial_config(experiment_config, port)
Deshui Yu's avatar
Deshui Yu committed
90
91
92
93
94
95
96

def set_remote_config(experiment_config, port):
    '''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)
97
    err_message = ''
98
    if not response or not check_response(response):
99
100
        if response is not None:
            err_message = response.text
101
102
            with open(STDERR_FULL_PATH, 'a+') as fout:
                fout.write(json.dumps(json.loads(err_message), indent=4, sort_keys=True, separators=(',', ':')))
103
        return False, err_message
Deshui Yu's avatar
Deshui Yu committed
104
105

    #set trial_config
106
    return set_trial_config(experiment_config, port), err_message
Deshui Yu's avatar
Deshui Yu committed
107

108
109
110
111
112
def set_pai_config(experiment_config, port):
    '''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)
113
    err_message = None
114
115
116
    if not response or not response.status_code == 200:
        if response is not None:
            err_message = response.text
chicm-ms's avatar
chicm-ms committed
117
118
            with open(STDERR_FULL_PATH, 'a+') as fout:
                fout.write(json.dumps(json.loads(err_message), indent=4, sort_keys=True, separators=(',', ':')))
119
120
121
122
123
        return False, err_message

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

Deshui Yu's avatar
Deshui Yu committed
124
125
126
127
128
129
130
131
def set_experiment(experiment_config, mode, port):
    '''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']
132
    request_data['searchSpace'] = experiment_config.get('searchSpace')
chicm-ms's avatar
chicm-ms committed
133
134
    if experiment_config.get('multiPhase'):
        request_data['multiPhase'] = experiment_config.get('multiPhase')
Deshui Yu's avatar
Deshui Yu committed
135
136
137
138
139
140
141
    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(
142
            {'key':'codeDir', 'value':experiment_config['trial']['codeDir']})
Deshui Yu's avatar
Deshui Yu committed
143
        request_data['clusterMetaData'].append(
144
            {'key': 'command', 'value': experiment_config['trial']['command']})
145
    elif experiment_config['trainingServicePlatform'] == 'remote':
Deshui Yu's avatar
Deshui Yu committed
146
147
148
        request_data['clusterMetaData'].append(
            {'key': 'machine_list', 'value': experiment_config['machineList']})
        value_dict = dict()
149
150
151
        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
152
153
        request_data['clusterMetaData'].append(
            {'key': 'trial_config', 'value': value_dict})
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
    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
173
174

    response = rest_post(experiment_url(port), json.dumps(request_data), 20)
175
176
177
178
179
    if check_response(response):
        return response
    else:
        with open(STDERR_FULL_PATH, 'a+') as fout:
            fout.write(json.dumps(json.loads(response.text), indent=4, sort_keys=True, separators=(',', ':')))
180
        print_error('Setting experiment error, error message is {}'.format(response.text))
181
        return None
Deshui Yu's avatar
Deshui Yu committed
182
183
184
185

def launch_experiment(args, experiment_config, mode, webuiport, experiment_id=None):
    '''follow steps to start rest server and start experiment'''
    nni_config = Config()
186
187
188
189
190
191
    #Check if there is an experiment running
    origin_rest_pid = nni_config.get_config('restServerPid')
    if origin_rest_pid and detect_process(origin_rest_pid):
        print_error('There is an experiment running, please stop it first...')
        print_normal('You can use \'nnictl stop\' command to stop an experiment!')
        exit(0)
Deshui Yu's avatar
Deshui Yu committed
192
    # start rest server
193
    rest_process = start_rest_server(REST_PORT, experiment_config['trainingServicePlatform'], mode, experiment_id)
Deshui Yu's avatar
Deshui Yu committed
194
195
196
    nni_config.set_config('restServerPid', rest_process.pid)
    # Deal with annotation
    if experiment_config.get('useAnnotation'):
liuzhe-lz's avatar
liuzhe-lz committed
197
        path = os.path.join(tempfile.gettempdir(), 'nni', 'annotation')
QuanluZhang's avatar
QuanluZhang committed
198
199
        if not os.path.isdir(path):
            os.makedirs(path)
liuzhe-lz's avatar
liuzhe-lz committed
200
201
202
203
        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)
204
        experiment_config['searchSpace'] = json.dumps(search_space)
Deshui Yu's avatar
Deshui Yu committed
205
        assert search_space, ERROR_INFO % 'Generated search space is empty'
206
207
208
    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
209
    else:
210
        experiment_config['searchSpace'] = json.dumps('')
Deshui Yu's avatar
Deshui Yu committed
211
212

    # check rest server
213
214
    running, _ = check_rest_server(REST_PORT)
    if running:
215
        print_normal('Successfully started Restful server!')
Deshui Yu's avatar
Deshui Yu committed
216
217
218
    else:
        print_error('Restful server start failed!')
        try:
219
220
            cmds = ['pkill', '-P', str(rest_process.pid)]
            call(cmds)
Deshui Yu's avatar
Deshui Yu committed
221
222
223
224
225
226
227
        except Exception:
            raise Exception(ERROR_INFO % 'Rest server stopped!')
        exit(0)

    # set remote config
    if experiment_config['trainingServicePlatform'] == 'remote':
        print_normal('Setting remote config...')
228
229
        config_result, err_msg = set_remote_config(experiment_config, REST_PORT)
        if config_result:
Deshui Yu's avatar
Deshui Yu committed
230
231
            print_normal('Success!')
        else:
232
            print_error('Failed! Error is: {}'.format(err_msg))
Deshui Yu's avatar
Deshui Yu committed
233
            try:
234
235
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
Deshui Yu's avatar
Deshui Yu committed
236
237
238
239
240
241
242
243
            except Exception:
                raise Exception(ERROR_INFO % 'Rest server stopped!')
            exit(0)

    # set local config
    if experiment_config['trainingServicePlatform'] == 'local':
        print_normal('Setting local config...')
        if set_local_config(experiment_config, REST_PORT):
244
            print_normal('Successfully set local config!')
Deshui Yu's avatar
Deshui Yu committed
245
246
247
        else:
            print_error('Failed!')
            try:
248
249
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
Deshui Yu's avatar
Deshui Yu committed
250
251
252
            except Exception:
                raise Exception(ERROR_INFO % 'Rest server stopped!')
            exit(0)
253
254
255
256
257
258
    
    #set pai config
    if experiment_config['trainingServicePlatform'] == 'pai':
        print_normal('Setting pai config...')
        config_result, err_msg = set_pai_config(experiment_config, REST_PORT)
        if config_result:
259
            print_normal('Successfully set pai config!')
260
        else:
261
262
            if err_msg:
                print_error('Failed! Error is: {}'.format(err_msg))
263
264
265
266
            try:
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
            except Exception:
267
                raise Exception(ERROR_INFO % 'Restful server stopped!')
268
            exit(0)
269
270
271
272
273
274
275
276
277
278
279
    
    #start webui
    if check_web_ui():
        print_warning('{0} {1}'.format(' '.join(nni_config.get_config('webuiUrl')),'is being used, please stop it first!'))
        print_normal('You can use \'nnictl webui stop\' to stop old Web UI process...')
    else:
        print_normal('Starting Web UI...')
        webui_process = start_web_ui(webuiport)
        if webui_process:
            nni_config.set_config('webuiPid', webui_process.pid)
            print_normal('Successfully started Web UI!')
Deshui Yu's avatar
Deshui Yu committed
280
281
282
283
284
285
286
287
288
289
290

    # start a new experiment
    print_normal('Starting experiment...')
    response = set_experiment(experiment_config, mode, REST_PORT)
    if response:
        if experiment_id is None:
            experiment_id = json.loads(response.text).get('experiment_id')
        nni_config.set_config('experimentId', experiment_id)
    else:
        print_error('Failed!')
        try:
291
292
            cmds = ['pkill', '-P', str(rest_process.pid)]
            call(cmds)
293
294
            cmds = ['pkill', '-P', str(webui_process.pid)]
            call(cmds)
Deshui Yu's avatar
Deshui Yu committed
295
        except Exception:
296
            raise Exception(ERROR_INFO % 'Restful server stopped!')
Deshui Yu's avatar
Deshui Yu committed
297
        exit(0)
298
    print_normal(EXPERIMENT_SUCCESS_INFO % (experiment_id, REST_PORT, '   '.join(nni_config.get_config('webuiUrl'))))
Deshui Yu's avatar
Deshui Yu committed
299
300
301
302
303
304
305
306
307
308
309

def resume_experiment(args):
    '''resume an experiment'''
    nni_config = Config()
    experiment_config = nni_config.get_config('experimentConfig')
    experiment_id = nni_config.get_config('experimentId')
    launch_experiment(args, experiment_config, 'resume', args.webuiport, experiment_id)

def create_experiment(args):
    '''start a new experiment'''
    nni_config = Config()
310
311
312
    config_path = os.path.abspath(args.config)
    experiment_config = get_yml_content(config_path)
    validate_all_content(experiment_config, config_path)
Deshui Yu's avatar
Deshui Yu committed
313
314
315
316

    nni_config.set_config('experimentConfig', experiment_config)
    launch_experiment(args, experiment_config, 'new', args.webuiport)
    nni_config.set_config('restServerPort', REST_PORT)