chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,112 @@
|
||||
"""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__]))
|
||||
@@ -0,0 +1,164 @@
|
||||
"""Tests for mirror config classes."""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.llm._internal.common.utils.cloud_utils import (
|
||||
CloudMirrorConfig,
|
||||
LoraMirrorConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestCloudMirrorConfig:
|
||||
"""Tests for the CloudMirrorConfig class."""
|
||||
|
||||
def test_valid_s3_uri(self):
|
||||
"""Test valid S3 URI."""
|
||||
config = CloudMirrorConfig(bucket_uri="s3://my-bucket/path")
|
||||
assert config.bucket_uri == "s3://my-bucket/path"
|
||||
assert config.storage_type == "s3"
|
||||
|
||||
def test_valid_gcs_uri(self):
|
||||
"""Test valid GCS URI."""
|
||||
config = CloudMirrorConfig(bucket_uri="gs://my-bucket/path")
|
||||
assert config.bucket_uri == "gs://my-bucket/path"
|
||||
assert config.storage_type == "gcs"
|
||||
|
||||
def test_valid_abfss_uri(self):
|
||||
"""Test valid ABFSS URI."""
|
||||
config = CloudMirrorConfig(
|
||||
bucket_uri="abfss://container@account.dfs.core.windows.net/path"
|
||||
)
|
||||
assert (
|
||||
config.bucket_uri == "abfss://container@account.dfs.core.windows.net/path"
|
||||
)
|
||||
assert config.storage_type == "abfss"
|
||||
|
||||
def test_valid_azure_uri(self):
|
||||
"""Test valid Azure URI."""
|
||||
config = CloudMirrorConfig(
|
||||
bucket_uri="azure://container@account.blob.core.windows.net/path"
|
||||
)
|
||||
assert (
|
||||
config.bucket_uri == "azure://container@account.blob.core.windows.net/path"
|
||||
)
|
||||
assert config.storage_type == "azure"
|
||||
|
||||
def test_none_uri(self):
|
||||
"""Test None URI."""
|
||||
config = CloudMirrorConfig(bucket_uri=None)
|
||||
assert config.bucket_uri is None
|
||||
assert config.storage_type is None
|
||||
|
||||
def test_invalid_uri(self):
|
||||
"""Test invalid URI."""
|
||||
with pytest.raises(
|
||||
ValueError, match='Got invalid value "file:///tmp" for bucket_uri'
|
||||
):
|
||||
CloudMirrorConfig(bucket_uri="file:///tmp")
|
||||
|
||||
def test_extra_files(self):
|
||||
"""Test extra files configuration."""
|
||||
config = CloudMirrorConfig(
|
||||
bucket_uri="s3://bucket/path",
|
||||
extra_files=[
|
||||
{"bucket_uri": "s3://bucket/file1", "destination_path": "/dest1"},
|
||||
{"bucket_uri": "s3://bucket/file2", "destination_path": "/dest2"},
|
||||
],
|
||||
)
|
||||
assert len(config.extra_files) == 2
|
||||
assert config.extra_files[0].bucket_uri == "s3://bucket/file1"
|
||||
assert config.extra_files[0].destination_path == "/dest1"
|
||||
|
||||
|
||||
class TestLoraMirrorConfig:
|
||||
"""Tests for the LoraMirrorConfig class."""
|
||||
|
||||
def test_valid_s3_config(self):
|
||||
"""Test valid S3 LoRA config."""
|
||||
config = LoraMirrorConfig(
|
||||
lora_model_id="test-model",
|
||||
bucket_uri="s3://my-bucket/lora-models",
|
||||
max_total_tokens=1000,
|
||||
)
|
||||
assert config.lora_model_id == "test-model"
|
||||
assert config.bucket_uri == "s3://my-bucket/lora-models"
|
||||
assert config.bucket_name == "my-bucket"
|
||||
assert config.bucket_path == "lora-models"
|
||||
|
||||
def test_valid_abfss_config(self):
|
||||
"""Test valid ABFSS LoRA config."""
|
||||
config = LoraMirrorConfig(
|
||||
lora_model_id="test-model",
|
||||
bucket_uri="abfss://container@account.dfs.core.windows.net/lora/models",
|
||||
max_total_tokens=1000,
|
||||
)
|
||||
assert config.lora_model_id == "test-model"
|
||||
assert (
|
||||
config.bucket_uri
|
||||
== "abfss://container@account.dfs.core.windows.net/lora/models"
|
||||
)
|
||||
assert config.bucket_name == "container"
|
||||
assert config.bucket_path == "lora/models"
|
||||
|
||||
def test_valid_azure_config(self):
|
||||
"""Test valid Azure LoRA config."""
|
||||
config = LoraMirrorConfig(
|
||||
lora_model_id="test-model",
|
||||
bucket_uri="azure://container@account.blob.core.windows.net/lora/models",
|
||||
max_total_tokens=1000,
|
||||
)
|
||||
assert config.lora_model_id == "test-model"
|
||||
assert (
|
||||
config.bucket_uri
|
||||
== "azure://container@account.blob.core.windows.net/lora/models"
|
||||
)
|
||||
assert config.bucket_name == "container"
|
||||
assert config.bucket_path == "lora/models"
|
||||
|
||||
def test_bucket_path_parsing(self):
|
||||
"""Test bucket path parsing for different URI formats."""
|
||||
# S3 with multiple path segments
|
||||
config = LoraMirrorConfig(
|
||||
lora_model_id="test",
|
||||
bucket_uri="s3://bucket/path/to/model",
|
||||
max_total_tokens=1000,
|
||||
)
|
||||
assert config.bucket_name == "bucket"
|
||||
assert config.bucket_path == "path/to/model"
|
||||
|
||||
# ABFSS with multiple path segments
|
||||
config = LoraMirrorConfig(
|
||||
lora_model_id="test",
|
||||
bucket_uri="abfss://container@account.dfs.core.windows.net/deep/nested/path",
|
||||
max_total_tokens=1000,
|
||||
)
|
||||
assert config.bucket_name == "container"
|
||||
assert config.bucket_path == "deep/nested/path"
|
||||
|
||||
def test_invalid_uri(self):
|
||||
"""Test invalid URI in LoRA config."""
|
||||
with pytest.raises(
|
||||
ValueError, match='Got invalid value "file:///tmp" for bucket_uri'
|
||||
):
|
||||
LoraMirrorConfig(
|
||||
lora_model_id="test-model",
|
||||
bucket_uri="file:///tmp",
|
||||
max_total_tokens=1000,
|
||||
)
|
||||
|
||||
def test_optional_fields(self):
|
||||
"""Test optional fields in LoRA config."""
|
||||
config = LoraMirrorConfig(
|
||||
lora_model_id="test-model",
|
||||
bucket_uri="s3://bucket/path",
|
||||
max_total_tokens=1000,
|
||||
sync_args=["--exclude", "*.tmp"],
|
||||
)
|
||||
assert config.max_total_tokens == 1000
|
||||
assert config.sync_args == ["--exclude", "*.tmp"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,600 @@
|
||||
"""Tests for PyArrowFileSystem class."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
import pyarrow.fs as pa_fs
|
||||
import pytest
|
||||
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.pyarrow_filesystem import (
|
||||
PyArrowFileSystem,
|
||||
)
|
||||
|
||||
|
||||
class TestPyArrowFileSystem:
|
||||
"""Tests for the PyArrowFileSystem class."""
|
||||
|
||||
@patch("pyarrow.fs.S3FileSystem")
|
||||
def test_get_file(self, mock_s3fs):
|
||||
"""Test getting a file from cloud storage."""
|
||||
# Setup mock filesystem and file
|
||||
mock_fs = MagicMock()
|
||||
mock_s3fs.return_value = mock_fs
|
||||
|
||||
# Mock file content and info
|
||||
mock_file = MagicMock()
|
||||
mock_file.read.return_value = b"test file content"
|
||||
mock_fs.open_input_file.return_value.__enter__.return_value = mock_file
|
||||
mock_fs.get_file_info.return_value.type = pa_fs.FileType.File
|
||||
|
||||
# Test getting file as string (default)
|
||||
content = PyArrowFileSystem.get_file("s3://bucket/test.txt")
|
||||
assert content == "test file content"
|
||||
|
||||
# Test getting file as bytes
|
||||
content_bytes = PyArrowFileSystem.get_file(
|
||||
"s3://bucket/test.txt", decode_as_utf_8=False
|
||||
)
|
||||
assert content_bytes == b"test file content"
|
||||
|
||||
# Test non-existent file
|
||||
mock_fs.get_file_info.return_value.type = pa_fs.FileType.NotFound
|
||||
assert PyArrowFileSystem.get_file("s3://bucket/nonexistent.txt") is None
|
||||
|
||||
@patch("pyarrow.fs.GcsFileSystem")
|
||||
def test_list_subfolders(self, mock_gcsfs):
|
||||
"""Test listing subfolders in cloud storage."""
|
||||
# Setup mock filesystem
|
||||
mock_fs = MagicMock()
|
||||
mock_gcsfs.return_value = mock_fs
|
||||
|
||||
# Create mock file infos for directory listing
|
||||
dir1 = MagicMock()
|
||||
dir1.type = pa_fs.FileType.Directory
|
||||
dir1.path = "bucket/parent/dir1"
|
||||
|
||||
dir2 = MagicMock()
|
||||
dir2.type = pa_fs.FileType.Directory
|
||||
dir2.path = "bucket/parent/dir2"
|
||||
|
||||
file1 = MagicMock()
|
||||
file1.type = pa_fs.FileType.File
|
||||
file1.path = "bucket/parent/file.txt"
|
||||
|
||||
mock_fs.get_file_info.return_value = [dir1, dir2, file1]
|
||||
|
||||
# Test listing subfolders
|
||||
folders = PyArrowFileSystem.list_subfolders("gs://bucket/parent")
|
||||
assert sorted(folders) == ["dir1", "dir2"]
|
||||
|
||||
@patch.object(PyArrowFileSystem, "get_fs_and_path")
|
||||
def test_list_subfolders_exception_handling(self, mock_get_fs_and_path):
|
||||
"""Test that list_subfolders returns empty list when get_fs_and_path raises exception."""
|
||||
# Make get_fs_and_path raise an exception
|
||||
mock_get_fs_and_path.side_effect = ValueError("Example exception")
|
||||
|
||||
# Test that list_subfolders handles the exception gracefully
|
||||
folders = PyArrowFileSystem.list_subfolders("gs://bucket/parent")
|
||||
assert folders == []
|
||||
|
||||
# Verify get_fs_and_path was called
|
||||
mock_get_fs_and_path.assert_called_once_with("gs://bucket/parent/")
|
||||
|
||||
@patch("pyarrow.fs.copy_files")
|
||||
@patch("pyarrow.fs.S3FileSystem")
|
||||
def test_download_files_no_filters(self, mock_s3fs, mock_copy_files):
|
||||
"""Test downloading files from cloud storage without filters."""
|
||||
|
||||
# Setup mock filesystem
|
||||
mock_fs = MagicMock()
|
||||
mock_s3fs.return_value = mock_fs
|
||||
|
||||
# Create temp directory for testing
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Test downloading files without filters
|
||||
PyArrowFileSystem.download_files(tempdir, "s3://bucket/dir")
|
||||
|
||||
# Verify copy_files was called with correct arguments
|
||||
mock_copy_files.assert_called_once_with(
|
||||
source="bucket/dir",
|
||||
destination=tempdir,
|
||||
source_filesystem=mock_fs,
|
||||
destination_filesystem=ANY,
|
||||
use_threads=True,
|
||||
chunk_size=64 * 1024 * 1024,
|
||||
)
|
||||
|
||||
@patch("pyarrow.fs.copy_files")
|
||||
@patch("pyarrow.fs.S3FileSystem")
|
||||
def test_download_files_with_filters(self, mock_s3fs, mock_copy_files):
|
||||
"""Test downloading files from cloud storage with filters."""
|
||||
|
||||
# Setup mock filesystem
|
||||
mock_fs = MagicMock()
|
||||
mock_s3fs.return_value = mock_fs
|
||||
|
||||
# Create mock file infos for listing
|
||||
file_info1 = MagicMock()
|
||||
file_info1.type = pa_fs.FileType.File
|
||||
file_info1.path = "bucket/dir/file1.txt"
|
||||
|
||||
file_info2 = MagicMock()
|
||||
file_info2.type = pa_fs.FileType.File
|
||||
file_info2.path = "bucket/dir/subdir/file2.json"
|
||||
|
||||
file_info3 = MagicMock()
|
||||
file_info3.type = pa_fs.FileType.File
|
||||
file_info3.path = "bucket/dir/file3.tmp"
|
||||
|
||||
dir_info = MagicMock()
|
||||
dir_info.type = pa_fs.FileType.Directory
|
||||
dir_info.path = "bucket/dir/subdir"
|
||||
|
||||
mock_fs.get_file_info.return_value = [
|
||||
file_info1,
|
||||
file_info2,
|
||||
file_info3,
|
||||
dir_info,
|
||||
]
|
||||
|
||||
# Create temp directory for testing
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Test downloading files with filters
|
||||
PyArrowFileSystem.download_files(
|
||||
tempdir,
|
||||
"s3://bucket/dir",
|
||||
substrings_to_include=["file1", "file2"],
|
||||
suffixes_to_exclude=[".tmp"],
|
||||
)
|
||||
|
||||
# Verify copy_files was called for each filtered file
|
||||
assert mock_copy_files.call_count == 2
|
||||
|
||||
# Get all calls to copy_files
|
||||
calls = mock_copy_files.call_args_list
|
||||
|
||||
# Verify the calls (order may vary due to threading)
|
||||
expected_sources = {"bucket/dir/file1.txt", "bucket/dir/subdir/file2.json"}
|
||||
expected_dests = {
|
||||
os.path.join(tempdir, "file1.txt"),
|
||||
os.path.join(tempdir, "subdir", "file2.json"),
|
||||
}
|
||||
|
||||
actual_sources = {call.kwargs["source"] for call in calls}
|
||||
actual_dests = {call.kwargs["destination"] for call in calls}
|
||||
|
||||
assert actual_sources == expected_sources
|
||||
assert actual_dests == expected_dests
|
||||
|
||||
# Verify all calls have correct filesystem and options
|
||||
for call in calls:
|
||||
assert call.kwargs["source_filesystem"] == mock_fs
|
||||
assert call.kwargs["destination_filesystem"] is not None
|
||||
assert call.kwargs["use_threads"] is True
|
||||
assert call.kwargs["chunk_size"] == 64 * 1024 * 1024
|
||||
|
||||
@patch("pyarrow.fs.copy_files")
|
||||
@patch("pyarrow.fs.S3FileSystem")
|
||||
def test_upload_files(self, mock_s3fs, mock_copy_files):
|
||||
"""Test uploading files to cloud storage."""
|
||||
|
||||
# Setup mock filesystem
|
||||
mock_fs = MagicMock()
|
||||
mock_s3fs.return_value = mock_fs
|
||||
|
||||
# Create temp directory for testing
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Test uploading files
|
||||
PyArrowFileSystem.upload_files(tempdir, "s3://bucket/dir")
|
||||
|
||||
# Check that the files are copied
|
||||
mock_copy_files.assert_called_once_with(
|
||||
source=tempdir,
|
||||
destination="bucket/dir",
|
||||
source_filesystem=ANY,
|
||||
destination_filesystem=ANY,
|
||||
)
|
||||
|
||||
|
||||
class TestFilterFiles:
|
||||
"""Tests for the _filter_files method in PyArrowFileSystem."""
|
||||
|
||||
def test_filter_files_no_filters(self):
|
||||
"""Test filtering files with no inclusion or exclusion filters."""
|
||||
# Setup mock filesystem
|
||||
mock_fs = MagicMock()
|
||||
|
||||
# Create mock file infos
|
||||
file_info1 = MagicMock()
|
||||
file_info1.type = pa_fs.FileType.File
|
||||
file_info1.path = "bucket/model/file1.txt"
|
||||
|
||||
file_info2 = MagicMock()
|
||||
file_info2.type = pa_fs.FileType.File
|
||||
file_info2.path = "bucket/model/subdir/file2.json"
|
||||
|
||||
dir_info = MagicMock()
|
||||
dir_info.type = pa_fs.FileType.Directory
|
||||
dir_info.path = "bucket/model/subdir"
|
||||
|
||||
mock_fs.get_file_info.return_value = [file_info1, file_info2, dir_info]
|
||||
|
||||
# Test filtering with no filters
|
||||
result = PyArrowFileSystem._filter_files(
|
||||
fs=mock_fs, source_path="bucket/model", destination_path="/local/dest"
|
||||
)
|
||||
|
||||
# Should include all files, exclude directories
|
||||
expected = [
|
||||
("bucket/model/file1.txt", "/local/dest/file1.txt"),
|
||||
("bucket/model/subdir/file2.json", "/local/dest/subdir/file2.json"),
|
||||
]
|
||||
assert sorted(result) == sorted(expected)
|
||||
|
||||
# Verify filesystem was called correctly
|
||||
mock_fs.get_file_info.assert_called_once()
|
||||
call_args = mock_fs.get_file_info.call_args[0][0]
|
||||
assert call_args.base_dir == "bucket/model"
|
||||
assert call_args.recursive is True
|
||||
|
||||
def test_filter_files_with_inclusion_substrings(self):
|
||||
"""Test filtering files with inclusion substrings."""
|
||||
# Setup mock filesystem
|
||||
mock_fs = MagicMock()
|
||||
|
||||
# Create mock file infos
|
||||
file_info1 = MagicMock()
|
||||
file_info1.type = pa_fs.FileType.File
|
||||
file_info1.path = "bucket/model/config.json"
|
||||
|
||||
file_info2 = MagicMock()
|
||||
file_info2.type = pa_fs.FileType.File
|
||||
file_info2.path = "bucket/model/weights.bin"
|
||||
|
||||
file_info3 = MagicMock()
|
||||
file_info3.type = pa_fs.FileType.File
|
||||
file_info3.path = "bucket/model/tokenizer.json"
|
||||
|
||||
mock_fs.get_file_info.return_value = [file_info1, file_info2, file_info3]
|
||||
|
||||
# Test filtering with inclusion substrings
|
||||
result = PyArrowFileSystem._filter_files(
|
||||
fs=mock_fs,
|
||||
source_path="bucket/model",
|
||||
destination_path="/local/dest",
|
||||
substrings_to_include=["config", "tokenizer"],
|
||||
)
|
||||
|
||||
# Should only include files with "config" or "tokenizer" in path
|
||||
expected = [
|
||||
("bucket/model/config.json", "/local/dest/config.json"),
|
||||
("bucket/model/tokenizer.json", "/local/dest/tokenizer.json"),
|
||||
]
|
||||
assert sorted(result) == sorted(expected)
|
||||
|
||||
def test_filter_files_with_exclusion_suffixes(self):
|
||||
"""Test filtering files with exclusion suffixes."""
|
||||
# Setup mock filesystem
|
||||
mock_fs = MagicMock()
|
||||
|
||||
# Create mock file infos
|
||||
file_info1 = MagicMock()
|
||||
file_info1.type = pa_fs.FileType.File
|
||||
file_info1.path = "bucket/model/model.bin"
|
||||
|
||||
file_info2 = MagicMock()
|
||||
file_info2.type = pa_fs.FileType.File
|
||||
file_info2.path = "bucket/model/config.json"
|
||||
|
||||
file_info3 = MagicMock()
|
||||
file_info3.type = pa_fs.FileType.File
|
||||
file_info3.path = "bucket/model/temp.tmp"
|
||||
|
||||
file_info4 = MagicMock()
|
||||
file_info4.type = pa_fs.FileType.File
|
||||
file_info4.path = "bucket/model/log.txt"
|
||||
|
||||
mock_fs.get_file_info.return_value = [
|
||||
file_info1,
|
||||
file_info2,
|
||||
file_info3,
|
||||
file_info4,
|
||||
]
|
||||
|
||||
# Test filtering with exclusion suffixes
|
||||
result = PyArrowFileSystem._filter_files(
|
||||
fs=mock_fs,
|
||||
source_path="bucket/model",
|
||||
destination_path="/local/dest",
|
||||
suffixes_to_exclude=[".tmp", ".txt"],
|
||||
)
|
||||
|
||||
# Should exclude files ending with .tmp or .txt
|
||||
expected = [
|
||||
("bucket/model/model.bin", "/local/dest/model.bin"),
|
||||
("bucket/model/config.json", "/local/dest/config.json"),
|
||||
]
|
||||
assert sorted(result) == sorted(expected)
|
||||
|
||||
def test_filter_files_with_both_filters(self):
|
||||
"""Test filtering files with both inclusion and exclusion filters."""
|
||||
# Setup mock filesystem
|
||||
mock_fs = MagicMock()
|
||||
|
||||
# Create mock file infos
|
||||
file_info1 = MagicMock()
|
||||
file_info1.type = pa_fs.FileType.File
|
||||
file_info1.path = "bucket/model/config.json"
|
||||
|
||||
file_info2 = MagicMock()
|
||||
file_info2.type = pa_fs.FileType.File
|
||||
file_info2.path = "bucket/model/config.tmp"
|
||||
|
||||
file_info3 = MagicMock()
|
||||
file_info3.type = pa_fs.FileType.File
|
||||
file_info3.path = "bucket/model/weights.bin"
|
||||
|
||||
file_info4 = MagicMock()
|
||||
file_info4.type = pa_fs.FileType.File
|
||||
file_info4.path = "bucket/model/tokenizer.json"
|
||||
|
||||
mock_fs.get_file_info.return_value = [
|
||||
file_info1,
|
||||
file_info2,
|
||||
file_info3,
|
||||
file_info4,
|
||||
]
|
||||
|
||||
# Test filtering with both inclusion and exclusion
|
||||
result = PyArrowFileSystem._filter_files(
|
||||
fs=mock_fs,
|
||||
source_path="bucket/model",
|
||||
destination_path="/local/dest",
|
||||
substrings_to_include=["config", "tokenizer"],
|
||||
suffixes_to_exclude=[".tmp"],
|
||||
)
|
||||
|
||||
# Should include files with "config" or "tokenizer" but exclude .tmp files
|
||||
expected = [
|
||||
("bucket/model/config.json", "/local/dest/config.json"),
|
||||
("bucket/model/tokenizer.json", "/local/dest/tokenizer.json"),
|
||||
]
|
||||
assert sorted(result) == sorted(expected)
|
||||
|
||||
|
||||
class TestPyArrowFileSystemAzureSupport:
|
||||
"""Tests for Azure/ABFSS support in PyArrowFileSystem."""
|
||||
|
||||
@patch("adlfs.AzureBlobFileSystem")
|
||||
@patch("azure.identity.DefaultAzureCredential")
|
||||
@patch("pyarrow.fs.PyFileSystem")
|
||||
@patch("pyarrow.fs.FSSpecHandler")
|
||||
def test_get_fs_and_path_abfss(
|
||||
self, mock_handler, mock_pyfs, mock_cred, mock_adlfs
|
||||
):
|
||||
"""Test getting ABFSS filesystem and path."""
|
||||
mock_adlfs_instance = MagicMock()
|
||||
mock_adlfs.return_value = mock_adlfs_instance
|
||||
mock_pyfs_instance = MagicMock()
|
||||
mock_pyfs.return_value = mock_pyfs_instance
|
||||
|
||||
fs, path = PyArrowFileSystem.get_fs_and_path(
|
||||
"abfss://container@account.dfs.core.windows.net/path/to/file"
|
||||
)
|
||||
|
||||
assert fs == mock_pyfs_instance
|
||||
assert path == "container/path/to/file"
|
||||
|
||||
# Verify the adlfs filesystem was created with correct parameters
|
||||
mock_adlfs.assert_called_once_with(
|
||||
account_name="account", credential=mock_cred.return_value
|
||||
)
|
||||
mock_handler.assert_called_once_with(mock_adlfs_instance)
|
||||
mock_pyfs.assert_called_once_with(mock_handler.return_value)
|
||||
|
||||
@patch("adlfs.AzureBlobFileSystem")
|
||||
@patch("azure.identity.DefaultAzureCredential")
|
||||
@patch("pyarrow.fs.PyFileSystem")
|
||||
@patch("pyarrow.fs.FSSpecHandler")
|
||||
def test_get_fs_and_path_azure(
|
||||
self, mock_handler, mock_pyfs, mock_cred, mock_adlfs
|
||||
):
|
||||
"""Test getting Azure filesystem and path."""
|
||||
mock_adlfs_instance = MagicMock()
|
||||
mock_adlfs.return_value = mock_adlfs_instance
|
||||
mock_pyfs_instance = MagicMock()
|
||||
mock_pyfs.return_value = mock_pyfs_instance
|
||||
|
||||
fs, path = PyArrowFileSystem.get_fs_and_path(
|
||||
"azure://container@account.blob.core.windows.net/path/to/file"
|
||||
)
|
||||
|
||||
assert fs == mock_pyfs_instance
|
||||
assert path == "container/path/to/file"
|
||||
|
||||
# Verify the adlfs filesystem was created with correct parameters
|
||||
mock_adlfs.assert_called_once_with(
|
||||
account_name="account", credential=mock_cred.return_value
|
||||
)
|
||||
|
||||
def test_abfss_uri_validation(self):
|
||||
"""Test ABFSS URI validation."""
|
||||
# Test valid URIs
|
||||
valid_uris = [
|
||||
"abfss://container@account.dfs.core.windows.net/path",
|
||||
"abfss://my-container@myaccount.dfs.core.windows.net/deep/nested/path",
|
||||
]
|
||||
|
||||
for uri in valid_uris:
|
||||
with patch("adlfs.AzureBlobFileSystem"), patch(
|
||||
"azure.identity.DefaultAzureCredential"
|
||||
), patch("pyarrow.fs.PyFileSystem"), patch("pyarrow.fs.FSSpecHandler"):
|
||||
# Should not raise an exception
|
||||
PyArrowFileSystem._create_abfss_filesystem(uri)
|
||||
|
||||
# Test invalid URIs
|
||||
invalid_uris = [
|
||||
"abfss://container", # Missing @account
|
||||
"abfss://@account.dfs.core.windows.net/path", # Empty container
|
||||
"abfss://container@account.wrong.domain/path", # Wrong domain
|
||||
"abfss://container@.dfs.core.windows.net/path", # Empty account
|
||||
"abfss://container@account.dfs.core.windows.net", # No path (but this is actually valid)
|
||||
]
|
||||
|
||||
for uri in invalid_uris[:-1]: # Skip the last one as it's actually valid
|
||||
with pytest.raises(ValueError):
|
||||
PyArrowFileSystem._create_abfss_filesystem(uri)
|
||||
|
||||
def test_azure_uri_validation(self):
|
||||
"""Test Azure URI validation."""
|
||||
# Test valid URIs
|
||||
valid_uris = [
|
||||
"azure://container@account.blob.core.windows.net/path",
|
||||
"azure://container@account.dfs.core.windows.net/path",
|
||||
"azure://my-container@myaccount.blob.core.windows.net/deep/nested/path",
|
||||
]
|
||||
|
||||
for uri in valid_uris:
|
||||
with patch("adlfs.AzureBlobFileSystem"), patch(
|
||||
"azure.identity.DefaultAzureCredential"
|
||||
), patch("pyarrow.fs.PyFileSystem"), patch("pyarrow.fs.FSSpecHandler"):
|
||||
# Should not raise an exception
|
||||
PyArrowFileSystem._create_azure_filesystem(uri)
|
||||
|
||||
# Test invalid URIs
|
||||
invalid_uris = [
|
||||
"azure://container", # Missing @account
|
||||
"azure://@account.blob.core.windows.net/path", # Empty container
|
||||
"azure://container@account.wrong.domain/path", # Wrong domain
|
||||
"azure://container@.blob.core.windows.net/path", # Empty account
|
||||
]
|
||||
|
||||
for uri in invalid_uris:
|
||||
with pytest.raises(ValueError):
|
||||
PyArrowFileSystem._create_azure_filesystem(uri)
|
||||
|
||||
def test_abfss_import_error(self):
|
||||
"""Test ImportError when adlfs is not available."""
|
||||
with patch(
|
||||
"builtins.__import__", side_effect=ImportError("No module named 'adlfs'")
|
||||
):
|
||||
with pytest.raises(
|
||||
ImportError, match="You must `pip install adlfs azure-identity`"
|
||||
):
|
||||
PyArrowFileSystem._create_abfss_filesystem(
|
||||
"abfss://container@account.dfs.core.windows.net/path"
|
||||
)
|
||||
|
||||
def test_azure_import_error(self):
|
||||
"""Test ImportError when adlfs is not available for Azure."""
|
||||
with patch(
|
||||
"builtins.__import__", side_effect=ImportError("No module named 'adlfs'")
|
||||
):
|
||||
with pytest.raises(
|
||||
ImportError, match="You must `pip install adlfs azure-identity`"
|
||||
):
|
||||
PyArrowFileSystem._create_azure_filesystem(
|
||||
"azure://container@account.blob.core.windows.net/path"
|
||||
)
|
||||
|
||||
@patch("adlfs.AzureBlobFileSystem")
|
||||
@patch("azure.identity.DefaultAzureCredential")
|
||||
@patch("pyarrow.fs.PyFileSystem")
|
||||
@patch("pyarrow.fs.FSSpecHandler")
|
||||
def test_abfss_anonymous_access_ignored(
|
||||
self, mock_handler, mock_pyfs, mock_cred, mock_adlfs
|
||||
):
|
||||
"""Test that anonymous access pattern is ignored for ABFSS URIs."""
|
||||
mock_adlfs_instance = MagicMock()
|
||||
mock_adlfs.return_value = mock_adlfs_instance
|
||||
mock_pyfs_instance = MagicMock()
|
||||
mock_pyfs.return_value = mock_pyfs_instance
|
||||
|
||||
# ABFSS URI with @ symbol should not trigger anonymous access logic
|
||||
fs, path = PyArrowFileSystem.get_fs_and_path(
|
||||
"abfss://container@account.dfs.core.windows.net/path"
|
||||
)
|
||||
|
||||
assert fs == mock_pyfs_instance
|
||||
assert path == "container/path"
|
||||
|
||||
# Verify that DefaultAzureCredential was used, not anonymous access
|
||||
mock_cred.assert_called_once()
|
||||
mock_adlfs.assert_called_once_with(
|
||||
account_name="account", credential=mock_cred.return_value
|
||||
)
|
||||
|
||||
@patch("adlfs.AzureBlobFileSystem")
|
||||
@patch("azure.identity.DefaultAzureCredential")
|
||||
@patch("pyarrow.fs.PyFileSystem")
|
||||
@patch("pyarrow.fs.FSSpecHandler")
|
||||
def test_get_file_abfss(self, mock_handler, mock_pyfs, mock_cred, mock_adlfs):
|
||||
"""Test getting a file from ABFSS storage."""
|
||||
# Setup mock filesystem and file
|
||||
mock_adlfs_instance = MagicMock()
|
||||
mock_adlfs.return_value = mock_adlfs_instance
|
||||
mock_fs = MagicMock()
|
||||
mock_pyfs.return_value = mock_fs
|
||||
|
||||
# Mock file content and info
|
||||
mock_file = MagicMock()
|
||||
mock_file.read.return_value = b"test abfss content"
|
||||
mock_fs.open_input_file.return_value.__enter__.return_value = mock_file
|
||||
mock_fs.get_file_info.return_value.type = pa_fs.FileType.File
|
||||
|
||||
# Test getting file as string (default)
|
||||
content = PyArrowFileSystem.get_file(
|
||||
"abfss://container@account.dfs.core.windows.net/test.txt"
|
||||
)
|
||||
assert content == "test abfss content"
|
||||
|
||||
# Verify the correct path was used
|
||||
mock_fs.get_file_info.assert_called_with("container/test.txt")
|
||||
mock_fs.open_input_file.assert_called_with("container/test.txt")
|
||||
|
||||
@patch("adlfs.AzureBlobFileSystem")
|
||||
@patch("azure.identity.DefaultAzureCredential")
|
||||
@patch("pyarrow.fs.PyFileSystem")
|
||||
@patch("pyarrow.fs.FSSpecHandler")
|
||||
def test_list_subfolders_abfss(
|
||||
self, mock_handler, mock_pyfs, mock_cred, mock_adlfs
|
||||
):
|
||||
"""Test listing subfolders in ABFSS storage."""
|
||||
# Setup mock filesystem
|
||||
mock_adlfs_instance = MagicMock()
|
||||
mock_adlfs.return_value = mock_adlfs_instance
|
||||
mock_fs = MagicMock()
|
||||
mock_pyfs.return_value = mock_fs
|
||||
|
||||
# Create mock file infos for directory listing
|
||||
dir1 = MagicMock()
|
||||
dir1.type = pa_fs.FileType.Directory
|
||||
dir1.path = "container/parent/subdir1"
|
||||
|
||||
dir2 = MagicMock()
|
||||
dir2.type = pa_fs.FileType.Directory
|
||||
dir2.path = "container/parent/subdir2"
|
||||
|
||||
file1 = MagicMock()
|
||||
file1.type = pa_fs.FileType.File
|
||||
file1.path = "container/parent/file.txt"
|
||||
|
||||
mock_fs.get_file_info.return_value = [dir1, dir2, file1]
|
||||
|
||||
# Test listing subfolders
|
||||
folders = PyArrowFileSystem.list_subfolders(
|
||||
"abfss://container@account.dfs.core.windows.net/parent"
|
||||
)
|
||||
assert sorted(folders) == ["subdir1", "subdir2"]
|
||||
|
||||
# Verify the correct path was used
|
||||
mock_fs.get_file_info.assert_called_once()
|
||||
call_args = mock_fs.get_file_info.call_args[0][0]
|
||||
assert call_args.base_dir == "container/parent/"
|
||||
assert call_args.recursive is False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,522 @@
|
||||
"""Tests for S3FileSystem class."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from botocore.exceptions import ClientError
|
||||
|
||||
from ray.llm._internal.common.utils.cloud_filesystem.s3_filesystem import (
|
||||
S3FileSystem,
|
||||
)
|
||||
|
||||
|
||||
class TestS3FileSystem:
|
||||
"""Tests for the S3FileSystem class."""
|
||||
|
||||
@patch("boto3.client")
|
||||
@pytest.mark.parametrize(
|
||||
"decode_as_utf_8,file_content,expected_content",
|
||||
[
|
||||
(True, b"test file content", "test file content"),
|
||||
(False, b"test file content", b"test file content"),
|
||||
],
|
||||
)
|
||||
def test_get_file(
|
||||
self,
|
||||
mock_boto_client,
|
||||
decode_as_utf_8,
|
||||
file_content,
|
||||
expected_content,
|
||||
):
|
||||
"""Test getting a file from S3 as string or bytes."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
# Mock get_object response
|
||||
mock_body = MagicMock()
|
||||
mock_body.read.return_value = file_content
|
||||
mock_s3_client.get_object.return_value = {
|
||||
"Body": mock_body,
|
||||
"ContentLength": len(file_content),
|
||||
}
|
||||
|
||||
# Test getting file
|
||||
content = S3FileSystem.get_file(
|
||||
"s3://bucket/test.txt", decode_as_utf_8=decode_as_utf_8
|
||||
)
|
||||
|
||||
assert content == expected_content
|
||||
mock_s3_client.get_object.assert_called_once_with(
|
||||
Bucket="bucket", Key="test.txt"
|
||||
)
|
||||
|
||||
@patch("boto3.client")
|
||||
def test_get_file_not_found(self, mock_boto_client):
|
||||
"""Test getting a non-existent file from S3."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
# Simulate NoSuchKey error
|
||||
mock_s3_client.get_object.side_effect = ClientError(
|
||||
{
|
||||
"Error": {
|
||||
"Code": "NoSuchKey",
|
||||
"Message": "The specified key does not exist.",
|
||||
}
|
||||
},
|
||||
"GetObject",
|
||||
)
|
||||
|
||||
assert S3FileSystem.get_file("s3://bucket/nonexistent.txt") is None
|
||||
|
||||
@patch("boto3.client")
|
||||
def test_get_file_anonymous(self, mock_boto_client):
|
||||
"""Test getting a file from S3 with anonymous access."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
# Mock get_object response
|
||||
mock_body = MagicMock()
|
||||
mock_body.read.return_value = b"anonymous content"
|
||||
mock_s3_client.get_object.return_value = {
|
||||
"Body": mock_body,
|
||||
}
|
||||
|
||||
# Test getting file with anonymous URI
|
||||
content = S3FileSystem.get_file("s3://anonymous@bucket/test.txt")
|
||||
|
||||
assert content == "anonymous content"
|
||||
# Verify anonymous config was used (UNSIGNED signature)
|
||||
assert mock_boto_client.called
|
||||
|
||||
@patch("boto3.client")
|
||||
@pytest.mark.parametrize(
|
||||
"uri,expected_prefix",
|
||||
[
|
||||
("s3://bucket/parent", "parent/"),
|
||||
("s3://bucket/parent/", "parent/"),
|
||||
],
|
||||
)
|
||||
def test_list_subfolders(self, mock_boto_client, uri, expected_prefix):
|
||||
"""Test listing subfolders in S3."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
# Mock list_objects_v2 response
|
||||
mock_s3_client.list_objects_v2.return_value = {
|
||||
"CommonPrefixes": [
|
||||
{"Prefix": f"{expected_prefix}folder1/"},
|
||||
{"Prefix": f"{expected_prefix}folder2/"},
|
||||
]
|
||||
}
|
||||
|
||||
folders = S3FileSystem.list_subfolders(uri)
|
||||
assert sorted(folders) == ["folder1", "folder2"]
|
||||
mock_s3_client.list_objects_v2.assert_called_once_with(
|
||||
Bucket="bucket", Prefix=expected_prefix, Delimiter="/"
|
||||
)
|
||||
|
||||
@patch("boto3.client")
|
||||
def test_list_subfolders_exception(self, mock_boto_client):
|
||||
"""Test listing subfolders when operation fails."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
mock_s3_client.list_objects_v2.side_effect = Exception("Network error")
|
||||
assert S3FileSystem.list_subfolders("s3://bucket/parent") == []
|
||||
|
||||
def test_list_subfolders_invalid_uri(self):
|
||||
"""Test listing subfolders with invalid URI."""
|
||||
# list_subfolders catches all exceptions and returns empty list
|
||||
result = S3FileSystem.list_subfolders("gs://bucket/parent")
|
||||
assert result == []
|
||||
|
||||
@patch("boto3.client")
|
||||
@pytest.mark.parametrize(
|
||||
"uri,expected_prefix",
|
||||
[
|
||||
("s3://bucket/dir", "dir/"),
|
||||
("s3://bucket/dir/", "dir/"),
|
||||
],
|
||||
)
|
||||
def test_download_files(self, mock_boto_client, uri, expected_prefix):
|
||||
"""Test downloading files from S3."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
# Mock paginator
|
||||
mock_paginator = MagicMock()
|
||||
mock_s3_client.get_paginator.return_value = mock_paginator
|
||||
mock_paginator.paginate.return_value = [
|
||||
{
|
||||
"Contents": [
|
||||
{"Key": f"{expected_prefix}file1.txt", "Size": 100},
|
||||
{"Key": f"{expected_prefix}file2.txt", "Size": 200},
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
# Mock download_file to do nothing
|
||||
mock_s3_client.download_file = MagicMock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
S3FileSystem.download_files(tempdir, uri, max_workers=2)
|
||||
|
||||
# Verify paginator was called correctly
|
||||
mock_s3_client.get_paginator.assert_called_with("list_objects_v2")
|
||||
mock_paginator.paginate.assert_called_once_with(
|
||||
Bucket="bucket", Prefix=expected_prefix
|
||||
)
|
||||
|
||||
# Verify files were downloaded
|
||||
assert mock_s3_client.download_file.call_count == 2
|
||||
|
||||
@patch("boto3.client")
|
||||
def test_download_files_with_filters(self, mock_boto_client):
|
||||
"""Test downloading files with inclusion and exclusion filters."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
# Mock paginator with various files
|
||||
mock_paginator = MagicMock()
|
||||
mock_s3_client.get_paginator.return_value = mock_paginator
|
||||
mock_paginator.paginate.return_value = [
|
||||
{
|
||||
"Contents": [
|
||||
{"Key": "dir/config.json", "Size": 100},
|
||||
{"Key": "dir/tokenizer.json", "Size": 200},
|
||||
{"Key": "dir/data.tmp", "Size": 300},
|
||||
{"Key": "dir/readme.txt", "Size": 400},
|
||||
{"Key": "dir/other.bin", "Size": 500},
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
# Mock download_file to do nothing
|
||||
mock_s3_client.download_file = MagicMock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
S3FileSystem.download_files(
|
||||
tempdir,
|
||||
"s3://bucket/dir",
|
||||
substrings_to_include=["config", "tokenizer"],
|
||||
suffixes_to_exclude=[".tmp", ".txt"],
|
||||
max_workers=2,
|
||||
)
|
||||
|
||||
# Should only download config.json and tokenizer.json
|
||||
# (included by substring, not excluded by suffix)
|
||||
assert mock_s3_client.download_file.call_count == 2
|
||||
|
||||
# Get the keys that were downloaded
|
||||
downloaded_keys = [
|
||||
call[0][1] for call in mock_s3_client.download_file.call_args_list
|
||||
]
|
||||
assert "dir/config.json" in downloaded_keys
|
||||
assert "dir/tokenizer.json" in downloaded_keys
|
||||
|
||||
@patch("boto3.client")
|
||||
def test_download_files_no_matching_files(self, mock_boto_client):
|
||||
"""Test downloading when no files match filters."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
# Mock paginator with files that won't match
|
||||
mock_paginator = MagicMock()
|
||||
mock_s3_client.get_paginator.return_value = mock_paginator
|
||||
mock_paginator.paginate.return_value = [
|
||||
{
|
||||
"Contents": [
|
||||
{"Key": "dir/file1.txt", "Size": 100},
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# This should not raise, just return without downloading
|
||||
S3FileSystem.download_files(
|
||||
tempdir,
|
||||
"s3://bucket/dir",
|
||||
substrings_to_include=["nonexistent"],
|
||||
max_workers=2,
|
||||
)
|
||||
|
||||
# Verify no files were downloaded
|
||||
mock_s3_client.download_file.assert_not_called()
|
||||
|
||||
@patch("boto3.client")
|
||||
@patch("ray.llm._internal.common.utils.cloud_filesystem.s3_filesystem.logger")
|
||||
def test_download_files_concurrent_failure(self, mock_logger, mock_boto_client):
|
||||
"""Test downloading files when some downloads fail."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
|
||||
# Mock paginator
|
||||
mock_paginator = MagicMock()
|
||||
mock_s3_client.get_paginator.return_value = mock_paginator
|
||||
mock_paginator.paginate.return_value = [
|
||||
{
|
||||
"Contents": [
|
||||
{"Key": "dir/file1.txt", "Size": 100},
|
||||
{"Key": "dir/file2.txt", "Size": 200},
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
# Make download_file fail
|
||||
mock_s3_client.download_file.side_effect = Exception("Download failed")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Should complete without raising, but log errors
|
||||
S3FileSystem.download_files(tempdir, "s3://bucket/dir", max_workers=2)
|
||||
|
||||
# Verify error was logged for failed downloads
|
||||
mock_logger.error.assert_called()
|
||||
error_call = mock_logger.error.call_args_list[0][0][0]
|
||||
assert "Failed to download" in error_call
|
||||
|
||||
def test_download_files_invalid_uri(self):
|
||||
"""Test downloading files with invalid URI."""
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
with pytest.raises(ValueError, match="Invalid S3 URI"):
|
||||
S3FileSystem.download_files(tempdir, "gs://bucket/dir")
|
||||
|
||||
@patch("boto3.client")
|
||||
@pytest.mark.parametrize(
|
||||
"uri,expected_prefix",
|
||||
[
|
||||
("s3://bucket/dir", "dir/"),
|
||||
("s3://bucket/dir/", "dir/"),
|
||||
],
|
||||
)
|
||||
def test_upload_files_directory(self, mock_boto_client, uri, expected_prefix):
|
||||
"""Test uploading a directory to S3."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
mock_s3_client.upload_file = MagicMock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Create some test files
|
||||
test_file1 = os.path.join(tempdir, "file1.txt")
|
||||
test_file2 = os.path.join(tempdir, "subdir", "file2.txt")
|
||||
os.makedirs(os.path.dirname(test_file2), exist_ok=True)
|
||||
|
||||
with open(test_file1, "w") as f:
|
||||
f.write("test1")
|
||||
with open(test_file2, "w") as f:
|
||||
f.write("test2")
|
||||
|
||||
S3FileSystem.upload_files(tempdir, uri)
|
||||
|
||||
# Verify files were uploaded
|
||||
assert mock_s3_client.upload_file.call_count == 2
|
||||
|
||||
# Check the S3 keys that were used
|
||||
uploaded_keys = [
|
||||
call[0][2] for call in mock_s3_client.upload_file.call_args_list
|
||||
]
|
||||
assert f"{expected_prefix}file1.txt" in uploaded_keys
|
||||
assert f"{expected_prefix}subdir/file2.txt" in uploaded_keys
|
||||
|
||||
@patch("boto3.client")
|
||||
def test_upload_files_single_file(self, mock_boto_client):
|
||||
"""Test uploading a single file to S3."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
mock_s3_client.upload_file = MagicMock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Create a test file
|
||||
test_file = os.path.join(tempdir, "single.txt")
|
||||
with open(test_file, "w") as f:
|
||||
f.write("test content")
|
||||
|
||||
S3FileSystem.upload_files(test_file, "s3://bucket/dir/")
|
||||
|
||||
# Verify single file was uploaded
|
||||
mock_s3_client.upload_file.assert_called_once()
|
||||
call_args = mock_s3_client.upload_file.call_args[0]
|
||||
assert call_args[2] == "dir/single.txt"
|
||||
|
||||
@patch("boto3.client")
|
||||
def test_upload_files_exception(self, mock_boto_client):
|
||||
"""Test uploading files when operation fails."""
|
||||
# Setup mock S3 client
|
||||
mock_s3_client = MagicMock()
|
||||
mock_boto_client.return_value = mock_s3_client
|
||||
mock_s3_client.upload_file.side_effect = Exception("Network error")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Create a test file
|
||||
test_file = os.path.join(tempdir, "test.txt")
|
||||
with open(test_file, "w") as f:
|
||||
f.write("test")
|
||||
|
||||
with pytest.raises(Exception, match="Network error"):
|
||||
S3FileSystem.upload_files(tempdir, "s3://bucket/dir")
|
||||
|
||||
def test_upload_files_invalid_uri(self):
|
||||
"""Test uploading files with invalid URI."""
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
with pytest.raises(ValueError, match="Invalid S3 URI"):
|
||||
S3FileSystem.upload_files(tempdir, "gs://bucket/dir")
|
||||
|
||||
def test_upload_files_nonexistent_path(self):
|
||||
"""Test uploading from a path that doesn't exist."""
|
||||
with pytest.raises(ValueError, match="does not exist"):
|
||||
S3FileSystem.upload_files("/nonexistent/path", "s3://bucket/dir")
|
||||
|
||||
def test_parse_s3_uri(self):
|
||||
"""Test parsing S3 URIs."""
|
||||
# Standard URI
|
||||
bucket, key, is_anon = S3FileSystem._parse_s3_uri(
|
||||
"s3://bucket/path/to/file.txt"
|
||||
)
|
||||
assert bucket == "bucket"
|
||||
assert key == "path/to/file.txt"
|
||||
assert is_anon is False
|
||||
|
||||
# Anonymous URI
|
||||
bucket, key, is_anon = S3FileSystem._parse_s3_uri(
|
||||
"s3://anonymous@bucket/file.txt"
|
||||
)
|
||||
assert bucket == "bucket"
|
||||
assert key == "file.txt"
|
||||
assert is_anon is True
|
||||
|
||||
# Bucket only
|
||||
bucket, key, is_anon = S3FileSystem._parse_s3_uri("s3://bucket")
|
||||
assert bucket == "bucket"
|
||||
assert key == ""
|
||||
assert is_anon is False
|
||||
|
||||
def test_calculate_optimal_workers(self):
|
||||
"""Test worker calculation based on file characteristics."""
|
||||
# Many small files (< 1MB)
|
||||
workers = S3FileSystem._calculate_optimal_workers(
|
||||
num_files=50, total_size=50 * 500 * 1024 # 50 files * 500KB each
|
||||
)
|
||||
assert workers == 50 # Should use many workers for small files
|
||||
|
||||
# Medium files (1-10MB)
|
||||
workers = S3FileSystem._calculate_optimal_workers(
|
||||
num_files=50, total_size=50 * 5 * 1024 * 1024 # 50 files * 5MB each
|
||||
)
|
||||
assert workers == 25 # Should use moderate workers
|
||||
|
||||
# Large files (> 10MB)
|
||||
workers = S3FileSystem._calculate_optimal_workers(
|
||||
num_files=50, total_size=50 * 50 * 1024 * 1024 # 50 files * 50MB each
|
||||
)
|
||||
assert workers == 20 # Should cap at 20 for large files
|
||||
|
||||
# Zero files
|
||||
workers = S3FileSystem._calculate_optimal_workers(num_files=0, total_size=0)
|
||||
assert workers == 10 # Should return default_min
|
||||
|
||||
|
||||
class TestS3FileSystemIntegration:
|
||||
"""Integration tests for S3FileSystem (requires actual S3 access)."""
|
||||
|
||||
def test_list_subfolders_real_s3(self):
|
||||
"""Test listing subfolders from real S3 bucket."""
|
||||
# Test listing subfolders in the parent directory which has actual subfolders
|
||||
folders = S3FileSystem.list_subfolders(
|
||||
"s3://anonymous@air-example-data/rayllm-ossci/"
|
||||
)
|
||||
# Verify we get expected subfolders
|
||||
assert isinstance(folders, list)
|
||||
assert "meta-Llama-3.2-1B-Instruct" in folders
|
||||
assert len(folders) > 0
|
||||
|
||||
def test_get_file_real_s3(self):
|
||||
"""Test getting a file from real S3 bucket."""
|
||||
# Test getting a small config file
|
||||
content = S3FileSystem.get_file(
|
||||
"s3://anonymous@air-example-data/rayllm-ossci/meta-Llama-3.2-1B-Instruct/config.json"
|
||||
)
|
||||
assert content is not None
|
||||
assert isinstance(content, str)
|
||||
# Verify it's valid JSON
|
||||
config = json.loads(content)
|
||||
assert "model_type" in config or "vocab_size" in config
|
||||
|
||||
def test_download_files_with_exclusion(self):
|
||||
"""Test downloading files with exclusion filter (exclude safetensors files)."""
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Download files excluding safetensors
|
||||
S3FileSystem.download_files(
|
||||
tempdir,
|
||||
"s3://anonymous@air-example-data/rayllm-ossci/meta-Llama-3.2-1B-Instruct/",
|
||||
suffixes_to_exclude=[".safetensors"],
|
||||
)
|
||||
|
||||
# Get list of downloaded files
|
||||
downloaded_files = set()
|
||||
for root, dirs, files in os.walk(tempdir):
|
||||
for file in files:
|
||||
rel_path = os.path.relpath(os.path.join(root, file), tempdir)
|
||||
downloaded_files.add(rel_path)
|
||||
|
||||
# Verify safetensors file is excluded
|
||||
assert (
|
||||
"model.safetensors" not in downloaded_files
|
||||
), "safetensors file should be excluded"
|
||||
|
||||
# Verify other files are downloaded
|
||||
assert "config.json" in downloaded_files
|
||||
assert "tokenizer.json" in downloaded_files
|
||||
assert len(downloaded_files) > 0
|
||||
|
||||
def test_download_files_with_inclusion(self):
|
||||
"""Test downloading files with inclusion filter (include only .json files)."""
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
# Download only .json files
|
||||
S3FileSystem.download_files(
|
||||
tempdir,
|
||||
"s3://anonymous@air-example-data/rayllm-ossci/meta-Llama-3.2-1B-Instruct/",
|
||||
substrings_to_include=[".json"],
|
||||
)
|
||||
|
||||
# Get list of downloaded files
|
||||
downloaded_files = set()
|
||||
for root, dirs, files in os.walk(tempdir):
|
||||
for file in files:
|
||||
rel_path = os.path.relpath(os.path.join(root, file), tempdir)
|
||||
downloaded_files.add(rel_path)
|
||||
|
||||
# Verify only .json files are downloaded
|
||||
expected_json_files = {
|
||||
"config.json",
|
||||
"generation_config.json",
|
||||
"special_tokens_map.json",
|
||||
"tokenizer.json",
|
||||
"tokenizer_config.json",
|
||||
}
|
||||
assert (
|
||||
downloaded_files == expected_json_files
|
||||
), f"Expected {expected_json_files}, got {downloaded_files}"
|
||||
|
||||
# Verify non-json files are excluded
|
||||
assert "model.safetensors" not in downloaded_files
|
||||
assert "README.md" not in downloaded_files
|
||||
assert "LICENSE.txt" not in downloaded_files
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,320 @@
|
||||
"""Utility tests for cloud functionality."""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.llm._internal.common.utils.cloud_utils import (
|
||||
CloudObjectCache,
|
||||
is_remote_path,
|
||||
remote_object_cache,
|
||||
)
|
||||
|
||||
|
||||
class MockSyncFetcher:
|
||||
def __init__(self):
|
||||
self.call_count = 0
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, key: str):
|
||||
self.call_count += 1
|
||||
self.calls.append(key)
|
||||
if key == "missing":
|
||||
return -1
|
||||
return f"value-{key}"
|
||||
|
||||
|
||||
class MockAsyncFetcher:
|
||||
def __init__(self):
|
||||
self.call_count = 0
|
||||
self.calls = []
|
||||
|
||||
async def __call__(self, key: str):
|
||||
self.call_count += 1
|
||||
self.calls.append(key)
|
||||
if key == "missing":
|
||||
return -1
|
||||
return f"value-{key}"
|
||||
|
||||
|
||||
class TestCloudObjectCache:
|
||||
"""Tests for the CloudObjectCache class."""
|
||||
|
||||
def test_sync_cache_basic(self):
|
||||
"""Test basic synchronous cache functionality."""
|
||||
fetcher = MockSyncFetcher()
|
||||
cache = CloudObjectCache(max_size=2, fetch_fn=fetcher)
|
||||
|
||||
# Test fetching a value (should be a miss)
|
||||
assert cache.get("key1") == "value-key1"
|
||||
assert fetcher.call_count == 1
|
||||
assert fetcher.calls == ["key1"]
|
||||
|
||||
# Test cache hit (should not call fetcher)
|
||||
assert cache.get("key1") == "value-key1"
|
||||
assert fetcher.call_count == 1 # Count should not increase
|
||||
assert fetcher.calls == ["key1"] # Calls should not change
|
||||
|
||||
# Test cache size limit
|
||||
assert cache.get("key2") == "value-key2" # Miss, should call fetcher
|
||||
assert fetcher.call_count == 2
|
||||
assert fetcher.calls == ["key1", "key2"]
|
||||
|
||||
assert (
|
||||
cache.get("key3") == "value-key3"
|
||||
) # Miss, should call fetcher and evict key1
|
||||
assert fetcher.call_count == 3
|
||||
assert fetcher.calls == ["key1", "key2", "key3"]
|
||||
|
||||
assert len(cache) == 2
|
||||
|
||||
# Verify key1 was evicted by checking if it's fetched again
|
||||
assert cache.get("key1") == "value-key1" # Miss, should call fetcher
|
||||
assert fetcher.call_count == 4
|
||||
assert fetcher.calls == ["key1", "key2", "key3", "key1"]
|
||||
|
||||
# Verify final cache state
|
||||
assert len(cache) == 2
|
||||
assert "key3" in cache._cache # key3 should still be in cache
|
||||
assert "key1" in cache._cache # key1 should be back in cache
|
||||
assert "key2" not in cache._cache # key2 should have been evicted
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_cache_missing_object_expiration(self):
|
||||
"""Test cache expiration for missing objects in async mode."""
|
||||
fetcher = MockAsyncFetcher()
|
||||
cache = CloudObjectCache(
|
||||
max_size=2,
|
||||
fetch_fn=fetcher,
|
||||
missing_expire_seconds=1, # 1 second to expire missing object
|
||||
exists_expire_seconds=3, # 3 seconds to expire existing object
|
||||
missing_object_value=-1,
|
||||
)
|
||||
|
||||
# Test missing object expiration
|
||||
assert await cache.aget("missing") is -1 # First fetch
|
||||
assert fetcher.call_count == 1
|
||||
assert fetcher.calls == ["missing"]
|
||||
|
||||
# Should still be cached
|
||||
assert await cache.aget("missing") is -1 # Cache hit
|
||||
assert fetcher.call_count == 1 # No new fetch
|
||||
assert fetcher.calls == ["missing"]
|
||||
|
||||
await asyncio.sleep(1.5) # Wait for missing object to expire
|
||||
assert await cache.aget("missing") is -1 # Should fetch again after expiration
|
||||
assert fetcher.call_count == 2 # New fetch
|
||||
assert fetcher.calls == ["missing", "missing"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_cache_existing_object_expiration(self):
|
||||
"""Test expiration of existing objects in async mode."""
|
||||
fetcher = MockAsyncFetcher()
|
||||
cache = CloudObjectCache(
|
||||
max_size=2,
|
||||
fetch_fn=fetcher,
|
||||
missing_expire_seconds=1, # 1 second to expire missing object
|
||||
exists_expire_seconds=3, # 3 seconds to expire existing object
|
||||
missing_object_value=-1,
|
||||
)
|
||||
|
||||
# Test existing object expiration
|
||||
assert await cache.aget("key1") == "value-key1" # First fetch
|
||||
assert fetcher.call_count == 1
|
||||
assert fetcher.calls == ["key1"]
|
||||
|
||||
# Should still be cached (not expired)
|
||||
assert await cache.aget("key1") == "value-key1" # Cache hit
|
||||
assert fetcher.call_count == 1 # No new fetch
|
||||
assert fetcher.calls == ["key1"]
|
||||
|
||||
await asyncio.sleep(1.5) # Not expired yet (exists_expire_seconds=3)
|
||||
assert await cache.aget("key1") == "value-key1" # Should still hit cache
|
||||
assert fetcher.call_count == 1 # No new fetch
|
||||
assert fetcher.calls == ["key1"] # No change in calls
|
||||
|
||||
await asyncio.sleep(2) # Now expired (total > 2 seconds)
|
||||
assert await cache.aget("key1") == "value-key1" # Should fetch again
|
||||
assert fetcher.call_count == 2 # New fetch
|
||||
assert fetcher.calls == ["key1", "key1"]
|
||||
|
||||
# Verify final cache state
|
||||
assert len(cache) == 1
|
||||
|
||||
|
||||
class TestRemoteObjectCacheDecorator:
|
||||
"""Tests for the remote_object_cache decorator."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_functionality(self):
|
||||
"""Test basic remote_object_cache decorator functionality."""
|
||||
call_count = 0
|
||||
MISSING = object()
|
||||
|
||||
@remote_object_cache(
|
||||
max_size=2,
|
||||
missing_expire_seconds=1,
|
||||
exists_expire_seconds=3,
|
||||
missing_object_value=MISSING,
|
||||
)
|
||||
async def fetch(key: str):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if key == "missing":
|
||||
return MISSING
|
||||
return f"value-{key}"
|
||||
|
||||
# Test cache hit
|
||||
assert await fetch("key1") == "value-key1"
|
||||
assert call_count == 1
|
||||
assert await fetch("key1") == "value-key1" # Should hit cache
|
||||
assert call_count == 1 # Count should not increase
|
||||
|
||||
# Test cache size limit
|
||||
assert await fetch("key2") == "value-key2"
|
||||
assert call_count == 2
|
||||
assert await fetch("key3") == "value-key3" # Should evict key1
|
||||
assert call_count == 3
|
||||
|
||||
# Verify key1 was evicted
|
||||
assert await fetch("key1") == "value-key1"
|
||||
assert call_count == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expiration(self):
|
||||
"""Test cache expiration for both missing and existing objects."""
|
||||
call_count = 0
|
||||
MISSING = object()
|
||||
|
||||
@remote_object_cache(
|
||||
max_size=2,
|
||||
missing_expire_seconds=1, # 1 second to expire missing object
|
||||
exists_expire_seconds=3, # 3 seconds to expire existing object
|
||||
missing_object_value=MISSING,
|
||||
)
|
||||
async def fetch(key: str):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if key == "missing":
|
||||
return MISSING
|
||||
return f"value-{key}"
|
||||
|
||||
# Test missing object expiration
|
||||
assert await fetch("missing") is MISSING
|
||||
assert call_count == 1
|
||||
assert await fetch("missing") is MISSING # Should hit cache
|
||||
assert call_count == 1
|
||||
|
||||
await asyncio.sleep(1.5) # Wait for missing object to expire
|
||||
assert await fetch("missing") is MISSING # Should fetch again
|
||||
assert call_count == 2
|
||||
|
||||
# Test existing object expiration
|
||||
assert await fetch("key1") == "value-key1"
|
||||
assert call_count == 3
|
||||
assert await fetch("key1") == "value-key1" # Should hit cache
|
||||
assert call_count == 3
|
||||
|
||||
await asyncio.sleep(1.5) # Not expired yet
|
||||
assert await fetch("key1") == "value-key1" # Should still hit cache
|
||||
assert call_count == 3
|
||||
|
||||
await asyncio.sleep(2) # Now expired (total > 3 seconds)
|
||||
assert await fetch("key1") == "value-key1" # Should fetch again
|
||||
assert call_count == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_handling(self):
|
||||
"""Test error handling in remote_object_cache decorator."""
|
||||
call_count = 0
|
||||
|
||||
@remote_object_cache(max_size=2)
|
||||
async def fetch(key: str):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if key == "error":
|
||||
raise ValueError("Test error")
|
||||
return f"value-{key}"
|
||||
|
||||
# Test successful case
|
||||
assert await fetch("key1") == "value-key1"
|
||||
assert call_count == 1
|
||||
|
||||
# Test error case
|
||||
with pytest.raises(ValueError, match="Test error"):
|
||||
await fetch("error")
|
||||
assert call_count == 2
|
||||
|
||||
# Verify error wasn't cached
|
||||
with pytest.raises(ValueError, match="Test error"):
|
||||
await fetch("error")
|
||||
assert call_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_access(self):
|
||||
"""Test concurrent access to cached function."""
|
||||
call_count = 0
|
||||
DELAY = 0.1
|
||||
|
||||
@remote_object_cache(max_size=2)
|
||||
async def slow_fetch(key: str):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
await asyncio.sleep(DELAY) # Simulate slow operation
|
||||
return f"value-{key}"
|
||||
|
||||
# Launch multiple concurrent calls
|
||||
tasks = [slow_fetch("key1") for _ in range(5)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# All results should be the same
|
||||
assert all(r == "value-key1" for r in results)
|
||||
# Should only call once despite multiple concurrent requests
|
||||
assert call_count == 1
|
||||
|
||||
|
||||
class TestIsRemotePath:
|
||||
"""Tests for the is_remote_path utility function."""
|
||||
|
||||
def test_s3_paths(self):
|
||||
"""Test S3 path detection."""
|
||||
assert is_remote_path("s3://bucket/path") is True
|
||||
assert is_remote_path("s3://bucket") is True
|
||||
assert is_remote_path("s3://anonymous@bucket/path") is True
|
||||
|
||||
def test_gcs_paths(self):
|
||||
"""Test GCS path detection."""
|
||||
assert is_remote_path("gs://bucket/path") is True
|
||||
assert is_remote_path("gs://bucket") is True
|
||||
assert is_remote_path("gs://anonymous@bucket/path") is True
|
||||
|
||||
def test_abfss_paths(self):
|
||||
"""Test ABFSS path detection."""
|
||||
assert (
|
||||
is_remote_path("abfss://container@account.dfs.core.windows.net/path")
|
||||
is True
|
||||
)
|
||||
assert is_remote_path("abfss://container@account.dfs.core.windows.net") is True
|
||||
|
||||
def test_azure_paths(self):
|
||||
"""Test Azure path detection."""
|
||||
assert (
|
||||
is_remote_path("azure://container@account.blob.core.windows.net/path")
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
is_remote_path("azure://container@account.dfs.core.windows.net/path")
|
||||
is True
|
||||
)
|
||||
|
||||
def test_local_paths(self):
|
||||
"""Test local path detection."""
|
||||
assert is_remote_path("/local/path") is False
|
||||
assert is_remote_path("./relative/path") is False
|
||||
assert is_remote_path("file:///local/path") is False
|
||||
assert is_remote_path("http://example.com") is False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Unit tests for ray.llm._internal.common.placement."""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from ray.llm._internal.common.placement import BundleConfig, PlacementGroupConfig
|
||||
|
||||
|
||||
def test_bundle_config_defaults_returns_float():
|
||||
b = BundleConfig()
|
||||
d = b.model_dump()
|
||||
assert d["CPU"] == 0.0 and d["GPU"] == 0.0
|
||||
|
||||
b2 = BundleConfig(CPU=2, GPU=1)
|
||||
assert b2.model_dump() == {"CPU": 2.0, "GPU": 1.0}
|
||||
|
||||
|
||||
def test_bundle_config_fractional_gpu():
|
||||
b = BundleConfig(GPU=0.5, CPU=1.0)
|
||||
assert b.model_dump()["GPU"] == 0.5
|
||||
|
||||
|
||||
def test_bundle_config_extra_custom_resource():
|
||||
b = BundleConfig(CPU=0.0, GPU=0.0, TPU=4.0)
|
||||
d = b.model_dump()
|
||||
assert d["TPU"] == 4.0
|
||||
|
||||
|
||||
def test_bundle_config_extra_resource_negative_rejected():
|
||||
with pytest.raises(ValueError, match="non-negative"):
|
||||
BundleConfig(TPU=-1.0)
|
||||
|
||||
|
||||
def test_bundle_config_extra_resource_non_numeric_rejected():
|
||||
with pytest.raises(ValueError, match="must be a number"):
|
||||
BundleConfig(CPU=0.0, GPU=0.0, bad="x")
|
||||
|
||||
|
||||
def test_placement_group_config_bundles_only():
|
||||
pg = PlacementGroupConfig(
|
||||
bundles=[BundleConfig(GPU=1.0, CPU=1.0)],
|
||||
strategy="PACK",
|
||||
)
|
||||
out = pg.model_dump()
|
||||
assert out["bundles"][0]["GPU"] == 1.0
|
||||
assert out["strategy"] == "PACK"
|
||||
|
||||
|
||||
def test_placement_group_config_bundle_per_worker_only():
|
||||
pg = PlacementGroupConfig(
|
||||
bundle_per_worker=BundleConfig(GPU=1.0, CPU=2.0),
|
||||
strategy="SPREAD",
|
||||
)
|
||||
out = pg.model_dump()
|
||||
assert out["bundle_per_worker"]["GPU"] == 1.0
|
||||
assert out["strategy"] == "SPREAD"
|
||||
|
||||
|
||||
def test_placement_group_config_rejects_neither_bundles_nor_per_worker():
|
||||
with pytest.raises(ValueError, match="either 'bundle_per_worker'"):
|
||||
PlacementGroupConfig(strategy="PACK")
|
||||
|
||||
|
||||
def test_placement_group_config_rejects_both_bundles_and_per_worker():
|
||||
with pytest.raises(ValueError, match="Cannot specify both"):
|
||||
PlacementGroupConfig(
|
||||
bundles=[BundleConfig(GPU=1.0)],
|
||||
bundle_per_worker=BundleConfig(GPU=1.0),
|
||||
)
|
||||
|
||||
|
||||
def test_placement_group_config_invalid_strategy_rejected():
|
||||
with pytest.raises(ValidationError):
|
||||
PlacementGroupConfig.model_validate(
|
||||
{
|
||||
"bundles": [{"GPU": 1.0, "CPU": 0.0}],
|
||||
"strategy": "INVALID",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_placement_group_config_from_raw_dict():
|
||||
pg = PlacementGroupConfig(
|
||||
**{"bundle_per_worker": {"CPU": 2, "GPU": 1}, "strategy": "SPREAD"}
|
||||
)
|
||||
assert pg.bundle_per_worker.CPU == 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,112 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.llm._internal.common.callbacks.base import (
|
||||
CallbackBase,
|
||||
)
|
||||
from ray.llm._internal.common.utils.download_utils import NodeModelDownloadable
|
||||
from ray.llm._internal.serve.core.configs.llm_config import (
|
||||
LLMConfig,
|
||||
ModelLoadingConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestingCallback(CallbackBase):
|
||||
def __init__(self, llm_config, raise_error_on_callback: bool = True, **kwargs):
|
||||
super().__init__(llm_config, raise_error_on_callback, **kwargs)
|
||||
self.before_init_called = False
|
||||
self.after_init_called = False
|
||||
self.before_init_ctx = None
|
||||
self.after_init_ctx = None
|
||||
assert kwargs["kwargs_test_key"] == "kwargs_test_value"
|
||||
|
||||
async def on_before_node_init(self) -> None:
|
||||
assert (
|
||||
self.ctx.worker_node_download_model
|
||||
== NodeModelDownloadable.MODEL_AND_TOKENIZER
|
||||
)
|
||||
self.ctx.worker_node_download_model = NodeModelDownloadable.NONE
|
||||
|
||||
self.ctx.custom_data["ctx_test_key"] = "ctx_test_value"
|
||||
self.before_init_called = True
|
||||
self.ctx.run_init_node = False
|
||||
|
||||
async def on_after_node_init(self) -> None:
|
||||
assert self.ctx.worker_node_download_model == NodeModelDownloadable.NONE
|
||||
|
||||
self.after_init_called = True
|
||||
assert self.ctx.custom_data["ctx_test_key"] == "ctx_test_value"
|
||||
|
||||
|
||||
class TestCallbackBase:
|
||||
@pytest.fixture
|
||||
def llm_config(self):
|
||||
config = LLMConfig(
|
||||
model_loading_config=ModelLoadingConfig(model_id="test-model"),
|
||||
llm_engine="vLLM",
|
||||
callback_config={
|
||||
"callback_class": TestingCallback,
|
||||
"callback_kwargs": {"kwargs_test_key": "kwargs_test_value"},
|
||||
},
|
||||
)
|
||||
return config
|
||||
|
||||
def test_callback_methods_called(self, llm_config):
|
||||
"""Test that callback methods are called during initialization."""
|
||||
|
||||
# Run initialization
|
||||
async def run_initialization():
|
||||
callback = llm_config.get_or_create_callback()
|
||||
await callback.run_callback("on_before_node_init")
|
||||
if callback.ctx.run_init_node:
|
||||
raise Exception("run_init_node is True")
|
||||
await callback.run_callback("on_after_node_init")
|
||||
|
||||
asyncio.run(run_initialization())
|
||||
# Verify callback was created and methods were called
|
||||
|
||||
callback = llm_config.get_or_create_callback()
|
||||
assert callback is not None
|
||||
assert isinstance(callback, TestingCallback)
|
||||
assert callback.before_init_called is True
|
||||
assert callback.after_init_called is True
|
||||
|
||||
def test_callback_singleton_behavior(self, llm_config):
|
||||
"""Test that callback instance is cached (singleton pattern)."""
|
||||
# Get callback multiple times
|
||||
callback1 = llm_config.get_or_create_callback()
|
||||
callback2 = llm_config.get_or_create_callback()
|
||||
|
||||
# Should be the same instance
|
||||
assert callback1 is callback2
|
||||
|
||||
def test_callback_must_inherit_from_callback_class(self):
|
||||
"""Test that callback_class must be a subclass of Callback, not just implement the same methods."""
|
||||
|
||||
class FakeCallback:
|
||||
"""A class that implements the same methods as Callback but doesn't inherit from it."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
async def on_before_node_init(self):
|
||||
pass
|
||||
|
||||
async def on_after_node_init(self):
|
||||
pass
|
||||
|
||||
# Should raise an error when trying to create callback
|
||||
with pytest.raises(Exception, match="is-subclass"):
|
||||
LLMConfig(
|
||||
model_loading_config=ModelLoadingConfig(model_id="test-model"),
|
||||
llm_engine="vLLM",
|
||||
callback_config={
|
||||
"callback_class": FakeCallback,
|
||||
"callback_kwargs": {},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main(["-v", __file__])
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import ANY, call, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.llm._internal.common.utils.upload_utils import upload_model_files
|
||||
|
||||
|
||||
@patch("pyarrow.fs.copy_files")
|
||||
def test_upload_undownloaded_model(mock_copy_files):
|
||||
model_id = "fake-model-id"
|
||||
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
model_dir = os.path.join(tempdir, model_id)
|
||||
# mock snapshot_download to create a subdir in the model_dir
|
||||
def create_subdir_side_effect(*args, **kwargs):
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
|
||||
with patch(
|
||||
"ray.llm._internal.common.utils.upload_utils.get_model_entrypoint",
|
||||
return_value=model_dir,
|
||||
), patch(
|
||||
"huggingface_hub.snapshot_download", side_effect=create_subdir_side_effect
|
||||
):
|
||||
# upload the model
|
||||
upload_model_files(model_id, "gs://bucket/model-id")
|
||||
|
||||
# check that the model was uploaded1
|
||||
mock_copy_files.assert_called_once_with(
|
||||
source=Path(model_dir),
|
||||
destination="bucket/model-id",
|
||||
source_filesystem=ANY,
|
||||
destination_filesystem=ANY,
|
||||
)
|
||||
|
||||
|
||||
@patch("pyarrow.fs.copy_files")
|
||||
def test_upload_downloaded_hf_model(mock_copy_files):
|
||||
model_id = "fake-model-id"
|
||||
hash = "0123456789abcdef"
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
model_dir = os.path.join(tempdir, model_id)
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
os.makedirs(os.path.join(model_dir, "snapshots", hash), exist_ok=True)
|
||||
model_rev_path = os.path.join(model_dir, "refs")
|
||||
os.makedirs(model_rev_path, exist_ok=True)
|
||||
with open(os.path.join(model_rev_path, "main"), "w") as f:
|
||||
f.write(hash)
|
||||
|
||||
with patch(
|
||||
"ray.llm._internal.common.utils.upload_utils.get_model_entrypoint",
|
||||
return_value=model_dir,
|
||||
):
|
||||
upload_model_files(model_id, "pyarrow-s3://bucket/model-id")
|
||||
|
||||
assert mock_copy_files.call_count == 2
|
||||
mock_copy_files.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
source=os.path.join(model_dir, "snapshots", hash),
|
||||
destination="bucket/model-id",
|
||||
source_filesystem=ANY,
|
||||
destination_filesystem=ANY,
|
||||
),
|
||||
call(
|
||||
source=os.path.join(model_rev_path, "main"),
|
||||
destination="bucket/model-id/hash",
|
||||
source_filesystem=ANY,
|
||||
destination_filesystem=ANY,
|
||||
),
|
||||
],
|
||||
any_order=True,
|
||||
)
|
||||
|
||||
|
||||
@patch("pyarrow.fs.copy_files")
|
||||
def test_upload_custom_model(mock_copy_files):
|
||||
model_id = "fake-model-id"
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
model_dir = os.path.join(tempdir, model_id)
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
|
||||
with patch(
|
||||
"ray.llm._internal.common.utils.upload_utils.get_model_entrypoint",
|
||||
return_value=model_dir,
|
||||
):
|
||||
upload_model_files(model_id, "pyarrow-s3://bucket/model-id")
|
||||
|
||||
mock_copy_files.assert_called_once_with(
|
||||
source=Path(model_dir),
|
||||
destination="bucket/model-id",
|
||||
source_filesystem=ANY,
|
||||
destination_filesystem=ANY,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user