lazi-fy huggingface, langchain serve, litellm loading

garak/generators/huggingface.py

@@ -14,6 +14,7 @@
  https://huggingface.co/docs/api-inference/quicktour
 """
 
+import importlib
 import logging
 import re
 import os
@@ -22,8 +23,6 @@
 
 import backoff
 import torch
-from PIL import Image
-from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
 
 from garak import _config
 from garak.exception import ModelNameMissingError
@@ -579,14 +578,40 @@ class LLaVA(Generator):
         "llava-hf/llava-v1.6-mistral-7b-hf",
     ]
 
+    # avoid attempt to pickle the client attribute
+    def __getstate__(self) -> object:
+        self._clear_client()
+        return dict(self.__dict__)
+
+    # restore the client attribute
+    def __setstate__(self, d) -> object:
+        self.__dict__.update(d)
+        self._load_client()
+
+    def _load_client(self):
+        PIL = importlib.import_module("PIL")
+        self.Image = PIL.Image
+
+        transformers = importlib.import_module("transformers")
+        self.LlavaNextProcessor = transformers.LlavaNextProcessor
+        self.LlavaNextForConditionalGeneration = (
+            transformers.LlavaNextForConditionalGeneration
+        )
+
+    def _clear_client(self):
+        self.Image = None
+        self.LlavaNextProcessor = None
+        self.LlavaNextForConditionalGeneration = None
+
     def __init__(self, name="", generations=10, config_root=_config):
         super().__init__(name, generations=generations, config_root=config_root)
         if self.name not in self.supported_models:
             raise ModelNameMissingError(
                 f"Invalid modal name {self.name}, current support: {self.supported_models}."
             )
-        self.processor = LlavaNextProcessor.from_pretrained(self.name)
-        self.model = LlavaNextForConditionalGeneration.from_pretrained(
+
+        self.processor = self.LlavaNextProcessor.from_pretrained(self.name)
+        self.model = self.LlavaNextForConditionalGeneration.from_pretrained(
             self.name,
             torch_dtype=self.torch_dtype,
             low_cpu_mem_usage=self.low_cpu_mem_usage,
@@ -597,15 +622,16 @@ def __init__(self, name="", generations=10, config_root=_config):
             raise RuntimeError(
                 "CUDA is not supported on this device. Please make sure CUDA is installed and configured properly."
             )
+        self._load_client()
 
     def generate(
         self, prompt: str, generations_this_call: int = 1
     ) -> List[Union[str, None]]:
         text_prompt = prompt["text"]
         try:
-            image_prompt = Image.open(prompt["image"])
-        except FileNotFoundError:
-            raise FileNotFoundError(f"Cannot open image {prompt['image']}.")
+            image_prompt = self.Image.open(prompt["image"])
+        except FileNotFoundError as exc:
+            raise FileNotFoundError(f"Cannot open image {prompt['image']}.") from exc
         except Exception as e:
             raise Exception(e)
 

garak/generators/litellm.py

@@ -31,23 +31,19 @@
 ```
 """
 
+import importlib
 import logging
 
 from os import getenv
 from typing import List, Union
 
 import backoff
 
-import litellm
 
 from garak import _config
+from garak.exception import APIKeyMissingError
 from garak.generators.base import Generator
 
-# Fix issue with Ollama which does not support `presence_penalty`
-litellm.drop_params = True
-# Suppress log messages from LiteLLM
-litellm.verbose_logger.disabled = True
-# litellm.set_verbose = True
 
 # Based on the param support matrix below:
 # https://docs.litellm.ai/docs/completion/input
@@ -109,6 +105,26 @@ class LiteLLMGenerator(Generator):
         "stop",
     )
 
+    # avoid attempt to pickle the client attribute
+    def __getstate__(self) -> object:
+        self._clear_client()
+        return dict(self.__dict__)
+
+    # restore the client attribute
+    def __setstate__(self, d) -> object:
+        self.__dict__.update(d)
+        self._load_client()
+
+    def _load_client(self):
+        self.litellm = importlib.import_module("litellm")
+        # Fix issue with Ollama which does not support `presence_penalty`
+        self.litellm.drop_params = True
+        # Suppress log messages from LiteLLM
+        self.litellm.verbose_logger.disabled = True
+
+    def _clear_client(self):
+        self.litellm = None
+
     def __init__(self, name: str = "", generations: int = 10, config_root=_config):
         self.name = name
         self.api_base = None
@@ -127,13 +143,10 @@ def __init__(self, name: str = "", generations: int = 10, config_root=_config):
             self.name, generations=self.generations, config_root=config_root
         )
 
-        if self.provider is None:
-            raise ValueError(
-                "litellm generator needs to have a provider value configured - see docs"
-            )
-        elif (
+        if (
             self.api_key is None
         ):  # TODO: special case where api_key is not always required
+            # TODO: add other providers
             if self.provider == "openai":
                 self.api_key = getenv(self.key_env_var, None)
                 if self.api_key is None:
@@ -142,6 +155,8 @@ def __init__(self, name: str = "", generations: int = 10, config_root=_config):
                         " or in the configuration file"
                     )
 
+        self._load_client()
+
     @backoff.on_exception(backoff.fibo, Exception, max_value=70)
     def _call_model(
         self, prompt: str, generations_this_call: int = 1
@@ -159,7 +174,7 @@ def _call_model(
             print(msg)
             return []
 
-        response = litellm.completion(
+        response = self.litellm.completion(
             model=self.name,
             messages=prompt,
             temperature=self.temperature,