CLI - Refine CLI handlers (#68)

* use absolute path of input file
* parse registry uri from image
* merge common parts for arguments processing

setup.py

@@ -146,6 +146,7 @@ def run(self):
             'mypy>=0.800',
             'pydocstyle>=5.1.1',
             'pytest-cov>=2.11.1',
+            'pytest-subtests>=0.4.0',
             'pytest>=6.2.2',
             'vcrpy>=4.1.1',
             'yapf>=0.30.0',

superbench/cli/_handler.py

@@ -21,42 +21,92 @@ def check_argument_file(name, file):
         name (str): argument name.
         file (str): file path.
 
+    Returns:
+        str: Absolute file path if it exists.
+
     Raises:
         CLIError: If file does not exist.
     """
-    if file and not Path(file).exists():
+    if file:
+        if not Path(file).exists():
             raise CLIError('{} {} does not exist.'.format(name, file))
+        return str(Path(file).resolve())
+    return file
 
 
-def version_command_handler():
-    """Print the current SuperBench tool version.
+def split_docker_domain(name):
+    """Split Docker image name to domain and remainder part.
+
+    Ported from https://github.com/distribution/distribution/blob/v2.7.1/reference/normalize.go#L62-L76.
+
+    Args:
+        name (str): Docker image name.
 
     Returns:
-        str: current SuperBench tool version.
+        str: Docker registry domain.
+        str: Remainder part.
     """
-    return superbench.__version__
+    legacy_default_domain = 'index.docker.io'
+    default_domain = 'docker.io'
 
+    i = name.find('/')
+    domain, remainder = '', ''
+    if i == -1 or ('.' not in name[:i] and ':' not in name[:i] and name[:i] != 'localhost'):
+        domain, remainder = default_domain, name
+    else:
+        domain, remainder = name[:i], name[i + 1:]
+    if domain == legacy_default_domain:
+        domain = default_domain
+    if domain == default_domain and '/' not in remainder:
+        remainder = 'library/{}'.format(remainder)
+    return domain, remainder
 
-def deploy_command_handler(
-    docker_image,
+
+def process_config_arguments(config_file=None, config_override=None):
+    """Process configuration arguments.
+
+    Args:
+        config_file (str, optional): Path to SuperBench config file. Defaults to None.
+        config_override (str, optional): Extra arguments to override config_file,
+            following [Hydra syntax](https://hydra.cc/docs/advanced/override_grammar/basic). Defaults to None.
+
+    Returns:
+        DictConfig: SuperBench config object.
+        str: Dir for output.
+
+    Raises:
+        CLIError: If input arguments are invalid.
+    """
+    config_file = check_argument_file('config_file', config_file)
+
+    # 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()
+
+    return sb_config, output_dir
+
+
+def process_runner_arguments(
+    docker_image='superbench/superbench',
     docker_username=None,
     docker_password=None,
     host_file=None,
     host_list=None,
     host_username=None,
     host_password=None,
-    private_key=None
+    private_key=None,
+    config_file=None,
+    config_override=None
 ):
-    """Deploy the SuperBench environments to all given nodes.
-
-    Deploy SuperBench environments on all nodes, including:
-    1. check drivers
-    2. install required system dependencies
-    3. install Docker and container runtime
-    4. pull Docker image
+    """Process runner related arguments.
 
     Args:
-        docker_image (str): Docker image URI.
+        docker_image (str, optional): Docker image URI. Defaults to superbench/superbench:latest.
         docker_username (str, optional): Docker registry username if authentication is needed. Defaults to None.
         docker_password (str, optional): Docker registry password if authentication is needed. Defaults to None.
         host_file (str, optional): Path to Ansible inventory host file. Defaults to None.
@@ -64,6 +114,15 @@ def deploy_command_handler(
         host_username (str, optional): Host username if needed. Defaults to None.
         host_password (str, optional): Host password or key passphase if needed. Defaults to None.
         private_key (str, optional): Path to private key if needed. Defaults to None.
+        config_file (str, optional): Path to SuperBench config file. Defaults to None.
+        config_override (str, optional): Extra arguments to override config_file,
+            following [Hydra syntax](https://hydra.cc/docs/advanced/override_grammar/basic). Defaults to None.
+
+    Returns:
+        DictConfig: SuperBench config object.
+        DictConfig: Docker config object.
+        DictConfig: Ansible config object.
+        str: Dir for output.
 
     Raises:
         CLIError: If input arguments are invalid.
@@ -72,21 +131,47 @@ def deploy_command_handler(
         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)
-    check_argument_file('private_key', private_key)
+    host_file = check_argument_file('host_file', host_file)
+    private_key = check_argument_file('private_key', private_key)
 
-    raise NotImplementedError
+    # Docker config
+    docker_config = OmegaConf.create(
+        {
+            'image': docker_image,
+            'username': docker_username,
+            'password': docker_password,
+            'registry': split_docker_domain(docker_image)[0],
+        }
+    )
+    # Ansible config
+    ansible_config = OmegaConf.create(
+        {
+            'host_file': host_file,
+            'host_list': host_list,
+            'host_username': host_username,
+            'host_password': host_password,
+            'private_key': private_key,
+        }
+    )
 
+    sb_config, output_dir = process_config_arguments(config_file=config_file, config_override=config_override)
 
-def exec_command_handler(
-    docker_image=None, docker_username=None, docker_password=None, config_file=None, config_override=None
-):
+    return docker_config, ansible_config, sb_config, output_dir
+
+
+def version_command_handler():
+    """Print the current SuperBench tool version.
+
+    Returns:
+        str: current SuperBench tool version.
+    """
+    return superbench.__version__
+
+
+def exec_command_handler(config_file=None, config_override=None):
     """Run the SuperBench benchmarks locally.
 
     Args:
-        docker_image (str, optional): Docker image URI.
-        docker_username (str, optional): Docker registry username if authentication is needed. Defaults to None.
-        docker_password (str, optional): Docker registry password if authentication is needed. Defaults to None.
         config_file (str, optional): Path to SuperBench config file. Defaults to None.
         config_override (str, optional): Extra arguments to override config_file,
             following [Hydra syntax](https://hydra.cc/docs/advanced/override_grammar/basic). Defaults to None.
@@ -94,29 +179,60 @@ def exec_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.')
-    check_argument_file('config_file', config_file)
+    sb_config, output_dir = process_config_arguments(config_file=config_file, config_override=config_override)
 
-    # Docker config
-    docker_config = OmegaConf.create()
-    for key in ['image', 'username', 'password']:
-        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)
+    executor = SuperBenchExecutor(sb_config, output_dir)
+    executor.exec()
 
-    # Create output directory
-    output_dir = create_output_dir()
 
-    executor = SuperBenchExecutor(sb_config, docker_config, output_dir)
-    executor.exec()
+def deploy_command_handler(
+    docker_image='superbench/superbench',
+    docker_username=None,
+    docker_password=None,
+    host_file=None,
+    host_list=None,
+    host_username=None,
+    host_password=None,
+    private_key=None
+):
+    """Deploy the SuperBench environments to all given nodes.
+
+    Deploy SuperBench environments on all nodes, including:
+    1. check drivers
+    2. install required system dependencies
+    3. install Docker and container runtime
+    4. pull Docker image
+
+    Args:
+        docker_image (str, optional): Docker image URI. Defaults to superbench/superbench:latest.
+        docker_username (str, optional): Docker registry username if authentication is needed. Defaults to None.
+        docker_password (str, optional): Docker registry password if authentication is needed. Defaults to None.
+        host_file (str, optional): Path to Ansible inventory host file. Defaults to None.
+        host_list (str, optional): Comma separated host list. Defaults to None.
+        host_username (str, optional): Host username if needed. Defaults to None.
+        host_password (str, optional): Host password or key passphase if needed. Defaults to None.
+        private_key (str, optional): Path to private key if needed. Defaults to None.
+
+    Raises:
+        CLIError: If input arguments are invalid.
+    """
+    docker_config, ansible_config, sb_config, output_dir = process_runner_arguments(
+        docker_image=docker_image,
+        docker_username=docker_username,
+        docker_password=docker_password,
+        host_file=host_file,
+        host_list=host_list,
+        host_username=host_username,
+        host_password=host_password,
+        private_key=private_key,
+    )
+
+    SuperBenchRunner(sb_config, docker_config, ansible_config, output_dir)
+    raise NotImplementedError
 
 
 def run_command_handler(
-    docker_image,
+    docker_image='superbench/superbench',
     docker_username=None,
     docker_password=None,
     host_file=None,
@@ -132,7 +248,7 @@ def run_command_handler(
     Run all benchmarks on given nodes.
 
     Args:
-        docker_image (str): Docker image URI.
+        docker_image (str, optional): Docker image URI. Defaults to superbench/superbench:latest.
         docker_username (str, optional): Docker registry username if authentication is needed. Defaults to None.
         docker_password (str, optional): Docker registry password if authentication is needed. Defaults to None.
         host_file (str, optional): Path to Ansible inventory host file. Defaults to None.
@@ -147,30 +263,18 @@ def run_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)
-    check_argument_file('private_key', private_key)
-    check_argument_file('config_file', config_file)
-
-    # Docker config
-    docker_config = OmegaConf.create()
-    for key in ['image', 'username', 'password']:
-        docker_config[key] = eval('docker_{}'.format(key))
-    # Ansible config
-    ansible_config = OmegaConf.create()
-    for key in ['file', 'list', 'username', 'password']:
-        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()
+    docker_config, ansible_config, sb_config, output_dir = process_runner_arguments(
+        docker_image=docker_image,
+        docker_username=docker_username,
+        docker_password=docker_password,
+        host_file=host_file,
+        host_list=host_list,
+        host_username=host_username,
+        host_password=host_password,
+        private_key=private_key,
+        config_file=config_file,
+        config_override=config_override,
+    )
 
     runner = SuperBenchRunner(sb_config, docker_config, ansible_config, output_dir)
     runner.run()

superbench/common/utils/file_handler.py

@@ -18,7 +18,7 @@ def create_output_dir():
         str: Output directory name.
     """
     output_name = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
-    output_path = Path('.', 'outputs', output_name)
+    output_path = Path('.', 'outputs', output_name).resolve()
     output_path.mkdir(mode=0o755, parents=True, exist_ok=True)
     return str(output_path)
 

superbench/executor/executor.py

@@ -13,16 +13,14 @@
 
 class SuperBenchExecutor():
     """SuperBench executor class."""
-    def __init__(self, sb_config, docker_config, output_dir):
+    def __init__(self, sb_config, output_dir):
         """Initilize.
 
         Args:
             sb_config (DictConfig): SuperBench config object.
-            docker_config (DictConfig): Docker config object.
             output_dir (str): Dir for output.
         """
         self._sb_config = sb_config
-        self._docker_config = docker_config
         self._output_dir = output_dir
 
         self.__set_logger('sb-exec.log')

tests/cli/test_handler.py

+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""CLI handler test."""
+
+import unittest
+
+import superbench.cli._handler as cli_handler
+
+
+class CLIHandlerTestCase(unittest.TestCase):
+    """A class for CLI handler test cases."""
+    def test_split_docker_domain(self):
+        """Test split_docker_domain function which splits Docker image name to domain and remainder part.
+
+        Test cases are ported from
+        https://github.com/distribution/distribution/blob/v2.7.1/reference/normalize_test.go#L468-L528.
+        """
+        test_cases = [
+            {
+                'input': 'test.com/foo',
+                'domain': 'test.com',
+                'name': 'foo',
+            },
+            {
+                'input': 'test_com/foo',
+                'domain': 'docker.io',
+                'name': 'test_com/foo',
+            },
+            {
+                'input': 'docker/migrator',
+                'domain': 'docker.io',
+                'name': 'docker/migrator',
+            },
+            {
+                'input': 'test.com:8080/foo',
+                'domain': 'test.com:8080',
+                'name': 'foo',
+            },
+            {
+                'input': 'test-com:8080/foo',
+                'domain': 'test-com:8080',
+                'name': 'foo',
+            },
+            {
+                'input': 'foo',
+                'domain': 'docker.io',
+                'name': 'library/foo',
+            },
+            {
+                'input': 'xn--n3h.com/foo',
+                'domain': 'xn--n3h.com',
+                'name': 'foo',
+            },
+            {
+                'input': 'xn--n3h.com:18080/foo',
+                'domain': 'xn--n3h.com:18080',
+                'name': 'foo',
+            },
+            {
+                'input': 'docker.io/foo',
+                'domain': 'docker.io',
+                'name': 'library/foo',
+            },
+            {
+                'input': 'docker.io/library/foo',
+                'domain': 'docker.io',
+                'name': 'library/foo',
+            },
+            {
+                'input': 'docker.io/library/foo/bar',
+                'domain': 'docker.io',
+                'name': 'library/foo/bar',
+            },
+        ]
+        for test_case in test_cases:
+            with self.subTest(msg='Testing with case', test_case=test_case):
+                domain, name = cli_handler.split_docker_domain(test_case['input'])
+                self.assertEqual(domain, test_case['domain'])
+                self.assertEqual(name, test_case['name'])

tests/cli/test_sb.py

@@ -53,13 +53,7 @@ def test_sb_version(self):
 
     def test_sb_deploy(self):
         """Test sb deploy."""
-        self.cmd('sb deploy --docker-image test:cuda11.1 --host-list localhost', expect_failure=True)
-
-    @capture_system_exit
-    def test_sb_deploy_no_docker_image(self):
-        """Test sb deploy, no --docker-image argument, should fail."""
-        self.cmd('sb deploy', expect_failure=True)
-        self.assertIn('sb deploy: error: the following arguments are required: --docker-image', self.stderr)
+        self.cmd('sb deploy --host-list localhost', expect_failure=True)
 
     def test_sb_exec(self):
         """Test sb exec."""
@@ -67,13 +61,12 @@ def test_sb_exec(self):
 
     def test_sb_run(self):
         """Test sb run."""
-        self.cmd('sb run --docker-image test:cuda11.1 --host-list localhost', checks=[NoneCheck()])
+        self.cmd('sb run --host-list localhost', checks=[NoneCheck()])
 
-    @capture_system_exit
-    def test_sb_run_no_docker_image(self):
-        """Test sb run, no --docker-image argument, should fail."""
-        self.cmd('sb run', expect_failure=True)
-        self.assertIn('sb run: error: the following arguments are required: --docker-image', self.stderr)
+    def test_sb_run_no_docker_auth(self):
+        """Test sb run, only --docker-username argument, should fail."""
+        result = self.cmd('sb run --docker-username test-user', expect_failure=True)
+        self.assertEqual(result.exit_code, 1)
 
     def test_sb_run_no_host(self):
         """Test sb run, no --host-file or --host-list, should fail."""
@@ -82,5 +75,5 @@ def test_sb_run_no_host(self):
 
     def test_sb_run_nonexist_host_file(self):
         """Test sb run, --host-file does not exist, should fail."""
-        result = self.cmd('sb run --docker-image test:cuda11.1 --host-file ./nonexist.yaml', expect_failure=True)
+        result = self.cmd('sb run --host-file ./nonexist.yaml', expect_failure=True)
         self.assertEqual(result.exit_code, 1)

tests/executor/test_executor.py

@@ -25,7 +25,7 @@ def setUp(self):
         self.default_config = OmegaConf.load(str(default_config_file))
         self.output_dir = tempfile.mkdtemp()
 
-        self.executor = SuperBenchExecutor(self.default_config, None, self.output_dir)
+        self.executor = SuperBenchExecutor(self.default_config, self.output_dir)
 
     def tearDown(self):
         """Hook method for deconstructing the test fixture after testing it."""