wehub-resource-sync
83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
import functools
|
|
import logging
|
|
import os
|
|
import time
|
|
|
|
from mlflow.environment_variables import _MLFLOW_TESTING
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
# NB: The maxsize=1 is added for encouraging the cache refresh so the user doesn't get stale
|
|
# commit hash from the cache. This doesn't work perfectly because it only updates cache
|
|
# when the user calls it with a different repo name, but it's better than nothing.
|
|
@functools.lru_cache(maxsize=1)
|
|
def get_latest_commit_for_repo(repo: str) -> str:
|
|
"""
|
|
Fetches the latest commit hash for a repository from the HuggingFace model hub.
|
|
"""
|
|
try:
|
|
import huggingface_hub as hub
|
|
except ImportError:
|
|
raise MlflowException(
|
|
"Unable to fetch model commit hash from the HuggingFace model hub. "
|
|
"This is required for saving a model without base model "
|
|
"weights, while ensuring the version consistency of the model. "
|
|
"Please install the `huggingface-hub` package and retry.",
|
|
error_code=RESOURCE_DOES_NOT_EXIST,
|
|
)
|
|
|
|
from huggingface_hub.errors import HfHubHTTPError
|
|
|
|
api = hub.HfApi()
|
|
for i in range(7):
|
|
try:
|
|
return api.model_info(repo).sha
|
|
except HfHubHTTPError as e:
|
|
if not _MLFLOW_TESTING.get():
|
|
raise
|
|
|
|
# Retry on rate limit error
|
|
if e.response.status_code == 429:
|
|
_logger.warning(
|
|
f"Rate limit exceeded while fetching commit hash for repo {repo}. "
|
|
f"Retrying in {2**i} seconds. Error: {e}",
|
|
)
|
|
time.sleep(2**i)
|
|
continue
|
|
raise
|
|
|
|
raise MlflowException(
|
|
"Unable to fetch model commit hash from the HuggingFace model hub. "
|
|
"This is required for saving a model without base model "
|
|
"weights, while ensuring the version consistency of the model. ",
|
|
error_code=RESOURCE_DOES_NOT_EXIST,
|
|
)
|
|
|
|
|
|
def is_valid_hf_repo_id(maybe_repo_id: str | None) -> bool:
|
|
"""
|
|
Check if the given string is a valid HuggingFace repo identifier e.g. "username/repo_id".
|
|
"""
|
|
|
|
if not maybe_repo_id or os.path.isdir(maybe_repo_id):
|
|
return False
|
|
|
|
try:
|
|
from huggingface_hub.utils import HFValidationError, validate_repo_id
|
|
except ImportError:
|
|
raise MlflowException(
|
|
"Unable to validate the repository identifier for the HuggingFace model hub "
|
|
"because the `huggingface-hub` package is not installed. Please install the "
|
|
"package with `pip install huggingface-hub` command and retry."
|
|
)
|
|
|
|
try:
|
|
validate_repo_id(maybe_repo_id)
|
|
return True
|
|
except HFValidationError as e:
|
|
_logger.warning(f"The repository identified {maybe_repo_id} is invalid: {e}")
|
|
return False
|