JafarUruc/orange_cube

1import io2import os3import time4from importlib.resources import files5from pathlib import Path6 7import gradio8import huggingface_hub9from gradio_client import Client, handle_file10from httpx import ReadTimeout11from huggingface_hub.errors import RepositoryNotFoundError12 13from trackio.sqlite_storage import SQLiteStorage14 15SPACE_URL = "https://huggingface.co/spaces/{space_id}"16 17 18def deploy_as_space(19    space_id: str,20    dataset_id: str | None = None,21):22    if (23        os.getenv("SYSTEM") == "spaces"24    ):  # in case a repo with this function is uploaded to spaces25        return26 27    trackio_path = files("trackio")28 29    hf_api = huggingface_hub.HfApi()30    whoami = None31    login = False32    try:33        whoami = hf_api.whoami()34        if whoami["auth"]["accessToken"]["role"] != "write":35            login = True36    except OSError:37        login = True38    if login:39        print("Need 'write' access token to create a Spaces repo.")40        huggingface_hub.login(add_to_git_credential=False)41        whoami = hf_api.whoami()42 43    huggingface_hub.create_repo(44        space_id,45        space_sdk="gradio",46        repo_type="space",47        exist_ok=True,48    )49 50    with open(Path(trackio_path, "README.md"), "r") as f:51        readme_content = f.read()52        readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)53        readme_buffer = io.BytesIO(readme_content.encode("utf-8"))54        hf_api.upload_file(55            path_or_fileobj=readme_buffer,56            path_in_repo="README.md",57            repo_id=space_id,58            repo_type="space",59        )60 61    huggingface_hub.utils.disable_progress_bars()62    hf_api.upload_folder(63        repo_id=space_id,64        repo_type="space",65        folder_path=trackio_path,66        ignore_patterns=["README.md"],67    )68 69    hf_token = huggingface_hub.utils.get_token()70    if hf_token is not None:71        huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)72    if dataset_id is not None:73        huggingface_hub.add_space_variable(space_id, "TRACKIO_DATASET_ID", dataset_id)74 75 76def create_space_if_not_exists(77    space_id: str,78    dataset_id: str | None = None,79) -> None:80    """81    Creates a new Hugging Face Space if it does not exist. If a dataset_id is provided, it will be added as a space variable.82 83    Args:84        space_id: The ID of the Space to create.85        dataset_id: The ID of the Dataset to add to the Space.86    """87    if "/" not in space_id:88        raise ValueError(89            f"Invalid space ID: {space_id}. Must be in the format: username/reponame or orgname/reponame."90        )91    if dataset_id is not None and "/" not in dataset_id:92        raise ValueError(93            f"Invalid dataset ID: {dataset_id}. Must be in the format: username/datasetname or orgname/datasetname."94        )95    try:96        huggingface_hub.repo_info(space_id, repo_type="space")97        print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")98        if dataset_id is not None:99            huggingface_hub.add_space_variable(100                space_id, "TRACKIO_DATASET_ID", dataset_id101            )102        return103    except RepositoryNotFoundError:104        pass105 106    print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")107    deploy_as_space(space_id, dataset_id)108 109 110def wait_until_space_exists(111    space_id: str,112) -> None:113    """114    Blocks the current thread until the space exists.115    May raise a TimeoutError if this takes quite a while.116 117    Args:118        space_id: The ID of the Space to wait for.119    """120    client = None121    for _ in range(30):122        try:123            client = Client(space_id, verbose=False)124            if client:125                break126        except ReadTimeout:127            time.sleep(5)128        except ValueError:129            time.sleep(5)130    raise TimeoutError("Waiting for space to exist took longer than expected")131 132 133def upload_db_to_space(project: str, space_id: str) -> None:134    """135    Uploads the database of a local Trackio project to a Hugging Face Space.136 137    Args:138        project: The name of the project to upload.139        space_id: The ID of the Space to upload to.140    """141    db_path = SQLiteStorage.get_project_db_path(project)142    client = Client(space_id, verbose=False)143    client.predict(144        api_name="/upload_db_to_space",145        project=project,146        uploaded_db=handle_file(db_path),147        hf_token=huggingface_hub.utils.get_token(),148    )149