add nasMode in pai config (#1222)

microsoft · Jun 28, 2019 · 19a675f · 19a675f
1 parent 0063dee
commit 19a675f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tools/nni_cmd/config_schema.py b/tools/nni_cmd/config_schema.py
@@ -215,6 +215,7 @@ def setPathCheck(key):
     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('classic_mode', 'enas_mode', 'oneshot_mode')
     }
 }
 
@@ -229,6 +230,7 @@ def setPathCheck(key):
 kubeflow_trial_schema = {
 'trial':{
         'codeDir':  setPathCheck('codeDir'),
+        Optional('nasMode'): setChoice('classic_mode', 'enas_mode', 'oneshot_mode'),
         Optional('ps'): {
             'replicas': setType('replicas', int),
             'command': setType('command', str),