CLI: Code Revision - Use omegaconf to replace hydra for configuration (#27)

Use omegaconf to replace hydra for configuration system:
* remove hydra
* use omegaconf to merge configurations

setup.py

@@ -133,9 +133,8 @@ def run(self):
     packages=find_packages(exclude=['tests']),
     python_requires='>=3.6, <4',
     install_requires=[
-        'hydra-colorlog>=1.0.0',
-        'hydra-core>=1.0.4',
         'knack>=0.7.2',
+        'omegaconf>=2.0.6',
     ],
     extras_require={
         'dev': ['pre-commit>=2.10.0'],

superbench/cli/__init__.py

@@ -3,6 +3,6 @@
 
 """SuperBench cli module."""
 
-from .sb import SuperBenchCLI
+from superbench.cli.sb import SuperBenchCLI
 
 __all__ = ['SuperBenchCLI']

superbench/cli/_commands.py

@@ -49,6 +49,7 @@ def load_arguments(self, command):
                 'config_override',
                 options_list=('--config-override', '-C'),
                 type=str,
-                help='Extra arguments to override config_file, following Hydra syntax.'
+                nargs='+',
+                help='Extra arguments to override config_file.'
             )
         super().load_arguments(command)

superbench/cli/_handler.py

@@ -4,13 +4,13 @@
 """SuperBench CLI command handler."""
 
 import os
-import yaml
 from pathlib import Path
 
 from knack.util import CLIError
+from omegaconf import OmegaConf
 
 import superbench
-from superbench.common.utils import get_sb_command, get_config, new_output_dir
+from superbench.common.utils import create_output_dir, get_sb_config
 
 
 def check_argument_file(name, file):
@@ -67,6 +67,8 @@ def deploy_command_handler(
     Raises:
         CLIError: If input arguments are invalid.
     """
+    if bool(docker_username) != bool(docker_password):
+        raise CLIError('Must specify both docker_username and docker_password if authentication is needed.')
     if not (host_file or host_list):
         raise CLIError('Must specify one of host_file or host_list.')
     check_argument_file('host_file', host_file)
@@ -95,15 +97,20 @@ def exec_command_handler(
         raise CLIError('Must specify both docker_username and docker_password if authentication is needed.')
     check_argument_file('config_file', config_file)
 
-    # Dump configs into outputs/date/config.merge.yaml
-    config = get_config(config_file)
-    config['docker'] = {}
+    # Docker config
+    docker_config = OmegaConf.create()
     for key in ['image', 'username', 'password']:
-        config['docker'][key] = eval('docker_{}'.format(key))
-    output_dir = new_output_dir()
-    with (Path(output_dir) / 'config.merge.yaml').open(mode='w') as f:
-        yaml.safe_dump(config, f)
-    os.system(get_sb_command('sb-exec', output_dir, config_override or ''))
+        docker_config[key] = eval('docker_{}'.format(key))
+    # SuperBench config
+    sb_config = get_sb_config(config_file)
+    if config_override:
+        sb_config_from_override = OmegaConf.from_dotlist(config_override)
+        sb_config = OmegaConf.merge(sb_config, sb_config_from_override)
+
+    # Create output directory
+    output_dir = create_output_dir()
+
+    os.system('sb-exec {}'.format(output_dir))
 
 
 def run_command_handler(
@@ -146,15 +153,21 @@ def run_command_handler(
     check_argument_file('private_key', private_key)
     check_argument_file('config_file', config_file)
 
-    # Dump configs into outputs/date/config.merge.yaml
-    config = get_config(config_file)
-    config['docker'] = {}
+    # Docker config
+    docker_config = OmegaConf.create()
     for key in ['image', 'username', 'password']:
-        config['docker'][key] = eval('docker_{}'.format(key))
-    config['ansible'] = {}
+        docker_config[key] = eval('docker_{}'.format(key))
+    # Ansible config
+    ansible_config = OmegaConf.create()
     for key in ['file', 'list', 'username', 'password']:
-        config['ansible']['host_{}'.format(key)] = eval('host_{}'.format(key))
-    output_dir = new_output_dir()
-    with (Path(output_dir) / 'config.merge.yaml').open(mode='w') as f:
-        yaml.safe_dump(config, f)
-    os.system(get_sb_command('sb-run', output_dir, config_override or ''))
+        ansible_config['host_{}'.format(key)] = eval('host_{}'.format(key))
+    # SuperBench config
+    sb_config = get_sb_config(config_file)
+    if config_override:
+        sb_config_from_override = OmegaConf.from_dotlist(config_override)
+        sb_config = OmegaConf.merge(sb_config, sb_config_from_override)
+
+    # Create output directory
+    output_dir = create_output_dir()
+
+    os.system('sb-run {}'.format(output_dir))

superbench/cli/sb_exec.py

@@ -5,15 +5,9 @@
 
 """SuperBench sb exec command."""
 
-import hydra
 
-from superbench.common.utils import logger
-
-
-@hydra.main(config_path='../config', config_name='default')
-def main(config):
+def main():
     """The main entrypoint for sb-exec."""
-    logger.info(config)
     # executor = SuperBenchExecutor(config)
     # executor.exec()
 

superbench/cli/sb_run.py

@@ -5,15 +5,9 @@
 
 """SuperBench sb run command."""
 
-import hydra
 
-from superbench.common.utils import logger
-
-
-@hydra.main(config_path='../config', config_name='default')
-def main(config):
+def main():
     """The main entrypoint for sb-run."""
-    logger.info(config)
     # runner = SuperBenchRunner(config)
     # runner.run()
 

superbench/common/utils/__init__.py

@@ -4,7 +4,6 @@
 """Exposes the interface of SuperBench common utilities."""
 
 from .logging import logger
-from .file_handler import new_output_dir, get_config
-from .command import get_sb_command
+from superbench.common.utils.file_handler import create_output_dir, get_sb_config
 
-__all__ = ['logger', 'new_output_dir', 'get_config', 'get_sb_command']
+__all__ = ['logger', 'create_output_dir', 'get_sb_config']

superbench/common/utils/command.py

@@ -2,23 +2,3 @@
 # Licensed under the MIT License.
 
 """Utilities for command."""
-
-
-def get_sb_command(cli, output_path, config_override):
-    """Get sb command for sb-run or sb-exec.
-
-    Args:
-        cli (str): CLI name.
-        output_path (str): Output directory path.
-        config_override (str): Extra arguments to override config.
-
-    Returns:
-        str: Command to run.
-    """
-    sb_cmd = '{cli} ' \
-        '--config-name=config.merge ' \
-        '--config-dir={path} ' \
-        'hydra.run.dir={path} ' \
-        'hydra.sweep.dir={path} ' \
-        '{args}'.format(cli=cli, path=output_path, args=config_override)
-    return sb_cmd.strip()

superbench/common/utils/file_handler.py

@@ -3,13 +3,14 @@
 
 """Utilities for file."""
 
-import yaml
 from pathlib import Path
 from datetime import datetime
 
+from omegaconf import OmegaConf
 
-def new_output_dir():
-    """Generate a new output directory.
+
+def create_output_dir():
+    """Create a new output directory.
 
     Generate a new output directory name based on current time and create it on filesystem.
 
@@ -22,7 +23,7 @@ def new_output_dir():
     return str(output_path)
 
 
-def get_config(config_file):
+def get_sb_config(config_file):
     """Read SuperBench config yaml.
 
     Read config file, use default config if None is provided.
@@ -31,12 +32,10 @@ def get_config(config_file):
         config_file (str): config file path.
 
     Returns:
-        dict: Config object, None if file does not exist.
+        OmegaConf: Config object, None if file does not exist.
     """
-    here = Path(__file__).parent.resolve()
-    p = Path(config_file) if config_file else here / '../../config/default.yaml'
+    default_config_file = Path(__file__).parent / '../../config/default.yaml'
+    p = Path(config_file) if config_file else default_config_file
     if not p.is_file():
         return None
-    with p.open() as f:
-        config = yaml.safe_load(f)
-    return config
+    return OmegaConf.load(str(p))