launcher.py 14.6 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')
133
134
135
136
    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
137
138
    if experiment_config.get('multiPhase'):
        request_data['multiPhase'] = experiment_config.get('multiPhase')
Deshui Yu's avatar
Deshui Yu committed
139
140
141
142
143
144
145
    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(
146
            {'key':'codeDir', 'value':experiment_config['trial']['codeDir']})
Deshui Yu's avatar
Deshui Yu committed
147
        request_data['clusterMetaData'].append(
148
            {'key': 'command', 'value': experiment_config['trial']['command']})
149
    elif experiment_config['trainingServicePlatform'] == 'remote':
Deshui Yu's avatar
Deshui Yu committed
150
151
152
        request_data['clusterMetaData'].append(
            {'key': 'machine_list', 'value': experiment_config['machineList']})
        value_dict = dict()
153
154
155
        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
156
157
        request_data['clusterMetaData'].append(
            {'key': 'trial_config', 'value': value_dict})
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
    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
177
178

    response = rest_post(experiment_url(port), json.dumps(request_data), 20)
179
180
181
182
183
    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=(',', ':')))
184
        print_error('Setting experiment error, error message is {}'.format(response.text))
185
        return None
Deshui Yu's avatar
Deshui Yu committed
186
187
188
189

def launch_experiment(args, experiment_config, mode, webuiport, experiment_id=None):
    '''follow steps to start rest server and start experiment'''
    nni_config = Config()
190
191
192
193
194
195
    #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
196
    # start rest server
197
    rest_process = start_rest_server(REST_PORT, experiment_config['trainingServicePlatform'], mode, experiment_id)
Deshui Yu's avatar
Deshui Yu committed
198
199
200
    nni_config.set_config('restServerPid', rest_process.pid)
    # Deal with annotation
    if experiment_config.get('useAnnotation'):
liuzhe-lz's avatar
liuzhe-lz committed
201
        path = os.path.join(tempfile.gettempdir(), 'nni', 'annotation')
QuanluZhang's avatar
QuanluZhang committed
202
203
        if not os.path.isdir(path):
            os.makedirs(path)
liuzhe-lz's avatar
liuzhe-lz committed
204
205
206
207
        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)
208
        experiment_config['searchSpace'] = json.dumps(search_space)
Deshui Yu's avatar
Deshui Yu committed
209
        assert search_space, ERROR_INFO % 'Generated search space is empty'
210
211
212
    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
213
    else:
214
        experiment_config['searchSpace'] = json.dumps('')
Deshui Yu's avatar
Deshui Yu committed
215
216

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

    # set remote config
    if experiment_config['trainingServicePlatform'] == 'remote':
        print_normal('Setting remote config...')
232
233
        config_result, err_msg = set_remote_config(experiment_config, REST_PORT)
        if config_result:
Deshui Yu's avatar
Deshui Yu committed
234
235
            print_normal('Success!')
        else:
236
            print_error('Failed! Error is: {}'.format(err_msg))
Deshui Yu's avatar
Deshui Yu committed
237
            try:
238
239
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
Deshui Yu's avatar
Deshui Yu committed
240
241
242
243
244
245
246
247
            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):
248
            print_normal('Successfully set local config!')
Deshui Yu's avatar
Deshui Yu committed
249
250
251
        else:
            print_error('Failed!')
            try:
252
253
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
Deshui Yu's avatar
Deshui Yu committed
254
255
256
            except Exception:
                raise Exception(ERROR_INFO % 'Rest server stopped!')
            exit(0)
257
258
259
260
261
262
    
    #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:
263
            print_normal('Successfully set pai config!')
264
        else:
265
266
            if err_msg:
                print_error('Failed! Error is: {}'.format(err_msg))
267
268
269
270
            try:
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
            except Exception:
271
                raise Exception(ERROR_INFO % 'Restful server stopped!')
272
            exit(0)
273
274
275
276
277
278
279
280
281
282
283
    
    #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
284
285
286
287
288
289
290
291
292
293
294

    # 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:
295
296
            cmds = ['pkill', '-P', str(rest_process.pid)]
            call(cmds)
297
298
            cmds = ['pkill', '-P', str(webui_process.pid)]
            call(cmds)
Deshui Yu's avatar
Deshui Yu committed
299
        except Exception:
300
            raise Exception(ERROR_INFO % 'Restful server stopped!')
Deshui Yu's avatar
Deshui Yu committed
301
        exit(0)
302
    print_normal(EXPERIMENT_SUCCESS_INFO % (experiment_id, REST_PORT, '   '.join(nni_config.get_config('webuiUrl'))))
Deshui Yu's avatar
Deshui Yu committed
303
304
305
306
307
308
309
310
311
312
313

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()
314
315
316
    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
317
318
319
320

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