Repository: instructlab/instructlab

Issue:
Improve backend auto-detection for vLLM
Today, ilab will check whether the model path is a directory and the system is Linux. If so, it will activate the vLLM backend to serve/chat/gen models.
The detection could be improved by:

* check for the presence of `.tensors` file
* a specific .json file that vLLM looks for (need to investigate this)

This will make the auto-detection more robus.


Discussion:

Patch:
diff --git a/src/instructlab/model/backends/backends.py b/src/instructlab/model/backends/backends.py
index 97987242cf..bbf6162627 100644
--- a/src/instructlab/model/backends/backends.py
+++ b/src/instructlab/model/backends/backends.py
@@ -6,6 +6,7 @@
 from typing import Optional, Tuple
 import abc
 import contextlib
+import json
 import logging
 import mmap
 import multiprocessing
@@ -111,6 +112,50 @@ def safe_close_all(resources: typing.Iterable[Closeable]):
             resource.close()
 
 
+def is_model_safetensors(model_path: pathlib.Path) -> bool:
+    """Check if model_path is a valid safe tensors directory
+
+    Check if provided path to model represents directory containing a safetensors representation
+    of a model. Directory must contain a specific set of files to qualify as a safetensors model directory
+    Args:
+        model_path (Path): The path to the model directory
+    Returns:
+        bool: True if the model is a safetensors model, False otherwise.
+    """
+    try:
+        files = list(model_path.iterdir())
+    except (FileNotFoundError, NotADirectoryError, PermissionError) as e:
+        logger.debug("Failed to read directory: %s", e)
+        return False
+
+    if not any(file.suffix == ".safetensors" for file in files):
+        logger.debug("'%s' has no *.safetensors files", model_path)
+        return False
+
+    basenames = {file.name for file in files}
+    requires_files = {
+        "config.json",
+        "tokenizer.json",
+        "tokenizer.model",
+        "tokenizer_config.json",
+    }
+    diff = requires_files.difference(basenames)
+    if diff:
+        logger.debug("'%s' is missing %s", model_path, diff)
+        return False
+
+    for file in model_path.glob("*.json"):
+        try:
+            with file.open(encoding="utf-8") as f:
+                json.load(f)
+        except (PermissionError, json.JSONDecodeError) as e:
+            logger.debug("'%s' is not a valid JSON file: e", file, e)
+            return False
+
+    # TODO: add check for safetensors file header (?)
+    return True
+
+
 def is_model_gguf(model_path: pathlib.Path) -> bool:
     """
     Check if the file is a GGUF file.
@@ -142,29 +187,29 @@ def determine_backend(model_path: pathlib.Path) -> Tuple[str, str]:
     """
     Determine the backend to use based on the model file properties.
     Args:
-        model_path (Path): The path to the model file.
+        model_path (Path): The path to the model file/directory.
     Returns:
         Tuple[str, str]: A tuple containing two strings:
                         - The backend to use.
                         - The reason why the backend was selected.
     """
-
-    if model_path.is_dir():
-        if sys.platform != "linux":
+    try:
+        if model_path.is_dir() and is_model_safetensors(model_path):
+            if sys.platform == "linux":
+                logger.debug(
+                    f"Model is huggingface safetensors and system is Linux, using {VLLM} backend."
+                )
+                return (
+                    VLLM,
+                    "model path is a directory containing huggingface safetensors files and running on Linux.",
+                )
             raise ValueError(
                 "Model is a directory containing huggingface safetensors files but the system is not Linux. "
                 "Using a directory with safetensors file will activate the vLLM serving backend, vLLM is only supported on Linux. "
                 "If you want to run the model on a different system (e.g. macOS), please use a GGUF file."
             )
-        # If the model is a directory, it's a VLLM model - it's kinda weak, but it's a start
-        if sys.platform == "linux":
-            logger.debug(
-                f"Model is a directory and system is Linux, using {VLLM} backend."
-            )
-            return (
-                VLLM,
-                "model path is a directory containing huggingface safetensors files and running on Linux.",
-            )
+    except Exception as e:
+        raise ValueError from e
 
     # Check if the model is a GGUF file
     try:


Test patch:
diff --git a/tests/test_backends.py b/tests/test_backends.py
index ea82fc20cb..f2b4c3b334 100644
--- a/tests/test_backends.py
+++ b/tests/test_backends.py
@@ -1,6 +1,8 @@
 # Standard
 from unittest import mock
 from unittest.mock import patch
+import json
+import os
 import pathlib
 import socket
 import sys
@@ -15,6 +17,27 @@
 from instructlab.model.backends.vllm import build_vllm_cmd
 
 
+# helper function to create dummy valid and invalid safetensor model directories
+def create_safetensors_model_files(safetensors_model_path: pathlib.Path, valid: bool):
+    test_json_dict = {"a": 1, "b": 2}
+    json_object = json.dumps(test_json_dict, indent=4)
+
+    for file in ["config.json", "tokenizer.json", "tokenizer_config.json"]:
+        os.makedirs(os.path.dirname(safetensors_model_path / file), exist_ok=True)
+        with open(safetensors_model_path / file, "a+", encoding="UTF-8") as f:
+            f.write(json_object)
+
+    with open(
+        safetensors_model_path / "model.safetensors", "a+", encoding="UTF-8"
+    ) as f:
+        f.write("")
+    if valid:
+        with open(
+            safetensors_model_path / "tokenizer.model", "a+", encoding="UTF-8"
+        ) as f:
+            f.write("")
+
+
 def test_free_port():
     host = "localhost"
     port = backends.free_tcp_ipv4_port(host)
@@ -49,33 +72,68 @@ def test_get_backend_auto_detection_success_gguf(
     m_is_model_gguf.assert_called_once_with(tmp_gguf)
 
 
-# this test fails because the model_path is a directory but the platform is not linux
-@patch("sys.platform", "darwin")
-def test_get_backend_auto_detection_failure_vllm_dir_but_not_linux(
-    tmp_path: pathlib.Path,
+# this tests both cases where a valid and invalid safetensors model directory is supplied
+@pytest.mark.parametrize(
+    "safetensors_dir,expected",
+    [
+        (
+            "valid_safetensors_model_dir",
+            True,
+        ),
+        (
+            "invalid_safetensors_model_dir",
+            False,
+        ),
+    ],
+)
+def test_is_model_safetensors_valid(
+    safetensors_dir: str, expected: bool, tmp_path: pathlib.Path
 ):
-    with pytest.raises(ValueError) as exc_info:
-        backends.get(tmp_path, None)
-    assert "Cannot determine which backend to use" in str(exc_info.value)
+    safetensors_model_path = tmp_path / safetensors_dir
+    create_safetensors_model_files(safetensors_model_path, expected)
 
+    val = backends.is_model_safetensors(safetensors_model_path)
+    assert val == expected
 
-# this test succeeds because the model_path is a directory and the platform is linux so vllm is
+
+# this test succeeds because the model_path is a valid safetensors directory and the platform is linux so vllm is
 # picked
 @patch("sys.platform", "linux")
-def test_get_backend_auto_detection_success_vllm_dir(tmp_path: pathlib.Path):
-    backend = backends.get(tmp_path, None)
-    assert backend == "vllm"
+@pytest.mark.parametrize(
+    "safetensors_dir,expected",
+    [
+        (
+            "valid_safetensors_model_dir",
+            "vllm",
+        ),
+    ],
+)
+def test_get_backend_auto_detection_success_vllm_dir(
+    safetensors_dir: str, expected: str, tmp_path: pathlib.Path
+):
+    safetensors_model_path = tmp_path / safetensors_dir
+    create_safetensors_model_files(safetensors_model_path, True)
+
+    backend = backends.get(safetensors_model_path, None)
+    assert backend == expected
 
 
 # this test fails because the OS is darwin and a directory is passed, only works on Linux
 @patch("sys.platform", "darwin")
-def test_get_backend_auto_detection_failed_vllm_dir_darwin(tmp_path: pathlib.Path):
+@pytest.mark.parametrize(
+    "safetensors_dir",
+    [
+        ("valid_safetensors_model_dir"),
+    ],
+)
+def test_get_backend_auto_detection_failed_vllm_dir_darwin(
+    safetensors_dir: str, tmp_path: pathlib.Path
+):
     with pytest.raises(ValueError) as exc_info:
-        backends.get(tmp_path, None)
-    assert (
-        "Model is a directory containing huggingface safetensors files but the system is not Linux."
-        in str(exc_info.value)
-    )
+        safetensors_model_path = tmp_path / safetensors_dir
+        create_safetensors_model_files(safetensors_model_path, True)
+        backends.get(safetensors_model_path, None)
+    assert "Cannot determine which backend to use" in str(exc_info.value)
 
 
 # this test succeeds even if the auto-detection picked a different backend, we continue with what