Port trial examples' config file to v2 (#3721)

Co-authored-by: liuzhe <zhe.liu@microsoft.com>

nni/experiment/config/common.py

@@ -29,6 +29,8 @@ class _AlgorithmConfig(ConfigBase):
         super().validate()
         _validate_algo(self)
 
+    _canonical_rules = {'code_directory': util.canonical_path}
+
 @dataclass(init=False)
 class AlgorithmConfig(_AlgorithmConfig):
     name: str
@@ -37,7 +39,7 @@ class AlgorithmConfig(_AlgorithmConfig):
 @dataclass(init=False)
 class CustomAlgorithmConfig(_AlgorithmConfig):
     class_name: str
-    class_directory: Optional[PathLike] = '.'
+    code_directory: Optional[PathLike] = '.'
     class_args: Optional[Dict[str, Any]] = None
 
 
@@ -67,7 +69,7 @@ class ExperimentConfig(ConfigBase):
     debug: bool = False
     log_level: Optional[str] = None
     experiment_working_directory: PathLike = '~/nni-experiments'
-    tuner_gpu_indices: Optional[Union[List[int], str]] = None
+    tuner_gpu_indices: Union[List[int], str, int, None] = None
     tuner: Optional[_AlgorithmConfig] = None
     assessor: Optional[_AlgorithmConfig] = None
     advisor: Optional[_AlgorithmConfig] = None
@@ -137,7 +139,7 @@ _canonical_rules = {
     'trial_code_directory': util.canonical_path,
     'max_experiment_duration': lambda value: f'{util.parse_time(value)}s' if value is not None else None,
     'experiment_working_directory': util.canonical_path,
-    'tuner_gpu_indices': lambda value: [int(idx) for idx in value.split(',')] if isinstance(value, str) else value,
+    'tuner_gpu_indices': util.canonical_gpu_indices,
     'tuner': lambda config: None if config is None or config.name == '_none_' else config.canonical(),
     'assessor': lambda config: None if config is None or config.name == '_none_' else config.canonical(),
     'advisor': lambda config: None if config is None or config.name == '_none_' else config.canonical(),

nni/experiment/config/convert.py

@@ -249,13 +249,13 @@ def convert_algo(algo_type, v1, v2):
         v2_algo = AlgorithmConfig(name=builtin_name, class_args=class_args)
 
     else:
-        class_directory = util.canonical_path(v1_algo.pop('codeDir'))
+        code_directory = util.canonical_path(v1_algo.pop('codeDir'))
         class_file_name = v1_algo.pop('classFileName')
         assert class_file_name.endswith('.py')
         class_name = class_file_name[:-3] + '.' + v1_algo.pop('className')
         v2_algo = CustomAlgorithmConfig(
             class_name=class_name,
-            class_directory=class_directory,
+            code_directory=code_directory,
             class_args=class_args
         )
 

nni/experiment/config/local.py

@@ -5,6 +5,7 @@ from dataclasses import dataclass
 from typing import List, Optional, Union
 
 from .common import TrainingServiceConfig
+from . import util
 
 __all__ = ['LocalConfig']
 
@@ -13,10 +14,10 @@ class LocalConfig(TrainingServiceConfig):
     platform: str = 'local'
     use_active_gpu: Optional[bool] = None
     max_trial_number_per_gpu: int = 1
-    gpu_indices: Optional[Union[List[int], str]] = None
+    gpu_indices: Union[List[int], str, int, None] = None
 
     _canonical_rules = {
-        'gpu_indices': lambda value: [int(idx) for idx in value.split(',')] if isinstance(value, str) else value
+        'gpu_indices': util.canonical_gpu_indices
     }
 
     _validation_rules = {

nni/experiment/config/remote.py

@@ -22,12 +22,12 @@ class RemoteMachineConfig(ConfigBase):
     ssh_passphrase: Optional[str] = None
     use_active_gpu: bool = False
     max_trial_number_per_gpu: int = 1
-    gpu_indices: Optional[Union[List[int], str]] = None
+    gpu_indices: Union[List[int], str, int, None] = None
     python_path: Optional[str] = None
 
     _canonical_rules = {
         'ssh_key_file': util.canonical_path,
-        'gpu_indices': lambda value: [int(idx) for idx in value.split(',')] if isinstance(value, str) else value,
+        'gpu_indices': util.canonical_gpu_indices
     }
 
     _validation_rules = {

nni/experiment/config/util.py

@@ -92,3 +92,10 @@ def _parse_unit(string, target_unit, all_units):
             value = float(number) * factor
             return math.ceil(value / all_units[target_unit])
     raise ValueError(f'Unsupported unit in "{string}"')
+
+def canonical_gpu_indices(indices: Union[List[int], str, int, None]) -> Optional[List[int]]:
+    if isinstance(indices, str):
+        return [int(idx) for idx in indices.split(',')]
+    if isinstance(indices, int):
+        return [indices]
+    return indices

nni/tools/nnictl/launcher.py

@@ -403,9 +403,13 @@ def launch_experiment(args, experiment_config, mode, experiment_id, config_versi
         if not os.path.isdir(path):
             os.makedirs(path)
         path = tempfile.mkdtemp(dir=path)
-        nas_mode = experiment_config['trial'].get('nasMode', 'classic_mode')
-        code_dir = expand_annotations(experiment_config['trial']['codeDir'], path, nas_mode=nas_mode)
-        experiment_config['trial']['codeDir'] = code_dir
+        if config_version == 1:
+            nas_mode = experiment_config['trial'].get('nasMode', 'classic_mode')
+            code_dir = expand_annotations(experiment_config['trial']['codeDir'], path, nas_mode=nas_mode)
+            experiment_config['trial']['codeDir'] = code_dir
+        else:
+            code_dir = expand_annotations(experiment_config['trialCodeDirectory'], path)
+            experiment_config['trialCodeDirectory'] = code_dir
         search_space = generate_search_space(code_dir)
         experiment_config['searchSpace'] = search_space
         assert search_space, ERROR_INFO % 'Generated search space is empty'

nni/tools/package_utils/__init__.py

@@ -187,7 +187,7 @@ def create_customized_class_instance(class_params):
         Returns customized class instance.
     """
 
-    code_dir = class_params.get('classDirectory')
+    code_dir = class_params.get('codeDirectory')
     qualified_class_name = class_params.get('className')
     class_args = class_params.get('classArgs')