From 1b86503d1edb93a9ab43adb98e04eba37213baaf Mon Sep 17 00:00:00 2001
From: Yifan Xiong <yifan.xiong@microsoft.com>
Date: Tue, 1 Nov 2022 10:46:19 +0800
Subject: [PATCH] CLI - Add non-zero return code for `sb [deploy,run]` (#425)

Add non-zero return code for `sb deploy` and `sb run` command when
there're Ansible failures in control plane.
Return code is set to count of failure.

For failures caused by benchmarks, return code is still set per benchmark
in results json file.
---
 superbench/cli/_handler.py   |  5 +++++
 superbench/runner/ansible.py |  2 ++
 superbench/runner/runner.py  |  8 ++++++++
 tests/cli/test_sb.py         | 13 ++++++++++---
 tests/runner/test_runner.py  |  4 ++++
 5 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/superbench/cli/_handler.py b/superbench/cli/_handler.py
index 03dba33b..859830fd 100644
--- a/superbench/cli/_handler.py
+++ b/superbench/cli/_handler.py
@@ -3,6 +3,7 @@
 
 """SuperBench CLI command handler."""
 
+import sys
 from pathlib import Path
 from importlib_metadata import version, PackageNotFoundError
 
@@ -252,6 +253,8 @@ def deploy_command_handler(
 
     runner = SuperBenchRunner(sb_config, docker_config, ansible_config, sb_output_dir)
     runner.deploy()
+    if runner.get_failure_count() != 0:
+        sys.exit(runner.get_failure_count())
 
 
 def run_command_handler(
@@ -307,3 +310,5 @@ def run_command_handler(
 
     runner = SuperBenchRunner(sb_config, docker_config, ansible_config, sb_output_dir)
     runner.run()
+    if runner.get_failure_count() != 0:
+        sys.exit(runner.get_failure_count())
diff --git a/superbench/runner/ansible.py b/superbench/runner/ansible.py
index f9f7b0fe..c012edc5 100644
--- a/superbench/runner/ansible.py
+++ b/superbench/runner/ansible.py
@@ -27,6 +27,7 @@ class AnsibleClient():
             'cmdline': '--forks 128',
         }
         self._head_host = None
+        self.failure_count = 0
         if config:
             inventory_file = getattr(config, 'host_file', None)
             inventory_list = getattr(config, 'host_list', None)
@@ -77,6 +78,7 @@ class AnsibleClient():
         if r.rc == 0:
             logger.info('Run succeed, return code {}.'.format(r.rc))
         else:
+            self.failure_count += 1
             logger.warning('Run failed, return code {}.'.format(r.rc))
         return r.rc
 
diff --git a/superbench/runner/runner.py b/superbench/runner/runner.py
index 3d4de0eb..cde88636 100644
--- a/superbench/runner/runner.py
+++ b/superbench/runner/runner.py
@@ -160,6 +160,14 @@ class SuperBenchRunner():
             logger.warning('Unknown mode %s.', mode.name)
         return mode_command.strip()
 
+    def get_failure_count(self):
+        """Get failure count during Ansible run.
+
+        Return:
+            int: Failure count.
+        """
+        return self._ansible_client.failure_count
+
     def deploy(self):    # pragma: no cover
         """Deploy SuperBench environment."""
         logger.info('Preparing SuperBench environment.')
diff --git a/tests/cli/test_sb.py b/tests/cli/test_sb.py
index fb975e08..d817bd86 100644
--- a/tests/cli/test_sb.py
+++ b/tests/cli/test_sb.py
@@ -8,6 +8,7 @@ import contextlib
 from functools import wraps
 from knack.testsdk import ScenarioTest, StringContainCheck, NoneCheck, JMESPathCheck
 from pathlib import Path
+from unittest import mock
 
 import superbench
 from superbench.cli import SuperBenchCLI
@@ -53,8 +54,10 @@ class SuperBenchCLIScenarioTest(ScenarioTest):
         """Test sb version."""
         self.cmd('sb version', checks=[StringContainCheck(superbench.__version__)])
 
-    def test_sb_deploy(self):
+    @mock.patch('superbench.runner.SuperBenchRunner.get_failure_count')
+    def test_sb_deploy(self, mocked_failure_count):
         """Test sb deploy."""
+        mocked_failure_count.return_value = 0
         self.cmd('sb deploy --host-list localhost', checks=[NoneCheck()])
 
     def test_sb_deploy_no_host(self):
@@ -65,12 +68,16 @@ class SuperBenchCLIScenarioTest(ScenarioTest):
         """Test sb exec."""
         self.cmd('sb exec --config-override superbench.enable=["none"]', checks=[NoneCheck()])
 
-    def test_sb_run(self):
+    @mock.patch('superbench.runner.SuperBenchRunner.get_failure_count')
+    def test_sb_run(self, mocked_failure_count):
         """Test sb run."""
+        mocked_failure_count.return_value = 0
         self.cmd('sb run --host-list localhost --config-override superbench.enable=none', checks=[NoneCheck()])
 
-    def test_sb_run_skipdocker(self):
+    @mock.patch('superbench.runner.SuperBenchRunner.get_failure_count')
+    def test_sb_run_skipdocker(self, mocked_failure_count):
         """Test sb run without docker."""
+        mocked_failure_count.return_value = 0
         self.cmd('sb run -l localhost -C superbench.enable=none --no-docker', checks=[NoneCheck()])
 
     def test_sb_run_no_docker_auth(self):
diff --git a/tests/runner/test_runner.py b/tests/runner/test_runner.py
index c9dd68ad..fa2bf162 100644
--- a/tests/runner/test_runner.py
+++ b/tests/runner/test_runner.py
@@ -41,6 +41,10 @@ class RunnerTestCase(unittest.TestCase):
         expected_log_file = Path(self.runner._sb_output_dir) / 'sb-run.log'
         self.assertTrue(expected_log_file.is_file())
 
+    def test_get_failure_count(self):
+        """Test get_failure_count."""
+        self.assertEqual(0, self.runner.get_failure_count())
+
     def test_get_mode_command(self):
         """Test __get_mode_command."""
         test_cases = [