diff --git a/python/openai/openai_frontend/engine/utils/triton.py b/python/openai/openai_frontend/engine/utils/triton.py index 3b0f04e3d9..636e58435d 100644 --- a/python/openai/openai_frontend/engine/utils/triton.py +++ b/python/openai/openai_frontend/engine/utils/triton.py @@ -28,6 +28,7 @@ import os import re from dataclasses import asdict, dataclass, field +from pathlib import Path from typing import Iterable, List, Optional, Union import numpy as np @@ -356,13 +357,24 @@ def _get_guided_json_from_tool( def _get_vllm_lora_names( model_repository: str | list[str], model_name: str, model_version: int ) -> None | List[str]: + if ( + len(model_name) == 0 + or model_name.isspace() + or "/" in model_name + or "\\" in model_name + ): + raise ValueError( + f"Invalid model name: '{model_name}'. Model names must be valid file-system-path segment names." + ) lora_names = [] repo_paths = model_repository if isinstance(repo_paths, str): repo_paths = [repo_paths] for repo_path in repo_paths: model_path = os.path.join(repo_path, model_name) - if os.path.normpath(model_path) != model_path: + if (not Path(model_path).is_relative_to(repo_path)) or ( + os.path.normpath(model_path) != model_path + ): raise ValueError( f"Invalid model name: '{model_name}'. Model names must be valid file-system-path segment names." ) diff --git a/python/openai/tests/test_lora.py b/python/openai/tests/test_lora.py index 4a5ff10361..550273cf52 100644 --- a/python/openai/tests/test_lora.py +++ b/python/openai/tests/test_lora.py @@ -29,12 +29,50 @@ import shutil import unittest +import pytest from huggingface_hub import snapshot_download from openai import BadRequestError, NotFoundError +from openai_frontend.engine.utils.triton import ( + _get_vllm_lora_names as get_vllm_lora_names, +) from .utils import OpenAIServer +@pytest.mark.parametrize( + "model_repository,model_name,expect_error", + [ + ("openai_model_repository", "", True), # Empty string as model name. + ("openai_model_repository", " ", True), # Whitespace-only model name. + ("openai_model_repository", "invalid/path", True), + ("openai_model_repository", "invalid\\path", True), + ("openai_model_repository", "../outside/repo", True), + ("openai_model_repository", "../test_models/identity_py", True), + ("test_models", "../test_models/identity_py", True), + ("test_models", "identity_py", False), + ("test_models", "mock_llm", False), + ], +) +def test_get_vllm_lora_name(model_repository: str, model_name: str, expect_error: bool): + try: + get_vllm_lora_names(model_repository, model_name, 1) + except ValueError as e: + if expect_error: + assert ( + f"Invalid model name: '{model_name}'. Model names must be valid file-system-path segment names." + == str(e) + ) + else: + raise pytest.fail( + f"(model_repository='{model_repository}', model_name='{model_name}') raised ValueError unexpectedly: {e}" + ) + else: + if expect_error: + raise pytest.fail( + f"(model_repository='{model_repository}', model_name='{model_name}') did not raise ValueError as expected." + ) + + def is_vllm_installed(): try: import vllm as _