OLS-1229: HuggingFace environment settings

Signed-off-by: Pavel Tisnovsky <[email protected]>

Author: tisnik

ols/utils/environments.py

@@ -3,6 +3,8 @@
 import os
 import tempfile
 
+import ols.app.models.config as config_model
+
 
 def configure_gradio_ui_envs() -> None:
     """Configure GradioUI framework environment variables."""
@@ -15,3 +17,17 @@ def configure_gradio_ui_envs() -> None:
     # Fixes: https://issues.redhat.com/browse/OLS-301
     tempdir = os.path.join(tempfile.gettempdir(), "matplotlib")
     os.environ["MPLCONFIGDIR"] = tempdir
+
+
+def configure_hugging_face_envs(ols_config: config_model.OLSConfig) -> None:
+    """Configure HuggingFace library environment variables."""
+    if (
+        ols_config
+        and hasattr(ols_config, "reference_content")
+        and hasattr(ols_config.reference_content, "embeddings_model_path")
+        and ols_config.reference_content.embeddings_model_path
+    ):
+        os.environ["TRANSFORMERS_CACHE"] = str(
+            ols_config.reference_content.embeddings_model_path
+        )
+        os.environ["TRANSFORMERS_OFFLINE"] = "1"

runner.py

@@ -5,28 +5,13 @@
 import threading
 from pathlib import Path
 
-import ols.app.models.config as config_model
 from ols.runners.uvicorn import start_uvicorn
 from ols.src.auth.auth import use_k8s_auth
 from ols.utils.certificates import generate_certificates_file
-from ols.utils.environments import configure_gradio_ui_envs
+from ols.utils.environments import configure_gradio_ui_envs, configure_hugging_face_envs
 from ols.utils.logging_configurator import configure_logging
 
 
-def configure_hugging_face_envs(ols_config: config_model.OLSConfig) -> None:
-    """Configure HuggingFace library environment variables."""
-    if (
-        ols_config
-        and hasattr(ols_config, "reference_content")
-        and hasattr(ols_config.reference_content, "embeddings_model_path")
-        and ols_config.reference_content.embeddings_model_path
-    ):
-        os.environ["TRANSFORMERS_CACHE"] = str(
-            ols_config.reference_content.embeddings_model_path
-        )
-        os.environ["TRANSFORMERS_OFFLINE"] = "1"
-
-
 def load_index():
     """Load the index."""
     # accessing the config's rag_index property will trigger the loading

tests/unit/utils/test_environments.py

@@ -3,7 +3,8 @@
 import os
 from unittest.mock import patch
 
-from ols.utils.environments import configure_gradio_ui_envs
+from ols.app.models.config import OLSConfig, ReferenceContent
+from ols.utils.environments import configure_gradio_ui_envs, configure_hugging_face_envs
 
 
 @patch.dict(os.environ, {"GRADIO_ANALYTICS_ENABLED": "", "MPLCONFIGDIR": ""})
@@ -19,3 +20,40 @@ def test_configure_gradio_ui_envs():
     # expected environment variables
     assert os.environ.get("GRADIO_ANALYTICS_ENABLED", None) == "false"
     assert os.environ.get("MPLCONFIGDIR", None) != ""
+
+
+@patch.dict(os.environ, {"TRANSFORMERS_CACHE": "", "TRANSFORMERS_OFFLINE": ""})
+def test_configure_hugging_face_env_no_reference_content_set():
+    """Test the function configure_hugging_face_envs."""
+    # setup before tested function is called
+    assert os.environ.get("TRANSFORMERS_CACHE", None) == ""
+    assert os.environ.get("TRANSFORMERS_OFFLINE", None) == ""
+
+    ols_config = OLSConfig()
+    ols_config.reference_content = None
+
+    # call the tested function
+    configure_hugging_face_envs(ols_config)
+
+    # expected environment variables
+    assert os.environ.get("TRANSFORMERS_CACHE", None) == ""
+    assert os.environ.get("TRANSFORMERS_OFFLINE", None) == ""
+
+
+@patch.dict(os.environ, {"TRANSFORMERS_CACHE": "", "TRANSFORMERS_OFFLINE": ""})
+def test_configure_hugging_face_env_reference_content_set():
+    """Test the function configure_hugging_face_envs."""
+    # setup before tested function is called
+    assert os.environ.get("TRANSFORMERS_CACHE", None) == ""
+    assert os.environ.get("TRANSFORMERS_OFFLINE", None) == ""
+
+    ols_config = OLSConfig()
+    ols_config.reference_content = ReferenceContent()
+    ols_config.reference_content.embeddings_model_path = "foo"
+
+    # call the tested function
+    configure_hugging_face_envs(ols_config)
+
+    # expected environment variables
+    assert os.environ.get("TRANSFORMERS_CACHE", None) == "foo"
+    assert os.environ.get("TRANSFORMERS_OFFLINE", None) == "1"