config_schema.py 22.6 KB
Newer Older
liuzhe-lz's avatar
liuzhe-lz committed
1
2
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
3
4

import os
chicm-ms's avatar
chicm-ms committed
5
from schema import Schema, And, Optional, Regex, Or
6
7
8
from .constants import SCHEMA_TYPE_ERROR, SCHEMA_RANGE_ERROR, SCHEMA_PATH_ERROR


chicm-ms's avatar
chicm-ms committed
9
def setType(key, valueType):
10
    '''check key type'''
chicm-ms's avatar
chicm-ms committed
11
    return And(valueType, error=SCHEMA_TYPE_ERROR % (key, valueType.__name__))
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

def setChoice(key, *args):
    '''check choice'''
    return And(lambda n: n in args, error=SCHEMA_RANGE_ERROR % (key, str(args)))

def setNumberRange(key, keyType, start, end):
    '''check number range'''
    return And(
        And(keyType, error=SCHEMA_TYPE_ERROR % (key, keyType.__name__)),
        And(lambda n: start <= n <= end, error=SCHEMA_RANGE_ERROR % (key, '(%s,%s)' % (start, end))),
    )

def setPathCheck(key):
    '''check if path exist'''
    return And(os.path.exists, error=SCHEMA_PATH_ERROR % key)
27

28
common_schema = {
29
30
31
32
    'authorName': setType('authorName', str),
    'experimentName': setType('experimentName', str),
    Optional('description'): setType('description', str),
    'trialConcurrency': setNumberRange('trialConcurrency', int, 1, 99999),
chicm-ms's avatar
chicm-ms committed
33
    Optional('maxExecDuration'): And(Regex(r'^[1-9][0-9]*[s|m|h|d]$', error='ERROR: maxExecDuration format is [digit]{s,m,h,d}')),
34
    Optional('maxTrialNum'): setNumberRange('maxTrialNum', int, 1, 99999),
George Cheng's avatar
George Cheng committed
35
36
    'trainingServicePlatform': setChoice(
        'trainingServicePlatform', 'remote', 'local', 'pai', 'kubeflow', 'frameworkcontroller', 'paiYarn', 'dlts'),
37
38
39
40
41
42
    Optional('searchSpacePath'): And(os.path.exists, error=SCHEMA_PATH_ERROR % 'searchSpacePath'),
    Optional('multiPhase'): setType('multiPhase', bool),
    Optional('multiThread'): setType('multiThread', bool),
    Optional('nniManagerIp'): setType('nniManagerIp', str),
    Optional('logDir'): And(os.path.isdir, error=SCHEMA_PATH_ERROR % 'logDir'),
    Optional('debug'): setType('debug', bool),
43
    Optional('versionCheck'): setType('versionCheck', bool),
44
45
46
47
48
49
50
    Optional('logLevel'): setChoice('logLevel', 'trace', 'debug', 'info', 'warning', 'error', 'fatal'),
    Optional('logCollection'): setChoice('logCollection', 'http', 'none'),
    'useAnnotation': setType('useAnnotation', bool),
    Optional('tuner'): dict,
    Optional('advisor'): dict,
    Optional('assessor'): dict,
    Optional('localConfig'): {
51
52
53
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
        Optional('maxTrialNumPerGpu'): setType('maxTrialNumPerGpu', int),
        Optional('useActiveGpu'): setType('useActiveGpu', bool)
54
55
56
    }
}
tuner_schema_dict = {
57
58
    'Anneal': {
        'builtinTunerName': 'Anneal',
59
60
61
        Optional('classArgs'): {
            'optimize_mode': setChoice('optimize_mode', 'maximize', 'minimize'),
        },
62
63
64
65
66
67
68
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
    },
    'SMAC': {
        'builtinTunerName': 'SMAC',
        Optional('classArgs'): {
            'optimize_mode': setChoice('optimize_mode', 'maximize', 'minimize'),
QuanluZhang's avatar
QuanluZhang committed
69
            Optional('config_dedup'): setType('config_dedup', bool)
70
        },
71
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
72
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
QuanluZhang's avatar
QuanluZhang committed
73
    },
74
75
76
77
    ('Evolution'): {
        'builtinTunerName': setChoice('builtinTunerName', 'Evolution'),
        Optional('classArgs'): {
            'optimize_mode': setChoice('optimize_mode', 'maximize', 'minimize'),
Lee's avatar
Lee committed
78
            Optional('population_size'): setNumberRange('population_size', int, 0, 99999),
79
        },
80
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
81
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
82
    },
83
84
    ('BatchTuner', 'GridSearch', 'Random'): {
        'builtinTunerName': setChoice('builtinTunerName', 'BatchTuner', 'GridSearch', 'Random'),
85
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
86
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
Shufan Huang's avatar
Shufan Huang committed
87
    },
xuehui's avatar
xuehui committed
88
89
    'TPE': {
        'builtinTunerName': 'TPE',
90
        Optional('classArgs'): {
xuehui's avatar
xuehui committed
91
92
93
94
95
            Optional('optimize_mode'): setChoice('optimize_mode', 'maximize', 'minimize'),
            Optional('parallel_optimize'): setType('parallel_optimize', bool),
            Optional('constant_liar_type'): setChoice('constant_liar_type', 'min', 'max', 'mean')
        },
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
96
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
xuehui's avatar
xuehui committed
97
    },
98
99
    'NetworkMorphism': {
        'builtinTunerName': 'NetworkMorphism',
100
        Optional('classArgs'): {
101
            Optional('optimize_mode'): setChoice('optimize_mode', 'maximize', 'minimize'),
chicm-ms's avatar
chicm-ms committed
102
            Optional('task'): setChoice('task', 'cv', 'nlp', 'common'),
103
104
105
106
            Optional('input_width'): setType('input_width', int),
            Optional('input_channel'): setType('input_channel', int),
            Optional('n_output_node'): setType('n_output_node', int),
            },
107
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
108
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
109
    },
110
111
    'MetisTuner': {
        'builtinTunerName': 'MetisTuner',
112
        Optional('classArgs'): {
113
114
115
116
117
118
            Optional('optimize_mode'): setChoice('optimize_mode', 'maximize', 'minimize'),
            Optional('no_resampling'): setType('no_resampling', bool),
            Optional('no_candidates'): setType('no_candidates', bool),
            Optional('selection_num_starting_points'):  setType('selection_num_starting_points', int),
            Optional('cold_start_num'): setType('cold_start_num', int),
            },
119
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
120
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
121
    },
Guoxin's avatar
Guoxin committed
122
123
    'GPTuner': {
        'builtinTunerName': 'GPTuner',
124
        Optional('classArgs'): {
Guoxin's avatar
Guoxin committed
125
126
127
128
129
130
131
132
133
134
            Optional('optimize_mode'): setChoice('optimize_mode', 'maximize', 'minimize'),
            Optional('utility'): setChoice('utility', 'ei', 'ucb', 'poi'),
            Optional('kappa'): setType('kappa', float),
            Optional('xi'): setType('xi', float),
            Optional('nu'): setType('nu', float),
            Optional('alpha'): setType('alpha', float),
            Optional('cold_start_num'): setType('cold_start_num', int),
            Optional('selection_num_warm_up'):  setType('selection_num_warm_up', int),
            Optional('selection_num_starting_points'):  setType('selection_num_starting_points', int),
            },
chicm-ms's avatar
chicm-ms committed
135
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
136
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
Guoxin's avatar
Guoxin committed
137
    },
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
    'PPOTuner': {
        'builtinTunerName': 'PPOTuner',
        'classArgs': {
            'optimize_mode': setChoice('optimize_mode', 'maximize', 'minimize'),
            Optional('trials_per_update'): setNumberRange('trials_per_update', int, 0, 99999),
            Optional('epochs_per_update'): setNumberRange('epochs_per_update', int, 0, 99999),
            Optional('minibatch_size'): setNumberRange('minibatch_size', int, 0, 99999),
            Optional('ent_coef'): setType('ent_coef', float),
            Optional('lr'): setType('lr', float),
            Optional('vf_coef'): setType('vf_coef', float),
            Optional('max_grad_norm'): setType('max_grad_norm', float),
            Optional('gamma'): setType('gamma', float),
            Optional('lam'): setType('lam', float),
            Optional('cliprange'): setType('cliprange', float),
        },
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
154
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
155
    },
RayMeng8's avatar
RayMeng8 committed
156
157
158
159
160
161
162
163
164
165
166
167
    'PBTTuner': {
        'builtinTunerName': 'PBTTuner',
        'classArgs': {
            'optimize_mode': setChoice('optimize_mode', 'maximize', 'minimize'),
            Optional('all_checkpoint_dir'): setType('all_checkpoint_dir', str),
            Optional('population_size'): setNumberRange('population_size', int, 0, 99999),
            Optional('factors'): setType('factors', tuple),
            Optional('fraction'): setType('fraction', float),
        },
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
    },
168
169
170
171
172
    'customized': {
        'codeDir': setPathCheck('codeDir'),
        'classFileName': setType('classFileName', str),
        'className': setType('className', str),
        Optional('classArgs'): dict,
173
        Optional('includeIntermediateResults'): setType('includeIntermediateResults', bool),
174
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
175
176
177
178
179
180
181
182
183
184
    }
}

advisor_schema_dict = {
    'Hyperband':{
        'builtinAdvisorName': Or('Hyperband'),
        'classArgs': {
            'optimize_mode': setChoice('optimize_mode', 'maximize', 'minimize'),
            Optional('R'): setType('R', int),
            Optional('eta'): setType('eta', int)
xuehui's avatar
xuehui committed
185
        },
186
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
187
    },
188
189
190
191
192
193
    'BOHB':{
        'builtinAdvisorName': Or('BOHB'),
        'classArgs': {
            'optimize_mode': setChoice('optimize_mode', 'maximize', 'minimize'),
            Optional('min_budget'): setNumberRange('min_budget', int, 0, 9999),
            Optional('max_budget'): setNumberRange('max_budget', int, 0, 9999),
194
            Optional('eta'):setNumberRange('eta', int, 0, 9999),
195
196
197
198
199
200
201
            Optional('min_points_in_model'): setNumberRange('min_points_in_model', int, 0, 9999),
            Optional('top_n_percent'): setNumberRange('top_n_percent', int, 1, 99),
            Optional('num_samples'): setNumberRange('num_samples', int, 1, 9999),
            Optional('random_fraction'): setNumberRange('random_fraction', float, 0, 9999),
            Optional('bandwidth_factor'): setNumberRange('bandwidth_factor', float, 0, 9999),
            Optional('min_bandwidth'): setNumberRange('min_bandwidth', float, 0, 9999),
        },
202
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
203
    },
204
205
206
207
208
    'customized':{
        'codeDir': setPathCheck('codeDir'),
        'classFileName': setType('classFileName', str),
        'className': setType('className', str),
        Optional('classArgs'): dict,
209
        Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
210
    }
211
}
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235

assessor_schema_dict = {
    'Medianstop': {
        'builtinAssessorName': 'Medianstop',
        Optional('classArgs'): {
            Optional('optimize_mode'): setChoice('optimize_mode', 'maximize', 'minimize'),
            Optional('start_step'): setNumberRange('start_step', int, 0, 9999),
        },
    },
    'Curvefitting': {
        'builtinAssessorName': 'Curvefitting',
        Optional('classArgs'): {
            'epoch_num': setNumberRange('epoch_num', int, 0, 9999),
            Optional('start_step'): setNumberRange('start_step', int, 0, 9999),
            Optional('threshold'): setNumberRange('threshold', float, 0, 9999),
            Optional('gap'): setNumberRange('gap', int, 1, 9999),
        },
    },
    'customized': {
        'codeDir': setPathCheck('codeDir'),
        'classFileName': setType('classFileName', str),
        'className': setType('className', str),
        Optional('classArgs'): dict,
    }
236
237
238
}

common_trial_schema = {
chicm-ms's avatar
chicm-ms committed
239
240
241
242
243
    'trial':{
        'command': setType('command', str),
        'codeDir': setPathCheck('codeDir'),
        Optional('gpuNum'): setNumberRange('gpuNum', int, 0, 99999),
        Optional('nasMode'): setChoice('nasMode', 'classic_mode', 'enas_mode', 'oneshot_mode', 'darts_mode')
244
245
246
    }
}

247
pai_yarn_trial_schema = {
chicm-ms's avatar
chicm-ms committed
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
    'trial':{
        'command': setType('command', str),
        'codeDir': setPathCheck('codeDir'),
        'gpuNum': setNumberRange('gpuNum', int, 0, 99999),
        'cpuNum': setNumberRange('cpuNum', int, 0, 99999),
        'memoryMB': setType('memoryMB', int),
        'image': setType('image', str),
        Optional('authFile'): And(os.path.exists, error=SCHEMA_PATH_ERROR % 'authFile'),
        Optional('shmMB'): setType('shmMB', int),
        Optional('dataDir'): And(Regex(r'hdfs://(([0-9]{1,3}.){3}[0-9]{1,3})(:[0-9]{2,5})?(/.*)?'),\
                            error='ERROR: dataDir format error, dataDir format is hdfs://xxx.xxx.xxx.xxx:xxx'),
        Optional('outputDir'): And(Regex(r'hdfs://(([0-9]{1,3}.){3}[0-9]{1,3})(:[0-9]{2,5})?(/.*)?'),\
                            error='ERROR: outputDir format error, outputDir format is hdfs://xxx.xxx.xxx.xxx:xxx'),
        Optional('virtualCluster'): setType('virtualCluster', str),
        Optional('nasMode'): setChoice('nasMode', 'classic_mode', 'enas_mode', 'oneshot_mode', 'darts_mode'),
        Optional('portList'): [{
            "label": setType('label', str),
            "beginAt": setType('beginAt', int),
            "portNumber": setType('portNumber', int)
        }]
268
269
270
    }
}

271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
pai_yarn_config_schema = {
    'paiYarnConfig': Or({
        'userName': setType('userName', str),
        'passWord': setType('passWord', str),
        'host': setType('host', str)
    }, {
        'userName': setType('userName', str),
        'token': setType('token', str),
        'host': setType('host', str)
    })
}


pai_trial_schema = {
    'trial':{
        'codeDir': setPathCheck('codeDir'),
        'nniManagerNFSMountPath': setPathCheck('nniManagerNFSMountPath'),
        'containerNFSMountPath': setType('containerNFSMountPath', str),
289
        Optional('command'): setType('command', str),
SparkSnail's avatar
SparkSnail committed
290
291
292
293
294
        Optional('gpuNum'): setNumberRange('gpuNum', int, 0, 99999),
        Optional('cpuNum'): setNumberRange('cpuNum', int, 0, 99999),
        Optional('memoryMB'): setType('memoryMB', int),
        Optional('image'): setType('image', str),
        Optional('virtualCluster'): setType('virtualCluster', str),
295
        Optional('paiStorageConfigName'): setType('paiStorageConfigName', str),
SparkSnail's avatar
SparkSnail committed
296
        Optional('paiConfigPath'): And(os.path.exists, error=SCHEMA_PATH_ERROR % 'paiConfigPath')
297
298
299
    }
}

300
pai_config_schema = {
301
    'paiConfig': Or({
302
303
304
        'userName': setType('userName', str),
        'passWord': setType('passWord', str),
        'host': setType('host', str)
305
306
307
308
309
    }, {
        'userName': setType('userName', str),
        'token': setType('token', str),
        'host': setType('host', str)
    })
310
311
}

George Cheng's avatar
George Cheng committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
dlts_trial_schema = {
    'trial':{
        'command': setType('command', str),
        'codeDir': setPathCheck('codeDir'),
        'gpuNum': setNumberRange('gpuNum', int, 0, 99999),
        'image': setType('image', str),
    }
}

dlts_config_schema = {
    'dltsConfig': {
        'dashboard': setType('dashboard', str),

        Optional('cluster'): setType('cluster', str),
        Optional('team'): setType('team', str),

        Optional('email'): setType('email', str),
        Optional('password'): setType('password', str),
    }
}

333
kubeflow_trial_schema = {
chicm-ms's avatar
chicm-ms committed
334
    'trial':{
335
        'codeDir':  setPathCheck('codeDir'),
336
        Optional('nasMode'): setChoice('nasMode', 'classic_mode', 'enas_mode', 'oneshot_mode', 'darts_mode'),
337
        Optional('ps'): {
338
339
340
341
342
            'replicas': setType('replicas', int),
            'command': setType('command', str),
            'gpuNum': setNumberRange('gpuNum', int, 0, 99999),
            'cpuNum': setNumberRange('cpuNum', int, 0, 99999),
            'memoryMB': setType('memoryMB', int),
343
344
            'image': setType('image', str),
            Optional('privateRegistryAuthPath'): And(os.path.exists, error=SCHEMA_PATH_ERROR % 'privateRegistryAuthPath')
345
        },
346
        Optional('master'): {
347
348
349
350
351
            'replicas': setType('replicas', int),
            'command': setType('command', str),
            'gpuNum': setNumberRange('gpuNum', int, 0, 99999),
            'cpuNum': setNumberRange('cpuNum', int, 0, 99999),
            'memoryMB': setType('memoryMB', int),
352
353
            'image': setType('image', str),
            Optional('privateRegistryAuthPath'): And(os.path.exists, error=SCHEMA_PATH_ERROR % 'privateRegistryAuthPath')
354
        },
355
        Optional('worker'):{
356
357
358
359
360
            'replicas': setType('replicas', int),
            'command': setType('command', str),
            'gpuNum': setNumberRange('gpuNum', int, 0, 99999),
            'cpuNum': setNumberRange('cpuNum', int, 0, 99999),
            'memoryMB': setType('memoryMB', int),
361
362
            'image': setType('image', str),
            Optional('privateRegistryAuthPath'): And(os.path.exists, error=SCHEMA_PATH_ERROR % 'privateRegistryAuthPath')
363
        }
364
365
366
367
    }
}

kubeflow_config_schema = {
SparkSnail's avatar
SparkSnail committed
368
    'kubeflowConfig':Or({
369
370
371
        'operator': setChoice('operator', 'tf-operator', 'pytorch-operator'),
        'apiVersion': setType('apiVersion', str),
        Optional('storage'): setChoice('storage', 'nfs', 'azureStorage'),
372
        'nfs': {
373
374
            'server': setType('server', str),
            'path': setType('path', str)
375
        }
chicm-ms's avatar
chicm-ms committed
376
    }, {
377
378
379
        'operator': setChoice('operator', 'tf-operator', 'pytorch-operator'),
        'apiVersion': setType('apiVersion', str),
        Optional('storage'): setChoice('storage', 'nfs', 'azureStorage'),
SparkSnail's avatar
SparkSnail committed
380
        'keyVault': {
381
382
383
384
            'vaultName': And(Regex('([0-9]|[a-z]|[A-Z]|-){1,127}'),\
                         error='ERROR: vaultName format error, vaultName support using (0-9|a-z|A-Z|-)'),
            'name': And(Regex('([0-9]|[a-z]|[A-Z]|-){1,127}'),\
                    error='ERROR: name format error, name support using (0-9|a-z|A-Z|-)')
SparkSnail's avatar
SparkSnail committed
385
386
        },
        'azureStorage': {
387
388
389
390
            'accountName': And(Regex('([0-9]|[a-z]|[A-Z]|-){3,31}'),\
                           error='ERROR: accountName format error, accountName support using (0-9|a-z|A-Z|-)'),
            'azureShare': And(Regex('([0-9]|[a-z]|[A-Z]|-){3,63}'),\
                          error='ERROR: azureShare format error, azureShare support using (0-9|a-z|A-Z|-)')
391
392
        },
        Optional('uploadRetryCount'): setNumberRange('uploadRetryCount', int, 1, 99999)
SparkSnail's avatar
SparkSnail committed
393
    })
394
395
}

396
397
frameworkcontroller_trial_schema = {
    'trial':{
398
        'codeDir':  setPathCheck('codeDir'),
399
        'taskRoles': [{
400
401
            'name': setType('name', str),
            'taskNum': setType('taskNum', int),
402
            'frameworkAttemptCompletionPolicy': {
403
404
                'minFailedTaskCount': setType('minFailedTaskCount', int),
                'minSucceededTaskCount': setType('minSucceededTaskCount', int),
405
            },
406
407
408
409
            'command': setType('command', str),
            'gpuNum': setNumberRange('gpuNum', int, 0, 99999),
            'cpuNum': setNumberRange('cpuNum', int, 0, 99999),
            'memoryMB': setType('memoryMB', int),
410
411
            'image': setType('image', str),
            Optional('privateRegistryAuthPath'): And(os.path.exists, error=SCHEMA_PATH_ERROR % 'privateRegistryAuthPath')
412
413
414
415
416
417
        }]
    }
}

frameworkcontroller_config_schema = {
    'frameworkcontrollerConfig':Or({
418
419
        Optional('storage'): setChoice('storage', 'nfs', 'azureStorage'),
        Optional('serviceAccountName'): setType('serviceAccountName', str),
420
        'nfs': {
421
422
            'server': setType('server', str),
            'path': setType('path', str)
423
        }
chicm-ms's avatar
chicm-ms committed
424
    }, {
425
426
        Optional('storage'): setChoice('storage', 'nfs', 'azureStorage'),
        Optional('serviceAccountName'): setType('serviceAccountName', str),
427
        'keyVault': {
428
429
430
431
            'vaultName': And(Regex('([0-9]|[a-z]|[A-Z]|-){1,127}'),\
                         error='ERROR: vaultName format error, vaultName support using (0-9|a-z|A-Z|-)'),
            'name': And(Regex('([0-9]|[a-z]|[A-Z]|-){1,127}'),\
                    error='ERROR: name format error, name support using (0-9|a-z|A-Z|-)')
432
433
        },
        'azureStorage': {
434
435
436
437
            'accountName': And(Regex('([0-9]|[a-z]|[A-Z]|-){3,31}'),\
                           error='ERROR: accountName format error, accountName support using (0-9|a-z|A-Z|-)'),
            'azureShare': And(Regex('([0-9]|[a-z]|[A-Z]|-){3,63}'),\
                          error='ERROR: azureShare format error, azureShare support using (0-9|a-z|A-Z|-)')
438
439
        },
        Optional('uploadRetryCount'): setNumberRange('uploadRetryCount', int, 1, 99999)
440
441
442
    })
}

443
machine_list_schema = {
444
445
    Optional('machineList'):[Or(
        {
chicm-ms's avatar
chicm-ms committed
446
447
448
449
450
451
452
453
            'ip': setType('ip', str),
            Optional('port'): setNumberRange('port', int, 1, 65535),
            'username': setType('username', str),
            'sshKeyPath': setPathCheck('sshKeyPath'),
            Optional('passphrase'): setType('passphrase', str),
            Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
            Optional('maxTrialNumPerGpu'): setType('maxTrialNumPerGpu', int),
            Optional('useActiveGpu'): setType('useActiveGpu', bool)
454
455
456
457
458
459
460
461
462
        },
        {
            'ip': setType('ip', str),
            Optional('port'): setNumberRange('port', int, 1, 65535),
            'username': setType('username', str),
            'passwd': setType('passwd', str),
            Optional('gpuIndices'): Or(int, And(str, lambda x: len([int(i) for i in x.split(',')]) > 0), error='gpuIndex format error!'),
            Optional('maxTrialNumPerGpu'): setType('maxTrialNumPerGpu', int),
            Optional('useActiveGpu'): setType('useActiveGpu', bool)
chicm-ms's avatar
chicm-ms committed
463
        })]
464
}
465
466
467

LOCAL_CONFIG_SCHEMA = Schema({**common_schema, **common_trial_schema})

468
REMOTE_CONFIG_SCHEMA = Schema({**common_schema, **common_trial_schema, **machine_list_schema})
469

470
471
PAI_CONFIG_SCHEMA = Schema({**common_schema, **pai_trial_schema, **pai_config_schema})

472
473
PAI_YARN_CONFIG_SCHEMA = Schema({**common_schema, **pai_yarn_trial_schema, **pai_yarn_config_schema})

George Cheng's avatar
George Cheng committed
474
475
DLTS_CONFIG_SCHEMA = Schema({**common_schema, **dlts_trial_schema, **dlts_config_schema})

476
KUBEFLOW_CONFIG_SCHEMA = Schema({**common_schema, **kubeflow_trial_schema, **kubeflow_config_schema})
477
478

FRAMEWORKCONTROLLER_CONFIG_SCHEMA = Schema({**common_schema, **frameworkcontroller_trial_schema, **frameworkcontroller_config_schema})