wehub-resource-sync
102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
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__]))
|