umair894/quickstart-trackio
0
1import importlib.metadata2import io3import os4import time5from importlib.resources import files6from pathlib import Path7 8import gradio9import huggingface_hub10from gradio_client import Client, handle_file11from httpx import ReadTimeout12from huggingface_hub.errors import RepositoryNotFoundError13from requests import HTTPError14 15import trackio16from trackio.sqlite_storage import SQLiteStorage17 18SPACE_HOST_URL = "https://{user_name}-{space_name}.hf.space/"19SPACE_URL = "https://huggingface.co/spaces/{space_id}"20 21 22def _is_trackio_installed_from_source() -> bool:23 """Check if trackio is installed from source/editable install vs PyPI."""24 try:25 trackio_file = trackio.__file__26 if "site-packages" not in trackio_file:27 return True28 29 dist = importlib.metadata.distribution("trackio")30 if dist.files:31 files = list(dist.files)32 has_pth = any(".pth" in str(f) for f in files)33 if has_pth:34 return True35 36 return False37 except (38 AttributeError,39 importlib.metadata.PackageNotFoundError,40 importlib.metadata.MetadataError,41 ValueError,42 TypeError,43 ):44 return True45 46 47def deploy_as_space(48 space_id: str,49 space_storage: huggingface_hub.SpaceStorage | None = None,50 dataset_id: str | None = None,51 private: bool | None = None,52):53 if (54 os.getenv("SYSTEM") == "spaces"55 ): # in case a repo with this function is uploaded to spaces56 return57 58 trackio_path = files("trackio")59 60 hf_api = huggingface_hub.HfApi()61 62 try:63 huggingface_hub.create_repo(64 space_id,65 private=private,66 space_sdk="gradio",67 space_storage=space_storage,68 repo_type="space",69 exist_ok=True,70 )71 except HTTPError as e:72 if e.response.status_code in [401, 403]: # unauthorized or forbidden73 print("Need 'write' access token to create a Spaces repo.")74 huggingface_hub.login(add_to_git_credential=False)75 huggingface_hub.create_repo(76 space_id,77 private=private,78 space_sdk="gradio",79 space_storage=space_storage,80 repo_type="space",81 exist_ok=True,82 )83 else:84 raise ValueError(f"Failed to create Space: {e}")85 86 with open(Path(trackio_path, "README.md"), "r") as f:87 readme_content = f.read()88 readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)89 readme_buffer = io.BytesIO(readme_content.encode("utf-8"))90 hf_api.upload_file(91 path_or_fileobj=readme_buffer,92 path_in_repo="README.md",93 repo_id=space_id,94 repo_type="space",95 )96 97 # We can assume pandas, gradio, and huggingface-hub are already installed in a Gradio Space.98 # Make sure necessary dependencies are installed by creating a requirements.txt.99 is_source_install = _is_trackio_installed_from_source()100 101 if is_source_install:102 requirements_content = """pyarrow>=21.0103plotly>=6.0.0,<7.0.0"""104 else:105 requirements_content = f"""pyarrow>=21.0106trackio=={trackio.__version__}107plotly>=6.0.0,<7.0.0"""108 109 requirements_buffer = io.BytesIO(requirements_content.encode("utf-8"))110 hf_api.upload_file(111 path_or_fileobj=requirements_buffer,112 path_in_repo="requirements.txt",113 repo_id=space_id,114 repo_type="space",115 )116 117 huggingface_hub.utils.disable_progress_bars()118 119 if is_source_install:120 hf_api.upload_folder(121 repo_id=space_id,122 repo_type="space",123 folder_path=trackio_path,124 ignore_patterns=["README.md"],125 )126 else:127 app_file_content = """import trackio128trackio.show()"""129 app_file_buffer = io.BytesIO(app_file_content.encode("utf-8"))130 hf_api.upload_file(131 path_or_fileobj=app_file_buffer,132 path_in_repo="ui/main.py",133 repo_id=space_id,134 repo_type="space",135 )136 137 if hf_token := huggingface_hub.utils.get_token():138 huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)139 if dataset_id is not None:140 huggingface_hub.add_space_variable(space_id, "TRACKIO_DATASET_ID", dataset_id)141 142 if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):143 huggingface_hub.add_space_variable(144 space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url145 )146 if logo_dark_url := os.environ.get("TRACKIO_LOGO_DARK_URL"):147 huggingface_hub.add_space_variable(148 space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url149 )150 151 if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):152 huggingface_hub.add_space_variable(space_id, "TRACKIO_PLOT_ORDER", plot_order)153 154 if theme := os.environ.get("TRACKIO_THEME"):155 huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)156 157 158def create_space_if_not_exists(159 space_id: str,160 space_storage: huggingface_hub.SpaceStorage | None = None,161 dataset_id: str | None = None,162 private: bool | None = None,163) -> None:164 """165 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.166 167 Args:168 space_id: The ID of the Space to create.169 dataset_id: The ID of the Dataset to add to the Space.170 private: Whether to make the Space private. If None (default), the repo will be171 public unless the organization's default is private. This value is ignored if172 the repo already exists.173 """174 if "/" not in space_id:175 raise ValueError(176 f"Invalid space ID: {space_id}. Must be in the format: username/reponame or orgname/reponame."177 )178 if dataset_id is not None and "/" not in dataset_id:179 raise ValueError(180 f"Invalid dataset ID: {dataset_id}. Must be in the format: username/datasetname or orgname/datasetname."181 )182 try:183 huggingface_hub.repo_info(space_id, repo_type="space")184 print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")185 if dataset_id is not None:186 huggingface_hub.add_space_variable(187 space_id, "TRACKIO_DATASET_ID", dataset_id188 )189 if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):190 huggingface_hub.add_space_variable(191 space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url192 )193 if logo_dark_url := os.environ.get("TRACKIO_LOGO_DARK_URL"):194 huggingface_hub.add_space_variable(195 space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url196 )197 198 if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):199 huggingface_hub.add_space_variable(200 space_id, "TRACKIO_PLOT_ORDER", plot_order201 )202 203 if theme := os.environ.get("TRACKIO_THEME"):204 huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)205 return206 except RepositoryNotFoundError:207 pass208 except HTTPError as e:209 if e.response.status_code in [401, 403]: # unauthorized or forbidden210 print("Need 'write' access token to create a Spaces repo.")211 huggingface_hub.login(add_to_git_credential=False)212 huggingface_hub.add_space_variable(213 space_id, "TRACKIO_DATASET_ID", dataset_id214 )215 else:216 raise ValueError(f"Failed to create Space: {e}")217 218 print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")219 deploy_as_space(space_id, space_storage, dataset_id, private)220 221 222def wait_until_space_exists(223 space_id: str,224) -> None:225 """226 Blocks the current thread until the space exists.227 May raise a TimeoutError if this takes quite a while.228 229 Args:230 space_id: The ID of the Space to wait for.231 """232 delay = 1233 for _ in range(10):234 try:235 Client(space_id, verbose=False)236 return237 except (ReadTimeout, ValueError):238 time.sleep(delay)239 delay = min(delay * 2, 30)240 raise TimeoutError("Waiting for space to exist took longer than expected")241 242 243def upload_db_to_space(project: str, space_id: str) -> None:244 """245 Uploads the database of a local Trackio project to a Hugging Face Space.246 247 Args:248 project: The name of the project to upload.249 space_id: The ID of the Space to upload to.250 """251 db_path = SQLiteStorage.get_project_db_path(project)252 client = Client(space_id, verbose=False)253 client.predict(254 api_name="/upload_db_to_space",255 project=project,256 uploaded_db=handle_file(db_path),257 hf_token=huggingface_hub.utils.get_token(),258 )259