Merge branch 'feature/aquav1.0.1' into update_telemetry_for_model_shapes

Author: VipulMascarenhas

ads/aqua/__init__.py

@@ -4,13 +4,11 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 
-import logging
 import os
-import sys
 
-from ads import set_auth
+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"
 
@@ -21,25 +19,7 @@ def get_logger_level():
     return level
 
 
-def configure_aqua_logger():
-    """Configures the AQUA logger."""
-    log_level = get_logger_level()
-    logger = logging.getLogger(__name__)
-    logger.setLevel(log_level)
-
-    handler = logging.StreamHandler(sys.stdout)
-    formatter = logging.Formatter(
-        "%(asctime)s - %(name)s.%(module)s - %(levelname)s - %(message)s"
-    )
-    handler.setFormatter(formatter)
-    handler.setLevel(log_level)
-
-    logger.addHandler(handler)
-    logger.propagate = False
-    return logger
-
-
-logger = configure_aqua_logger()
+logger.setLevel(get_logger_level())
 
 
 def set_log_level(log_level: str):
@@ -56,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)

ads/cli.py

@@ -8,6 +8,7 @@
 import sys
 
 import fire
+from dataclasses import is_dataclass
 from ads.common import logger
 
 try:
@@ -70,11 +71,28 @@ def _SeparateFlagArgs(args):
 fire.core.parser.SeparateFlagArgs = _SeparateFlagArgs
 
 
+def serialize(data):
+    """Serialize dataclass objects or lists of dataclass objects.
+    Parameters:
+        data: A dataclass object or a list of dataclass objects.
+    Returns:
+        None
+    Prints:
+        The string representation of each dataclass object.
+    """
+    if isinstance(data, list):
+        [print(str(item)) for item in data]
+    else:
+        print(str(data))
+
+
 def cli():
     if len(sys.argv) > 1 and sys.argv[1] == "aqua":
         from ads.aqua.cli import AquaCommand
 
-        fire.Fire(AquaCommand, command=sys.argv[2:], name="ads aqua")
+        fire.Fire(
+            AquaCommand, command=sys.argv[2:], name="ads aqua", serialize=serialize
+        )
     else:
         click_cli()
 

ads/common/serializer.py

@@ -195,7 +195,10 @@ def to_json(
             `None` in case when `uri` provided.
         """
         json_string = json.dumps(
-            self.to_dict(**kwargs), cls=encoder, default=default or self.serialize
+            self.to_dict(**kwargs),
+            cls=encoder,
+            default=default or self.serialize,
+            indent=4,
         )
         if uri:
             self._write_to_file(s=json_string, uri=uri, **kwargs)
@@ -463,9 +466,7 @@ def from_dict(
                 "These fields will be ignored."
             )
 
-        obj = cls(
-            **{key: obj_dict.get(key) for key in allowed_fields}
-        )
+        obj = cls(**{key: obj_dict.get(key) for key in allowed_fields})
 
         for key, value in obj_dict.items():
             if (

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)

tests/unitary/with_extras/aqua/test_global.py

@@ -6,7 +6,7 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
-from ads.aqua import configure_aqua_logger, get_logger_level, set_log_level
+from ads.aqua import get_logger_level, set_log_level
 
 
 class TestAquaLogging(unittest.TestCase):
@@ -22,18 +22,6 @@ def test_get_logger_level_from_env(self):
         """Test log level is correctly read from environment variable."""
         self.assertEqual(get_logger_level(), "DEBUG")
 
-    @patch("logging.getLogger")
-    @patch("logging.StreamHandler")
-    def test_configure_aqua_logger(self, mock_handler, mock_get_logger):
-        """Test that logger is correctly configured."""
-        mock_logger = MagicMock()
-        mock_get_logger.return_value = mock_logger
-
-        logger = configure_aqua_logger()
-
-        mock_get_logger.assert_called_once_with("ads.aqua")
-        mock_logger.setLevel.assert_called_with(self.DEFAULT_AQUA_LOG_LEVEL)
-
     @patch("ads.aqua.logger", create=True)
     def test_set_log_level(self, mock_logger):
         """Test that the log level of the logger is set correctly."""