Don't attempt batching with InstructLab's llama-cpp-python

src/instructlab/sdg/llmblock.py

@@ -9,6 +9,7 @@
 # Third Party
 from datasets import Dataset
 from tqdm import tqdm
+import httpx
 import openai
 
 # Local
@@ -40,16 +41,31 @@ def server_supports_batched(client, model_id: str) -> bool:
     supported = getattr(client, "server_supports_batched", None)
     if supported is not None:
         return supported
-    try:
-        # Make a test call to the server to determine whether it supports
-        # multiple input prompts per request and also the n parameter
-        response = client.completions.create(
-            model=model_id, prompt=["test1", "test2"], max_tokens=1, n=3
-        )
-        # Number outputs should be 2 * 3 = 6
-        supported = len(response.choices) == 6
-    except openai.InternalServerError:
-        supported = False
+    # Start looking for InstructLab's default llama-cpp-python so we
+    # can avoid throwing an assertion error in the server, as
+    # llama-cpp-python does not like us explicitly testing batches
+    if "/v1" in client.base_url.path:
+        try:
+            # The root (without /v1) will have InstructLab's welcome
+            # message
+            http_res = client.get("../", cast_to=httpx.Response)
+            if "Hello from InstructLab" in http_res.text:
+                # The server is llama-cpp-python, so disable batching
+                supported = False
+        except openai.APIStatusError:
+            # The server is not InstructLab's llama-cpp-python
+            pass
+    if supported is None:
+        try:
+            # Make a test call to the server to determine whether it supports
+            # multiple input prompts per request and also the n parameter
+            response = client.completions.create(
+                model=model_id, prompt=["test1", "test2"], max_tokens=1, n=3
+            )
+            # Number outputs should be 2 * 3 = 6
+            supported = len(response.choices) == 6
+        except openai.InternalServerError:
+            supported = False
     client.server_supports_batched = supported
     logger.info(f"LLM server supports batched inputs: {client.server_supports_batched}")
     return supported

tests/test_llmblock.py

@@ -6,9 +6,11 @@
 
 # Third Party
 from datasets import Dataset, Features, Value
+from httpx import URL
+from openai import InternalServerError, NotFoundError, OpenAI
 
 # First Party
-from src.instructlab.sdg.llmblock import LLMBlock
+from src.instructlab.sdg.llmblock import LLMBlock, server_supports_batched
 
 
 class TestLLMBlockModelPrompt(unittest.TestCase):
@@ -103,3 +105,55 @@ def test_max_num_tokens_override(self, mock_load_config):
         )
         num_tokens = block.gen_kwargs["max_tokens"]
         assert num_tokens == 512
+
+    def test_server_supports_batched_llama_cpp(self):
+        resp_text = """{"message":"Hello from InstructLab! Visit us at https://instructlab.ai"}"""
+        mock_client = MagicMock()
+        mock_client.server_supports_batched = None
+        mock_client.base_url = URL("http://localhost:8000/v1")
+        mock_client.get = MagicMock()
+        mock_client.get.return_value = MagicMock()
+        mock_client.get().text = resp_text
+        self.mock_ctx.client = mock_client
+        supports_batched = server_supports_batched(self.mock_ctx.client, "my-model")
+        assert not supports_batched
+
+    def test_server_supports_batched_other_llama_cpp(self):
+        resp_text = "another server"
+        mock_client = MagicMock()
+        mock_client.server_supports_batched = None
+        mock_client.base_url = URL("http://localhost:8000/v1")
+        mock_client.get = MagicMock()
+        mock_client.get.return_value = MagicMock()
+        mock_client.get().text = resp_text
+        mock_completion = MagicMock()
+        mock_completion.create = MagicMock()
+        mock_completion.create.side_effect = InternalServerError(
+            "mock error",
+            response=MagicMock(),
+            body=MagicMock(),
+        )
+        mock_client.completions = mock_completion
+        self.mock_ctx.client = mock_client
+        supports_batched = server_supports_batched(self.mock_ctx.client, "my-model")
+        assert not supports_batched
+
+    def test_server_supports_batched_vllm(self):
+        mock_client = MagicMock()
+        mock_client.server_supports_batched = None
+        mock_client.base_url = URL("http://localhost:8000/v1")
+        mock_client.get = MagicMock()
+        mock_client.get.side_effect = NotFoundError(
+            "mock error",
+            response=MagicMock(),
+            body=MagicMock(),
+        )
+        mock_completion_resp = MagicMock()
+        mock_completion_resp.choices = ["a", "b", "c", "d", "e", "f"]
+        mock_completion = MagicMock()
+        mock_completion.create = MagicMock()
+        mock_completion.create.return_value = mock_completion_resp
+        mock_client.completions = mock_completion
+        self.mock_ctx.client = mock_client
+        supports_batched = server_supports_batched(self.mock_ctx.client, "my-model")
+        assert supports_batched