add tests

Author: fxmarty

.gitignore

@@ -1,2 +1,3 @@
 .idea
 target
+__pycache__/

backends/python/server/text_embeddings_server/layers/attention/__init__.py

@@ -2,7 +2,10 @@
 import os
 
 if os.getenv("USE_FLASH_ATTENTION", "").lower() == "false":
-    raise ImportError("`USE_FLASH_ATTENTION` is false.")
+    class Attention:
+        def __getattr__(self, name):
+            raise RuntimeError(f"TEI is used with USE_FLASH_ATTENTION=false, accessing `attention` is prohibited")
+    attention = Attention()
 if SYSTEM == "cuda":
     from .cuda import attention
 elif SYSTEM == "rocm":

backends/python/server/text_embeddings_server/models/flash_bert.py

@@ -233,6 +233,5 @@ def embed(self, batch: FlashBatch) -> List[Embedding]:
                 )
                 for i in range(len(batch))
             ]
-
         else:
             raise NotImplementedError(f"Pooling {self.pooling_mode} is not implemented in the python backend")

router/src/lib.rs

@@ -105,7 +105,7 @@ pub async fn run(
         serde_json::from_str(&config).context("Failed to parse `config.json`")?;
 
     // Set model type from config
-    let backend_model_type = get_backend_model_type(&config, &model_root, &pooling)?;
+    let (backend_model_type, inferred_pooling) = get_backend_model_type(&config, &model_root, &pooling)?;
 
     // Info model type
     let model_type = match &backend_model_type {
@@ -191,7 +191,7 @@ pub async fn run(
         }
     });
 
-    let pooling_str = match pooling {
+    let pooling_str = match inferred_pooling {
         Some(pool) => pool.to_string(),
         None => "none".to_string(),
     };
@@ -313,19 +313,19 @@ fn get_backend_model_type(
     config: &ModelConfig,
     model_root: &Path,
     pooling: &Option<text_embeddings_backend::Pool>,
-) -> Result<text_embeddings_backend::ModelType> {
+) -> Result<(text_embeddings_backend::ModelType, Option<text_embeddings_backend::Pool>)> {
     for arch in &config.architectures {
         if Some(text_embeddings_backend::Pool::Splade) == *pooling && arch.ends_with("MaskedLM") {
-            return Ok(text_embeddings_backend::ModelType::Embedding(
+            return Ok((text_embeddings_backend::ModelType::Embedding(
                 text_embeddings_backend::Pool::Splade,
-            ));
+            ), Some(text_embeddings_backend::Pool::Splade)));
         } else if arch.ends_with("Classification") {
             if pooling.is_some() {
                 tracing::warn!(
                     "`--pooling` arg is set but model is a classifier. Ignoring `--pooling` arg."
                 );
             }
-            return Ok(text_embeddings_backend::ModelType::Classifier);
+            return Ok((text_embeddings_backend::ModelType::Classifier, None));
         }
     }
 
@@ -353,7 +353,7 @@ fn get_backend_model_type(
             }
         }
     };
-    Ok(text_embeddings_backend::ModelType::Embedding(pool))
+    Ok((text_embeddings_backend::ModelType::Embedding(pool.clone()), Some(pool)))
 }
 
 #[derive(Debug, Deserialize)]

tests/__init__.py

@@ -0,0 +1,113 @@
+import pytest
+import asyncio
+import contextlib
+import random
+import os
+import tempfile
+import subprocess
+import shutil
+import sys
+from typing import Optional
+from aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError
+import requests
+import time
+from requests.exceptions import ConnectionError as RequestsConnectionError
+
+@pytest.fixture(scope="module")
+def event_loop():
+    loop = asyncio.get_event_loop()
+    yield loop
+    loop.close()
+
+class ProcessLauncherHandle:
+    def __init__(self, process, port: int):
+        self.port = port
+        self.process = process
+
+    def _inner_health(self) -> bool:
+        return self.process.poll() is None
+    
+    def health(self, timeout: int = 60):
+        assert timeout > 0
+        for _ in range(timeout):
+            if not self._inner_health():
+                raise RuntimeError("Launcher crashed")
+
+            try:
+                url = f"http://0.0.0.0:{self.port}/health"
+                headers = {"Content-Type": "application/json"}
+
+                response = requests.post(url, headers=headers)
+                return
+            except (ClientConnectorError, ClientOSError, ServerDisconnectedError, RequestsConnectionError) as e:
+                print("Connecting")
+                time.sleep(1)
+        raise RuntimeError("Health check failed")
+
+@pytest.fixture(scope="module")
+def launcher(event_loop):
+    @contextlib.contextmanager
+    def local_launcher(
+        model_id: str,
+        trust_remote_code: bool = False,
+        use_flash_attention: bool = True,
+        dtype: Optional[str] = None,
+        revision: Optional[str] = None,
+        pooling: Optional[str] = None,
+    ):
+        port = random.randint(8000, 10_000)
+        shard_uds_path = (
+            f"/tmp/tei-tests-{model_id.split('/')[-1]}-server"
+        )
+
+        args = [
+            "text-embeddings-router",
+            "--model-id",
+            model_id,
+            "--port",
+            str(port),
+            "--uds-path",
+            shard_uds_path,
+        ]
+
+        env = os.environ
+
+        if dtype is not None:
+            args.append("--dtype")
+            args.append(dtype)
+        if revision is not None:
+            args.append("--revision")
+            args.append(revision)
+        if trust_remote_code:
+            args.append("--trust-remote-code")
+        if pooling:
+            args.append("--pooling")
+            args.append(str(max_input_length))
+
+        env["LOG_LEVEL"] = "debug"
+
+        if not use_flash_attention:
+            env["USE_FLASH_ATTENTION"] = "false"
+
+        with tempfile.TemporaryFile("w+") as tmp:
+            # We'll output stdout/stderr to a temporary file. Using a pipe
+            # cause the process to block until stdout is read.
+            print("call subprocess.Popen, with args", args)
+            with subprocess.Popen(
+                args,
+                stdout=tmp,
+                stderr=subprocess.STDOUT,
+                env=env,
+            ) as process:
+                yield ProcessLauncherHandle(process, port)
+
+                process.terminate()
+                process.wait(60)
+
+                tmp.seek(0)
+                shutil.copyfileobj(tmp, sys.stderr)
+
+        if not use_flash_attention:
+            del env["USE_FLASH_ATTENTION"]
+    
+    return local_launcher

tests/assets/default_bert.pt

@@ -0,0 +1,2 @@
+[pytest]
+asyncio_mode = auto

tests/assets/flash_bert.pt

@@ -0,0 +1,28 @@
+import pytest
+import requests
+import json
+import torch
+
+@pytest.fixture(scope="module")
+def default_model_handle(launcher):
+    with launcher("sentence-transformers/all-MiniLM-L6-v2", use_flash_attention=False) as handle:
+        yield handle
+
+@pytest.fixture(scope="module")
+async def default_model(default_model_handle):
+    default_model_handle.health(300)
+    return default_model_handle
+
+@pytest.mark.asyncio
+@pytest.mark.private
+async def test_single_query(default_model):
+    url = f"http://0.0.0.0:{default_model.port}/embed"
+    data = {"inputs": "What is Deep Learning?"}
+    headers = {"Content-Type": "application/json"}
+
+    response = requests.post(url, json=data, headers=headers)
+
+    embedding = torch.Tensor(json.loads(response.text))
+    # reference_embedding = torch.load("assets/default_model.pt")
+
+    # assert torch.allclose(embedding, reference_embedding)