wehub-resource-sync
296 lines
8.8 KiB
Python
296 lines
8.8 KiB
Python
import base64
|
|
import io
|
|
import os
|
|
import tempfile
|
|
from typing import Generator, List
|
|
|
|
import PIL.Image
|
|
import pytest
|
|
import requests
|
|
|
|
from ray import serve
|
|
from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_llm_deployment
|
|
|
|
S3_ARTIFACT_URL = "https://air-example-data.s3.amazonaws.com/"
|
|
S3_ARTIFACT_LLM_OSSCI_URL = S3_ARTIFACT_URL + "rayllm-ossci/"
|
|
S3_ARTIFACT_ASSETS_URL = S3_ARTIFACT_LLM_OSSCI_URL + "assets/"
|
|
|
|
|
|
def download_model_from_s3(
|
|
remote_url: str, file_list: List[str]
|
|
) -> Generator[str, None, None]:
|
|
"""
|
|
Download the model checkpoint and tokenizer from S3 for testing
|
|
The reason to download the model from S3 is to avoid downloading the model
|
|
from HuggingFace hub during testing, which is flaky because of the rate
|
|
limit and HF hub downtime.
|
|
|
|
Args:
|
|
remote_url: The remote URL to download the model from.
|
|
file_list: The list of files to download.
|
|
|
|
Yields:
|
|
str: The path to the downloaded model checkpoint and tokenizer.
|
|
"""
|
|
with tempfile.TemporaryDirectory(prefix="ray-llm-test-model") as checkpoint_dir:
|
|
print(f"Downloading model from {remote_url} to {checkpoint_dir}", flush=True)
|
|
for file_name in file_list:
|
|
response = requests.get(remote_url + file_name)
|
|
with open(os.path.join(checkpoint_dir, file_name), "wb") as fp:
|
|
fp.write(response.content)
|
|
yield os.path.abspath(checkpoint_dir)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_opt_125m():
|
|
"""The small decoder model for testing."""
|
|
REMOTE_URL = f"{S3_ARTIFACT_URL}facebook-opt-125m/"
|
|
FILE_LIST = [
|
|
"config.json",
|
|
"flax_model.msgpack",
|
|
"generation_config.json",
|
|
"merges.txt",
|
|
"pytorch_model.bin",
|
|
"special_tokens_map.json",
|
|
"tokenizer_config.json",
|
|
"vocab.json",
|
|
]
|
|
yield from download_model_from_s3(REMOTE_URL, FILE_LIST)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_llava_354m():
|
|
"""The vision language model for testing."""
|
|
REMOTE_URL = f"{S3_ARTIFACT_URL}llava-354M/"
|
|
FILE_LIST = [
|
|
"added_tokens.json",
|
|
"chat_template.json",
|
|
"config.json",
|
|
"generation_config.json",
|
|
"model.safetensors",
|
|
"preprocessor_config.json",
|
|
"processor_config.json",
|
|
"special_tokens_map.json",
|
|
"tokenizer.json",
|
|
"tokenizer.model",
|
|
"tokenizer_config.json",
|
|
]
|
|
yield from download_model_from_s3(REMOTE_URL, FILE_LIST)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_smolvlm_256m():
|
|
"""The vision language model for testing."""
|
|
REMOTE_URL = f"{S3_ARTIFACT_LLM_OSSCI_URL}smolvlm-256m-instruct/"
|
|
FILE_LIST = [
|
|
"added_tokens.json",
|
|
"chat_template.json",
|
|
"config.json",
|
|
"generation_config.json",
|
|
"merges.txt",
|
|
"model.safetensors",
|
|
"preprocessor_config.json",
|
|
"processor_config.json",
|
|
"special_tokens_map.json",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
"vocab.json",
|
|
]
|
|
yield from download_model_from_s3(REMOTE_URL, FILE_LIST)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_llama_3_2_216M():
|
|
"""The llama 3.2 216M model for testing."""
|
|
REMOTE_URL = f"{S3_ARTIFACT_URL}llama-3.2-216M-dummy/"
|
|
FILE_LIST = [
|
|
"config.json",
|
|
"generation_config.json",
|
|
"special_tokens_map.json",
|
|
"tokenizer_config.json",
|
|
"tokenizer.json",
|
|
"model.safetensors",
|
|
]
|
|
yield from download_model_from_s3(REMOTE_URL, FILE_LIST)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_llama_3_2_216M_lora():
|
|
"""The LoRA model for testing."""
|
|
REMOTE_URL = f"{S3_ARTIFACT_URL}llama-3.2-216M-lora-dummy/"
|
|
FILE_LIST = [
|
|
"adapter_config.json",
|
|
"adapter_model.safetensors",
|
|
]
|
|
yield from download_model_from_s3(REMOTE_URL, FILE_LIST)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_pixtral_12b():
|
|
"""The Pixtral 12B model for testing."""
|
|
REMOTE_URL = f"{S3_ARTIFACT_URL}mistral-community-pixtral-12b/"
|
|
FILE_LIST = [
|
|
"config.json",
|
|
"chat_template.json",
|
|
"preprocessor_config.json",
|
|
"processor_config.json",
|
|
"special_tokens_map.json",
|
|
"tokenizer_config.json",
|
|
"tokenizer.json",
|
|
]
|
|
yield from download_model_from_s3(REMOTE_URL, FILE_LIST)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_llama_3_2_1B_instruct():
|
|
"""The llama 3.2 1B Instruct model for testing."""
|
|
REMOTE_URL = f"{S3_ARTIFACT_URL}unsloth-Llama-3.2-1B-Instruct/"
|
|
FILE_LIST = [
|
|
"config.json",
|
|
"generation_config.json",
|
|
"model.safetensors",
|
|
"special_tokens_map.json",
|
|
"tokenizer_config.json",
|
|
"tokenizer.json",
|
|
]
|
|
yield from download_model_from_s3(REMOTE_URL, FILE_LIST)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_internlm2_1_8b():
|
|
"""
|
|
Yields the S3 URI so that download_model_files exercises the cloud download
|
|
path where the TOKENIZER_ONLY vs. EXCLUDE_SAFETENSORS filtering applies.
|
|
"""
|
|
yield "s3://anonymous@air-example-data/rayllm-ossci/internlm2-1_8b/"
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_qwen_2_5_omni_3b():
|
|
REMOTE_URL = f"{S3_ARTIFACT_LLM_OSSCI_URL}Qwen2.5-Omni-3B/"
|
|
FILE_LIST = [
|
|
"added_tokens.json",
|
|
"config.json",
|
|
"chat_template.json",
|
|
"generation_config.json",
|
|
"merges.txt",
|
|
"model-00001-of-00003.safetensors",
|
|
"model-00002-of-00003.safetensors",
|
|
"model-00003-of-00003.safetensors",
|
|
"model.safetensors.index.json",
|
|
"preprocessor_config.json",
|
|
"special_tokens_map.json",
|
|
"spk_dict.json",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
"vocab.json",
|
|
]
|
|
yield from download_model_from_s3(REMOTE_URL, FILE_LIST)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def gpu_type():
|
|
"""Get the GPU type used for testing."""
|
|
|
|
try:
|
|
import torch
|
|
|
|
print(f"{torch.version.cuda=}", flush=True)
|
|
name = torch.cuda.get_device_name()
|
|
# The name of the GPU is in the format of "NVIDIA L4" or "Tesla T4"
|
|
# or "NVIDIA H100 80GB HBM3"
|
|
type_name = name.split(" ")[1]
|
|
print(f"GPU type: {type_name}", flush=True)
|
|
yield type_name
|
|
except ImportError:
|
|
print("Failed to import torch to get GPU type", flush=True)
|
|
except ValueError as err:
|
|
print(f"Failed to get the GPU type: {err}", flush=True)
|
|
|
|
|
|
@pytest.fixture
|
|
def serve_cleanup():
|
|
yield
|
|
serve.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def create_model_opt_125m_deployment(gpu_type, model_opt_125m, serve_cleanup):
|
|
"""Create a serve deployment for testing."""
|
|
app_name = "test_serve_deployment_processor_app"
|
|
deployment_name = "test_deployment_name"
|
|
|
|
chat_template = """
|
|
{% if messages[0]['role'] == 'system' %}
|
|
{% set offset = 1 %}
|
|
{% else %}
|
|
{% set offset = 0 %}
|
|
{% endif %}
|
|
|
|
{{ bos_token }}
|
|
{% for message in messages %}
|
|
{% if (message['role'] == 'user') != (loop.index0 % 2 == offset) %}
|
|
{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
|
|
{% endif %}
|
|
|
|
{{ '<|im_start|>' + message['role'] + '\n' + message['content'] | trim + '<|im_end|>\n' }}
|
|
{% endfor %}
|
|
|
|
{% if add_generation_prompt %}
|
|
{{ '<|im_start|>assistant\n' }}
|
|
{% endif %}
|
|
"""
|
|
|
|
# Create a vLLM serve deployment
|
|
llm_config = LLMConfig(
|
|
model_loading_config=ModelLoadingConfig(
|
|
model_id=model_opt_125m,
|
|
model_source=model_opt_125m,
|
|
),
|
|
accelerator_type=gpu_type,
|
|
deployment_config=dict(
|
|
name="test_deployment_name", # This is not necessarily the final deployment name
|
|
autoscaling_config=dict(
|
|
min_replicas=1,
|
|
max_replicas=1,
|
|
),
|
|
),
|
|
engine_kwargs=dict(
|
|
enable_prefix_caching=True,
|
|
enable_chunked_prefill=True,
|
|
max_num_batched_tokens=4096,
|
|
# Add chat template for OPT model to enable chat API
|
|
chat_template=chat_template,
|
|
),
|
|
)
|
|
|
|
llm_app = build_llm_deployment(
|
|
llm_config, override_serve_options=dict(name=deployment_name)
|
|
)
|
|
serve.run(llm_app, name=app_name)
|
|
yield deployment_name, app_name
|
|
|
|
|
|
@pytest.fixture
|
|
def image_asset():
|
|
image_url = S3_ARTIFACT_ASSETS_URL + "cherry_blossom.jpg"
|
|
with requests.get(image_url) as response:
|
|
response.raise_for_status()
|
|
image_pil = PIL.Image.open(io.BytesIO(response.content))
|
|
yield image_url, image_pil
|
|
|
|
|
|
@pytest.fixture
|
|
def audio_asset():
|
|
audio_url = S3_ARTIFACT_ASSETS_URL + "winning_call.ogg"
|
|
with requests.get(audio_url) as response:
|
|
response.raise_for_status()
|
|
audio_data = base64.b64encode(response.content).decode("utf-8")
|
|
yield audio_url, audio_data
|
|
|
|
|
|
@pytest.fixture
|
|
def video_asset():
|
|
video_url = S3_ARTIFACT_ASSETS_URL + "free-videos.mp4"
|
|
yield video_url
|