wehub-resource-sync
127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
from pathlib import Path
|
|
|
|
import typer
|
|
from filelock import FileLock
|
|
from typing_extensions import Annotated
|
|
|
|
from ray.llm._internal.common.observability.logging import get_logger
|
|
from ray.llm._internal.common.utils.cloud_utils import (
|
|
CloudFileSystem,
|
|
CloudMirrorConfig,
|
|
CloudModelAccessor,
|
|
is_remote_path,
|
|
)
|
|
from ray.llm._internal.common.utils.download_utils import (
|
|
get_model_entrypoint,
|
|
)
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class CloudModelUploader(CloudModelAccessor):
|
|
"""Unified uploader to upload models to cloud storage (S3 or GCS).
|
|
|
|
Args:
|
|
model_id: The model id to upload.
|
|
mirror_config: The mirror config for the model.
|
|
"""
|
|
|
|
def upload_model(self) -> str:
|
|
"""Upload the model to cloud storage (s3 or gcs).
|
|
|
|
Returns:
|
|
The remote path of the uploaded model.
|
|
"""
|
|
bucket_uri = self.mirror_config.bucket_uri
|
|
|
|
lock_path = self._get_lock_path()
|
|
path = self._get_model_path()
|
|
storage_type = self.mirror_config.storage_type
|
|
|
|
try:
|
|
# Timeout 0 means there will be only one attempt to acquire
|
|
# the file lock. If it cannot be acquired, a TimeoutError
|
|
# will be thrown.
|
|
# This ensures that subsequent processes don't duplicate work.
|
|
with FileLock(lock_path, timeout=0):
|
|
try:
|
|
CloudFileSystem.upload_model(
|
|
local_path=path,
|
|
bucket_uri=bucket_uri,
|
|
)
|
|
logger.info(
|
|
"Finished uploading %s to %s storage",
|
|
self.model_id,
|
|
storage_type.upper() if storage_type else "cloud",
|
|
)
|
|
except RuntimeError:
|
|
logger.exception(
|
|
"Failed to upload model %s to %s storage",
|
|
self.model_id,
|
|
storage_type.upper() if storage_type else "cloud",
|
|
)
|
|
except TimeoutError:
|
|
# If the directory is already locked, then wait but do not do anything.
|
|
with FileLock(lock_path, timeout=-1):
|
|
pass
|
|
return bucket_uri
|
|
|
|
|
|
def upload_model_files(model_id: str, bucket_uri: str) -> str:
|
|
"""Upload the model files to cloud storage (s3 or gcs).
|
|
|
|
If `model_id` is a local path, the files will be uploaded to the cloud storage.
|
|
If `model_id` is a huggingface model id, the model will be downloaded from huggingface
|
|
and then uploaded to the cloud storage.
|
|
|
|
Args:
|
|
model_id: The huggingface model id, or local model path to upload.
|
|
bucket_uri: The bucket uri to upload the model to, must start with `s3://` or `gs://`.
|
|
|
|
Returns:
|
|
The remote path of the uploaded model.
|
|
"""
|
|
assert not is_remote_path(
|
|
model_id
|
|
), f"model_id must NOT be a remote path: {model_id}"
|
|
assert is_remote_path(bucket_uri), f"bucket_uri must be a remote path: {bucket_uri}"
|
|
|
|
if not Path(model_id).exists():
|
|
maybe_downloaded_model_path = get_model_entrypoint(model_id)
|
|
if not Path(maybe_downloaded_model_path).exists():
|
|
logger.info(
|
|
"Assuming %s is huggingface model id, and downloading it.", model_id
|
|
)
|
|
import huggingface_hub
|
|
|
|
huggingface_hub.snapshot_download(repo_id=model_id)
|
|
# Try to get the model path again after downloading.
|
|
maybe_downloaded_model_path = get_model_entrypoint(model_id)
|
|
assert Path(
|
|
maybe_downloaded_model_path
|
|
).exists(), f"Failed to download the model {model_id} to {maybe_downloaded_model_path}"
|
|
return upload_model_files(maybe_downloaded_model_path, bucket_uri)
|
|
else:
|
|
return upload_model_files(maybe_downloaded_model_path, bucket_uri)
|
|
|
|
uploader = CloudModelUploader(model_id, CloudMirrorConfig(bucket_uri=bucket_uri))
|
|
return uploader.upload_model()
|
|
|
|
|
|
def upload_model_cli(
|
|
model_source: Annotated[
|
|
str,
|
|
typer.Option(
|
|
help="HuggingFace model ID to download, or local model path to upload",
|
|
),
|
|
],
|
|
bucket_uri: Annotated[
|
|
str,
|
|
typer.Option(
|
|
help="The bucket uri to upload the model to, must start with `s3://` or `gs://`",
|
|
),
|
|
],
|
|
):
|
|
"""Upload the model files to cloud storage (s3 or gcs)."""
|
|
upload_model_files(model_source, bucket_uri)
|