Tonic/l-android-control

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 RepositoryNotFoundError12from requests import HTTPError13 14from trackio.sqlite_storage import SQLiteStorage15 16SPACE_URL = "https://huggingface.co/spaces/{space_id}"17 18 19def deploy_as_space(20    space_id: str,21    dataset_id: str | None = None,22):23    if (24        os.getenv("SYSTEM") == "spaces"25    ):  # in case a repo with this function is uploaded to spaces26        return27 28    trackio_path = files("trackio")29 30    hf_api = huggingface_hub.HfApi()31 32    try:33        huggingface_hub.create_repo(34            space_id,35            space_sdk="gradio",36            repo_type="space",37            exist_ok=True,38        )39    except HTTPError as e:40        if e.response.status_code in [401, 403]:  # unauthorized or forbidden41            print("Need 'write' access token to create a Spaces repo.")42            huggingface_hub.login(add_to_git_credential=False)43            huggingface_hub.create_repo(44                space_id,45                space_sdk="gradio",46                repo_type="space",47                exist_ok=True,48            )49        else:50            raise ValueError(f"Failed to create Space: {e}")51 52    with open(Path(trackio_path, "README.md"), "r") as f:53        readme_content = f.read()54        readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)55        readme_buffer = io.BytesIO(readme_content.encode("utf-8"))56        hf_api.upload_file(57            path_or_fileobj=readme_buffer,58            path_in_repo="README.md",59            repo_id=space_id,60            repo_type="space",61        )62 63    # We can assume pandas, gradio, and huggingface-hub are already installed in a Gradio Space.64    # Make sure necessary dependencies are installed by creating a requirements.txt.65    requirements_content = """66pyarrow>=21.067    """68    requirements_buffer = io.BytesIO(requirements_content.encode("utf-8"))69    hf_api.upload_file(70        path_or_fileobj=requirements_buffer,71        path_in_repo="requirements.txt",72        repo_id=space_id,73        repo_type="space",74    )75 76    huggingface_hub.utils.disable_progress_bars()77    hf_api.upload_folder(78        repo_id=space_id,79        repo_type="space",80        folder_path=trackio_path,81        ignore_patterns=["README.md"],82    )83 84    hf_token = huggingface_hub.utils.get_token()85    if hf_token is not None:86        huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)87    if dataset_id is not None:88        huggingface_hub.add_space_variable(space_id, "TRACKIO_DATASET_ID", dataset_id)89 90 91def create_space_if_not_exists(92    space_id: str,93    dataset_id: str | None = None,94) -> None:95    """96    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.97 98    Args:99        space_id: The ID of the Space to create.100        dataset_id: The ID of the Dataset to add to the Space.101    """102    if "/" not in space_id:103        raise ValueError(104            f"Invalid space ID: {space_id}. Must be in the format: username/reponame or orgname/reponame."105        )106    if dataset_id is not None and "/" not in dataset_id:107        raise ValueError(108            f"Invalid dataset ID: {dataset_id}. Must be in the format: username/datasetname or orgname/datasetname."109        )110    try:111        huggingface_hub.repo_info(space_id, repo_type="space")112        print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")113        if dataset_id is not None:114            huggingface_hub.add_space_variable(115                space_id, "TRACKIO_DATASET_ID", dataset_id116            )117        return118    except RepositoryNotFoundError:119        pass120    except HTTPError as e:121        if e.response.status_code in [401, 403]:  # unauthorized or forbidden122            print("Need 'write' access token to create a Spaces repo.")123            huggingface_hub.login(add_to_git_credential=False)124            huggingface_hub.add_space_variable(125                space_id, "TRACKIO_DATASET_ID", dataset_id126            )127        else:128            raise ValueError(f"Failed to create Space: {e}")129 130    print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")131    deploy_as_space(space_id, dataset_id)132 133 134def wait_until_space_exists(135    space_id: str,136) -> None:137    """138    Blocks the current thread until the space exists.139    May raise a TimeoutError if this takes quite a while.140 141    Args:142        space_id: The ID of the Space to wait for.143    """144    delay = 1145    for _ in range(10):146        try:147            Client(space_id, verbose=False)148            return149        except (ReadTimeout, ValueError):150            time.sleep(delay)151            delay = min(delay * 2, 30)152    raise TimeoutError("Waiting for space to exist took longer than expected")153 154 155def upload_db_to_space(project: str, space_id: str) -> None:156    """157    Uploads the database of a local Trackio project to a Hugging Face Space.158 159    Args:160        project: The name of the project to upload.161        space_id: The ID of the Space to upload to.162    """163    db_path = SQLiteStorage.get_project_db_path(project)164    client = Client(space_id, verbose=False)165    client.predict(166        api_name="/upload_db_to_space",167        project=project,168        uploaded_db=handle_file(db_path),169        hf_token=huggingface_hub.utils.get_token(),170    )171