def get_hf_hub_cache_repos()

in optimum/neuron/utils/cache_utils.py [0:0]

• 23 lines of code
• 13 McCabe index (conditional complexity)

def get_hf_hub_cache_repos(log_warnings: bool = False) -> List[str]:
    """
    Retrieves the name of the Hugging Face Hub model repo to use as remote cache.
    Priority:
        - If a repo is provided via the `CUSTOM_CACHE_REPO` environment variable, it will be used,
        - Else, if a custom cache repo has been set locally, it will be used,
        - Otherwise, it uses the default cache repo (on which most people do not have write access)
    """
    # Default hub repos.
    hf_hub_repos = HF_HUB_CACHE_REPOS

    # Locally saved hub repo.
    saved_custom_cache_repo = load_custom_cache_repo_name_from_hf_home()
    if saved_custom_cache_repo is not None and saved_custom_cache_repo not in hf_hub_repos:
        hf_hub_repos = [saved_custom_cache_repo] + hf_hub_repos

    # Hub repo set via the environment variable CUSTOM_CACHE_REPO.
    custom_cache_repo = os.environ.get("CUSTOM_CACHE_REPO", None)
    if custom_cache_repo is not None and custom_cache_repo not in hf_hub_repos:
        hf_hub_repos = [custom_cache_repo] + hf_hub_repos

    if log_warnings and is_main_worker() and saved_custom_cache_repo is None and custom_cache_repo is None:
        warn_once(
            logger,
            "No Neuron cache name is saved locally. This means that only the official Neuron cache will be used. You "
            "can create a Neuron cache repo by running the following command: `optimum-cli neuron cache create`. If "
            "the Neuron cache already exists you can set it by running the following command: `optimum-cli neuron cache "
            "set -n [name]`.",
        )

    if log_warnings and is_main_worker() and hf_hub_repos and not has_write_access_to_repo(hf_hub_repos[0]):
        warn_once(
            logger,
            f"You do not have write access to {hf_hub_repos[0]} so you will not be able to push any cached compilation "
            "files. Please log in and/or use a custom Neuron cache.",
        )
    return hf_hub_repos