launcher.py 16.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
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
37
import random
import string
Deshui Yu's avatar
Deshui Yu committed
38

39
def start_rest_server(port, platform, mode, config_file_name, experiment_id=None):
Deshui Yu's avatar
Deshui Yu committed
40
    '''Run nni manager process'''
41
    nni_config = Config(config_file_name)
42
43
44
    if detect_port(port):
        print_error('Port %s is used by another process, please reset the port!' % port)
        exit(1)
Deshui Yu's avatar
Deshui Yu committed
45
46

    print_normal('Starting restful server...')
47
    manager = os.environ.get('NNI_MANAGER', 'nnimanager')
Deshui Yu's avatar
Deshui Yu committed
48
49
50
    cmds = [manager, '--port', str(port), '--mode', platform, '--start_mode', mode]
    if mode == 'resume':
        cmds += ['--experiment_id', experiment_id]
51
52
    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')
goooxu's avatar
goooxu committed
53
54
    stdout_file = open(stdout_full_path, 'a+')
    stderr_file = open(stderr_full_path, 'a+')
55
56
57
58
59
    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)
Deshui Yu's avatar
Deshui Yu committed
60
    process = Popen(cmds, stdout=stdout_file, stderr=stderr_file)
61
    return process, str(time_now)
Deshui Yu's avatar
Deshui Yu committed
62

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

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

95
def set_remote_config(experiment_config, port, config_file_name):
Deshui Yu's avatar
Deshui Yu committed
96
97
98
99
100
    '''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)
101
    err_message = ''
102
    if not response or not check_response(response):
103
104
        if response is not None:
            err_message = response.text
105
            stderr_full_path = os.path.join(NNICTL_HOME_DIR, config_file_name, 'stderr')
goooxu's avatar
goooxu committed
106
            with open(stderr_full_path, 'a+') as fout:
107
                fout.write(json.dumps(json.loads(err_message), indent=4, sort_keys=True, separators=(',', ':')))
108
        return False, err_message
Deshui Yu's avatar
Deshui Yu committed
109
110

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

113
def set_pai_config(experiment_config, port, config_file_name):
114
115
116
117
    '''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)
118
    err_message = None
119
120
121
    if not response or not response.status_code == 200:
        if response is not None:
            err_message = response.text
122
            stderr_full_path = os.path.join(NNICTL_HOME_DIR, config_file_name, 'stderr')
123
            with open(stderr_full_path, 'a+') as fout:
chicm-ms's avatar
chicm-ms committed
124
                fout.write(json.dumps(json.loads(err_message), indent=4, sort_keys=True, separators=(',', ':')))
125
126
127
        return False, err_message

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

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

    response = rest_post(experiment_url(port), json.dumps(request_data), 20)
185
186
187
    if check_response(response):
        return response
    else:
188
        stderr_full_path = os.path.join(NNICTL_HOME_DIR, config_file_name, 'stderr')
goooxu's avatar
goooxu committed
189
        with open(stderr_full_path, 'a+') as fout:
190
            fout.write(json.dumps(json.loads(response.text), indent=4, sort_keys=True, separators=(',', ':')))
191
        print_error('Setting experiment error, error message is {}'.format(response.text))
192
        return None
Deshui Yu's avatar
Deshui Yu committed
193

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

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

    # set remote config
    if experiment_config['trainingServicePlatform'] == 'remote':
        print_normal('Setting remote config...')
233
        config_result, err_msg = set_remote_config(experiment_config, args.port, config_file_name)
234
        if config_result:
235
            print_normal('Successfully set remote config!')
Deshui Yu's avatar
Deshui Yu committed
236
        else:
237
            print_error('Failed! Error is: {}'.format(err_msg))
Deshui Yu's avatar
Deshui Yu committed
238
            try:
239
240
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
Deshui Yu's avatar
Deshui Yu committed
241
242
            except Exception:
                raise Exception(ERROR_INFO % 'Rest server stopped!')
goooxu's avatar
goooxu committed
243
            exit(1)
Deshui Yu's avatar
Deshui Yu committed
244
245
246
247

    # set local config
    if experiment_config['trainingServicePlatform'] == 'local':
        print_normal('Setting local config...')
248
        if set_local_config(experiment_config, args.port, config_file_name):
249
            print_normal('Successfully set local config!')
Deshui Yu's avatar
Deshui Yu committed
250
251
252
        else:
            print_error('Failed!')
            try:
253
254
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
Deshui Yu's avatar
Deshui Yu committed
255
256
            except Exception:
                raise Exception(ERROR_INFO % 'Rest server stopped!')
goooxu's avatar
goooxu committed
257
            exit(1)
258
259
260
261
    
    #set pai config
    if experiment_config['trainingServicePlatform'] == 'pai':
        print_normal('Setting pai config...')
262
        config_result, err_msg = set_pai_config(experiment_config, args.port, config_file_name)
263
        if config_result:
264
            print_normal('Successfully set pai config!')
265
        else:
266
267
            if err_msg:
                print_error('Failed! Error is: {}'.format(err_msg))
268
269
270
271
            try:
                cmds = ['pkill', '-P', str(rest_process.pid)]
                call(cmds)
            except Exception:
272
                raise Exception(ERROR_INFO % 'Restful server stopped!')
goooxu's avatar
goooxu committed
273
            exit(1)
Deshui Yu's avatar
Deshui Yu committed
274
275
276

    # start a new experiment
    print_normal('Starting experiment...')
277
    response = set_experiment(experiment_config, mode, args.port, config_file_name)
Deshui Yu's avatar
Deshui Yu committed
278
279
280
281
282
283
284
    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:
285
286
            cmds = ['pkill', '-P', str(rest_process.pid)]
            call(cmds)
Deshui Yu's avatar
Deshui Yu committed
287
        except Exception:
288
            raise Exception(ERROR_INFO % 'Restful server stopped!')
goooxu's avatar
goooxu committed
289
        exit(1)
SparkSnail's avatar
SparkSnail committed
290
291
    web_ui_url_list = get_local_urls(args.port)
    nni_config.set_config('webuiUrl', web_ui_url_list)
292
293
294
    
    #save experiment information
    experiment_config = Experiments()
295
    experiment_config.add_experiment(experiment_id, args.port, start_time, config_file_name)
296
297

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

299
300
301
302
303
304
305
306
307
def cmp_time(time1, time2):
    '''compare the time'''
    try:
        time1 = time.strptime(time1,'%Y-%m-%d %H:%M:%S')
        time2 = time.strptime(time2,'%Y-%m-%d %H:%M:%S')
        return int(time1) - int(time2)
    except:
        return 0

Deshui Yu's avatar
Deshui Yu committed
308
309
def resume_experiment(args):
    '''resume an experiment'''
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
    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:
        for key in experiment_dict.keys():
            if experiment_dict[key]['status'] == 'stopped':
                if experiment_id is None:
                    experiment_id = key
                    experiment_endTime = experiment_dict[key]['endTime']
                else:
                    if cmp_time(experiment_dict[key]['endTime'], experiment_endTime) > 0:
                        experiment_id = key
                        experiment_endTime = experiment_dict[key]['endTime']
        if experiment_id is None:
            print_error('There is no experiment stopped!')
            exit(1)
    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
338
339
    experiment_config = nni_config.get_config('experimentConfig')
    experiment_id = nni_config.get_config('experimentId')
340
    launch_experiment(args, experiment_config, 'resume', experiment_dict[experiment_id]['fileName'], experiment_id)
Deshui Yu's avatar
Deshui Yu committed
341
342
343

def create_experiment(args):
    '''start a new experiment'''
344
345
    config_file_name = ''.join(random.sample(string.ascii_letters + string.digits, 8))
    nni_config = Config(config_file_name)
346
    config_path = os.path.abspath(args.config)
goooxu's avatar
goooxu committed
347
348
349
    if not os.path.exists(config_path):
        print_error('Please set correct config path!')
        exit(1)
350
351
    experiment_config = get_yml_content(config_path)
    validate_all_content(experiment_config, config_path)
Deshui Yu's avatar
Deshui Yu committed
352
353

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