wehub-resource-sync
113 lines
4.9 KiB
Python
113 lines
4.9 KiB
Python
"""Tests for CloudFileSystem class."""
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from ray.llm._internal.common.utils.cloud_utils import CloudFileSystem
|
|
|
|
|
|
class TestCloudFileSystem:
|
|
"""Tests for the CloudFileSystem class."""
|
|
|
|
@patch("ray.llm._internal.common.utils.cloud_utils.GCSFileSystem")
|
|
def test_download_model(self, mock_gcs_filesystem):
|
|
"""Test downloading a model from cloud storage."""
|
|
# Mock GCSFileSystem.get_file to return hash content
|
|
mock_gcs_filesystem.get_file.return_value = "abcdef1234567890"
|
|
|
|
# Create temp directory for testing
|
|
with tempfile.TemporaryDirectory() as tempdir:
|
|
# Test downloading model
|
|
with patch.object(CloudFileSystem, "download_files") as mock_download:
|
|
CloudFileSystem.download_model(tempdir, "gs://bucket/model", False)
|
|
|
|
# Check that hash file was processed
|
|
assert os.path.exists(os.path.join(tempdir, "refs", "main"))
|
|
with open(os.path.join(tempdir, "refs", "main"), "r") as f:
|
|
assert f.read() == "abcdef1234567890"
|
|
|
|
# Verify get_file was called for hash file
|
|
mock_gcs_filesystem.get_file.assert_called_once_with(
|
|
"gs://bucket/model/hash", decode_as_utf_8=True
|
|
)
|
|
|
|
# Check that download_files was called correctly
|
|
mock_download.assert_called_once()
|
|
call_args = mock_download.call_args[1]
|
|
assert call_args["path"] == os.path.join(
|
|
tempdir, "snapshots", "abcdef1234567890"
|
|
)
|
|
assert call_args["bucket_uri"] == "gs://bucket/model"
|
|
assert call_args["substrings_to_include"] == []
|
|
assert call_args["suffixes_to_exclude"] is None
|
|
|
|
@patch("ray.llm._internal.common.utils.cloud_utils.GCSFileSystem")
|
|
def test_download_model_tokenizer_only_includes_chat_template_jinja(
|
|
self, mock_gcs_filesystem
|
|
):
|
|
"""Regression test for chat_template.jinja not downloaded in tokenizer-only mode.
|
|
|
|
Models like Qwen3.5 use an external chat_template.jinja file instead of the
|
|
legacy chat_template field inside tokenizer_config.json. Without "chat_template"
|
|
in the substring filter, the .jinja file is skipped and ChatTemplateStage fails
|
|
with: ValueError: Cannot use apply_chat_template because this processor does not
|
|
have a chat template.
|
|
"""
|
|
mock_gcs_filesystem.get_file.return_value = "abcdef1234567890"
|
|
|
|
with tempfile.TemporaryDirectory() as tempdir:
|
|
with patch.object(CloudFileSystem, "download_files") as mock_download:
|
|
CloudFileSystem.download_model(tempdir, "gs://bucket/model", True)
|
|
|
|
mock_download.assert_called_once()
|
|
substrings = mock_download.call_args[1]["substrings_to_include"]
|
|
# "chat_template" matches chat_template.jinja, which is the HuggingFace
|
|
# convention used by newer models (e.g. Qwen3.5) for external chat templates.
|
|
assert (
|
|
"chat_template" in substrings
|
|
), "chat_template.jinja files must be downloaded in tokenizer-only mode"
|
|
|
|
@patch("ray.llm._internal.common.utils.cloud_utils.GCSFileSystem")
|
|
def test_upload_model(self, mock_gcs_filesystem):
|
|
"""Test uploading a model to cloud storage."""
|
|
# Create temp directory for testing
|
|
with tempfile.TemporaryDirectory() as tempdir:
|
|
hash = "abcdef1234567890"
|
|
# Create refs/main file
|
|
os.makedirs(os.path.join(tempdir, "refs"), exist_ok=True)
|
|
model_rev_path = os.path.join(tempdir, "refs", "main")
|
|
with open(model_rev_path, "w") as f:
|
|
f.write(hash)
|
|
|
|
# Create snapshots/<hash> folder
|
|
model_asset_path = os.path.join(tempdir, "snapshots", hash)
|
|
os.makedirs(model_asset_path)
|
|
|
|
# Test uploading model
|
|
CloudFileSystem.upload_model(tempdir, "gs://bucket/model")
|
|
|
|
# Check that upload_files was called twice - once for model assets and once for hash file
|
|
assert mock_gcs_filesystem.upload_files.call_count == 2
|
|
|
|
# Verify the calls were made with correct arguments
|
|
calls = mock_gcs_filesystem.upload_files.call_args_list
|
|
call_paths = {
|
|
call[0][0] for call in calls
|
|
} # Extract local_path from each call
|
|
call_uris = {
|
|
call[0][1] for call in calls
|
|
} # Extract bucket_uri from each call
|
|
|
|
assert model_asset_path in call_paths
|
|
assert model_rev_path in call_paths
|
|
assert "gs://bucket/model" in call_uris
|
|
assert "gs://bucket/model/hash" in call_uris
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|