Address Copilot review

Author: Douglas Blank

comet_for_mlflow/comet_for_mlflow.py

@@ -130,28 +130,13 @@ def __init__(
         try:
             self.store = _get_store(mlflow_store_uri)
         except RestException as e:
-            # Check HTTP status code for authentication errors
-            status_code = (
-                e.get_http_status_code() if hasattr(e, "get_http_status_code") else None
-            )
-            error_msg = str(e)
-            if (
-                status_code == 401
-                or "401" in error_msg
-                or "Credential" in error_msg
-                or "authentication" in error_msg.lower()
-            ):
+            if self._is_authentication_error(e):
                 self._log_authentication_error(
                     mlflow_store_uri, "connecting to MLflow store"
                 )
             raise
         except Exception as e:
-            error_msg = str(e)
-            if (
-                "401" in error_msg
-                or "Credential" in error_msg
-                or "authentication" in error_msg.lower()
-            ):
+            if self._is_authentication_error(e):
                 self._log_authentication_error(
                     mlflow_store_uri, "connecting to MLflow store"
                 )
@@ -165,28 +150,13 @@ def __init__(
         try:
             self.mlflow_experiments = search_mlflow_store_experiments(self.store)
         except RestException as e:
-            # Check HTTP status code for authentication errors
-            status_code = (
-                e.get_http_status_code() if hasattr(e, "get_http_status_code") else None
-            )
-            error_msg = str(e)
-            if (
-                status_code == 401
-                or "401" in error_msg
-                or "Credential" in error_msg
-                or "authentication" in error_msg.lower()
-            ):
+            if self._is_authentication_error(e):
                 self._log_authentication_error(
                     mlflow_store_uri, "accessing MLflow experiments"
                 )
             raise
         except Exception as e:
-            error_msg = str(e)
-            if (
-                "401" in error_msg
-                or "Credential" in error_msg
-                or "authentication" in error_msg.lower()
-            ):
+            if self._is_authentication_error(e):
                 self._log_authentication_error(
                     mlflow_store_uri, "accessing MLflow experiments"
                 )
@@ -722,6 +692,22 @@ def get_api_key_or_login(self, api_key):
 
         return (api_key, None)
 
+    def _is_authentication_error(self, exception):
+        """Check if an exception is an authentication error (401)."""
+        error_msg = str(exception)
+        status_code = None
+
+        # Check HTTP status code for RestException
+        if hasattr(exception, "get_http_status_code"):
+            status_code = exception.get_http_status_code()
+
+        return (
+            status_code == 401
+            or "401" in error_msg
+            or "Credential" in error_msg
+            or "authentication" in error_msg.lower()
+        )
+
     def _log_authentication_error(self, mlflow_store_uri, context):
         """Log helpful error message for MLflow authentication errors."""
         LOGGER.error("")

examples/keras-example/run.py

@@ -11,7 +11,11 @@
 from tensorflow.keras.models import Sequential
 from tensorflow.keras.preprocessing.text import Tokenizer
 
-# When using when from databricks:
+# The following MLflow settings are specific to Databricks.
+# They should only be used when running in a Databricks environment.
+# If running outside Databricks, you may need to set a different tracking URI
+# and experiment, or remove these lines entirely.
+
 mlflow.set_tracking_uri("databricks")
 mlflow.set_experiment("/Shared/keras_example")