Graceful exit for aqua cli (#802)

Author: VipulMascarenhas

ads/aqua/__init__.py

@@ -8,7 +8,7 @@
 
 from ads import logger, set_auth
 from ads.aqua.utils import fetch_service_compartment
-from ads.config import NB_SESSION_OCID, OCI_RESOURCE_PRINCIPAL_VERSION
+from ads.config import OCI_RESOURCE_PRINCIPAL_VERSION
 
 ENV_VAR_LOG_LEVEL = "ADS_AQUA_LOG_LEVEL"
 
@@ -36,9 +36,3 @@ def set_log_level(log_level: str):
 ODSC_MODEL_COMPARTMENT_OCID = (
     os.environ.get("ODSC_MODEL_COMPARTMENT_OCID") or fetch_service_compartment()
 )
-if not ODSC_MODEL_COMPARTMENT_OCID:
-    if NB_SESSION_OCID:
-        logger.error(
-            f"Aqua is not available for this notebook session {NB_SESSION_OCID}."
-        )
-    exit()

ads/aqua/cli.py

@@ -4,12 +4,19 @@
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 import os
-
-from ads.aqua import ENV_VAR_LOG_LEVEL, set_log_level
+import sys
+
+from ads.aqua import (
+    ENV_VAR_LOG_LEVEL,
+    set_log_level,
+    ODSC_MODEL_COMPARTMENT_OCID,
+    logger,
+)
 from ads.aqua.deployment import AquaDeploymentApp
 from ads.aqua.evaluation import AquaEvaluationApp
 from ads.aqua.finetune import AquaFineTuningApp
 from ads.aqua.model import AquaModelApp
+from ads.config import NB_SESSION_OCID
 
 
 class AquaCommand:
@@ -41,3 +48,13 @@ def __init__(
             'WARNING', 'ERROR', and 'CRITICAL'.
         """
         set_log_level(log_level)
+        # gracefully exit if env var is not set
+        if not ODSC_MODEL_COMPARTMENT_OCID:
+            logger.error(
+                "ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
+            )
+            if NB_SESSION_OCID:
+                logger.error(
+                    f"Aqua is not available for the notebook session {NB_SESSION_OCID}."
+                )
+            sys.exit(1)

ads/aqua/utils.py

@@ -582,7 +582,7 @@ def fetch_service_compartment() -> Union[str, None]:
         )
     except AquaFileNotFoundError:
         logger.error(
-            f"ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
+            f"Config file {config_file_name}/{CONTAINER_INDEX} to fetch service compartment OCID could not be found."
         )
         return
     compartment_mapping = config.get(COMPARTMENT_MAPPING_KEY)

tests/unitary/with_extras/aqua/test_cli.py

@@ -4,26 +4,30 @@
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import os
 import logging
 import subprocess
 from unittest import TestCase
 from unittest.mock import patch
-
+from importlib import reload
 from parameterized import parameterized
 
+import ads.aqua
+import ads.config
 from ads.aqua.cli import AquaCommand
 
 
 class TestAquaCLI(TestCase):
     """Tests the AQUA CLI."""
 
-    DEFAUL_AQUA_CLI_LOGGING_LEVEL = "ERROR"
+    DEFAULT_AQUA_CLI_LOGGING_LEVEL = "ERROR"
     logger = logging.getLogger(__name__)
     logging.basicConfig(
         format="%(asctime)s %(module)s %(levelname)s: %(message)s",
         datefmt="%m/%d/%Y %I:%M:%S %p",
         level=logging.INFO,
     )
+    SERVICE_COMPARTMENT_ID = "ocid1.compartment.oc1..<OCID>"
 
     def test_entrypoint(self):
         """Tests CLI entrypoint."""
@@ -33,15 +37,55 @@ def test_entrypoint(self):
 
     @parameterized.expand(
         [
-            ("default", None, DEFAUL_AQUA_CLI_LOGGING_LEVEL),
+            ("default", None, DEFAULT_AQUA_CLI_LOGGING_LEVEL),
             ("set logging level", "info", "info"),
         ]
     )
-    @patch("ads.aqua.cli.set_log_level")
-    def test_aquacommand(self, name, arg, expected, mock_setting_log):
-        """Tests aqua command initailzation."""
-        if arg:
-            AquaCommand(arg)
-        else:
-            AquaCommand()
-        mock_setting_log.assert_called_with(expected)
+    def test_aquacommand(self, name, arg, expected):
+        """Tests aqua command initialization."""
+        with patch.dict(
+            os.environ,
+            {"ODSC_MODEL_COMPARTMENT_OCID": TestAquaCLI.SERVICE_COMPARTMENT_ID},
+        ):
+            reload(ads.config)
+            reload(ads.aqua)
+            reload(ads.aqua.cli)
+            with patch("ads.aqua.cli.set_log_level") as mock_setting_log:
+                if arg:
+                    AquaCommand(arg)
+                else:
+                    AquaCommand()
+                mock_setting_log.assert_called_with(expected)
+
+    @parameterized.expand(
+        [
+            ("default", None),
+            ("using jupyter instance", "nb-session-ocid"),
+        ]
+    )
+    def test_aqua_command_without_compartment_env_var(self, name, session_ocid):
+        """Test whether exit is called when ODSC_MODEL_COMPARTMENT_OCID is not set. Also check if NB_SESSION_OCID is
+        set then log the appropriate message."""
+
+        with patch("sys.exit") as mock_exit:
+            env_dict = {"ODSC_MODEL_COMPARTMENT_OCID": ""}
+            if session_ocid:
+                env_dict.update({"NB_SESSION_OCID": session_ocid})
+            with patch.dict(os.environ, env_dict):
+                reload(ads.config)
+                reload(ads.aqua)
+                reload(ads.aqua.cli)
+                with patch("ads.aqua.cli.set_log_level") as mock_setting_log:
+                    with patch("ads.aqua.logger.error") as mock_logger_error:
+                        AquaCommand()
+                        mock_setting_log.assert_called_with(
+                            TestAquaCLI.DEFAULT_AQUA_CLI_LOGGING_LEVEL
+                        )
+                        mock_logger_error.assert_any_call(
+                            "ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
+                        )
+                        if session_ocid:
+                            mock_logger_error.assert_any_call(
+                                f"Aqua is not available for the notebook session {session_ocid}."
+                            )
+                        mock_exit.assert_called_with(1)