Shivam249/function-gemma
0
1import importlib.metadata2import io3import os4import threading5import time6from importlib.resources import files7from pathlib import Path8 9import gradio10import huggingface_hub11from gradio_client import Client, handle_file12from httpx import ReadTimeout13from huggingface_hub.errors import HfHubHTTPError, RepositoryNotFoundError14 15import trackio16from trackio.sqlite_storage import SQLiteStorage17from trackio.utils import get_or_create_project_hash, preprocess_space_and_dataset_ids18 19SPACE_HOST_URL = "https://{user_name}-{space_name}.hf.space/"20SPACE_URL = "https://huggingface.co/spaces/{space_id}"21 22 23def _is_trackio_installed_from_source() -> bool:24 """Check if trackio is installed from source/editable install vs PyPI."""25 try:26 trackio_file = trackio.__file__27 if "site-packages" not in trackio_file:28 return True29 30 dist = importlib.metadata.distribution("trackio")31 if dist.files:32 files = list(dist.files)33 has_pth = any(".pth" in str(f) for f in files)34 if has_pth:35 return True36 37 return False38 except (39 AttributeError,40 importlib.metadata.PackageNotFoundError,41 importlib.metadata.MetadataError,42 ValueError,43 TypeError,44 ):45 return True46 47 48def deploy_as_space(49 space_id: str,50 space_storage: huggingface_hub.SpaceStorage | None = None,51 dataset_id: str | None = None,52 private: bool | None = None,53):54 if (55 os.getenv("SYSTEM") == "spaces"56 ): # in case a repo with this function is uploaded to spaces57 return58 59 trackio_path = files("trackio")60 61 hf_api = huggingface_hub.HfApi()62 63 try:64 huggingface_hub.create_repo(65 space_id,66 private=private,67 space_sdk="gradio",68 space_storage=space_storage,69 repo_type="space",70 exist_ok=True,71 )72 except HfHubHTTPError as e:73 if e.response.status_code in [401, 403]: # unauthorized or forbidden74 print("Need 'write' access token to create a Spaces repo.")75 huggingface_hub.login(add_to_git_credential=False)76 huggingface_hub.create_repo(77 space_id,78 private=private,79 space_sdk="gradio",80 space_storage=space_storage,81 repo_type="space",82 exist_ok=True,83 )84 else:85 raise ValueError(f"Failed to create Space: {e}")86 87 with open(Path(trackio_path, "README.md"), "r") as f:88 readme_content = f.read()89 readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)90 readme_buffer = io.BytesIO(readme_content.encode("utf-8"))91 hf_api.upload_file(92 path_or_fileobj=readme_buffer,93 path_in_repo="README.md",94 repo_id=space_id,95 repo_type="space",96 )97 98 # We can assume pandas, gradio, and huggingface-hub are already installed in a Gradio Space.99 # Make sure necessary dependencies are installed by creating a requirements.txt.100 is_source_install = _is_trackio_installed_from_source()101 102 if is_source_install:103 requirements_content = """pyarrow>=21.0104plotly>=6.0.0,<7.0.0"""105 else:106 requirements_content = f"""pyarrow>=21.0107trackio=={trackio.__version__}108plotly>=6.0.0,<7.0.0"""109 110 requirements_buffer = io.BytesIO(requirements_content.encode("utf-8"))111 hf_api.upload_file(112 path_or_fileobj=requirements_buffer,113 path_in_repo="requirements.txt",114 repo_id=space_id,115 repo_type="space",116 )117 118 huggingface_hub.utils.disable_progress_bars()119 120 if is_source_install:121 hf_api.upload_folder(122 repo_id=space_id,123 repo_type="space",124 folder_path=trackio_path,125 ignore_patterns=["README.md"],126 )127 else:128 app_file_content = """import trackio129trackio.show()"""130 app_file_buffer = io.BytesIO(app_file_content.encode("utf-8"))131 hf_api.upload_file(132 path_or_fileobj=app_file_buffer,133 path_in_repo="ui/main.py",134 repo_id=space_id,135 repo_type="space",136 )137 138 if hf_token := huggingface_hub.utils.get_token():139 huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)140 if dataset_id is not None:141 huggingface_hub.add_space_variable(space_id, "TRACKIO_DATASET_ID", dataset_id)142 143 if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):144 huggingface_hub.add_space_variable(145 space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url146 )147 if logo_dark_url := os.environ.get("TRACKIO_LOGO_DARK_URL"):148 huggingface_hub.add_space_variable(149 space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url150 )151 152 if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):153 huggingface_hub.add_space_variable(space_id, "TRACKIO_PLOT_ORDER", plot_order)154 155 if theme := os.environ.get("TRACKIO_THEME"):156 huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)157 158 huggingface_hub.add_space_variable(space_id, "GRADIO_MCP_SERVER", "True")159 160 161def create_space_if_not_exists(162 space_id: str,163 space_storage: huggingface_hub.SpaceStorage | None = None,164 dataset_id: str | None = None,165 private: bool | None = None,166) -> None:167 """168 Creates a new Hugging Face Space if it does not exist.169 170 Args:171 space_id (`str`):172 The ID of the Space to create.173 space_storage ([`~huggingface_hub.SpaceStorage`], *optional*):174 Choice of persistent storage tier for the Space.175 dataset_id (`str`, *optional*):176 The ID of the Dataset to add to the Space as a space variable.177 private (`bool`, *optional*):178 Whether to make the Space private. If `None` (default), the repo will be179 public unless the organization's default is private. This value is ignored180 if the repo already exists.181 """182 if "/" not in space_id:183 raise ValueError(184 f"Invalid space ID: {space_id}. Must be in the format: username/reponame or orgname/reponame."185 )186 if dataset_id is not None and "/" not in dataset_id:187 raise ValueError(188 f"Invalid dataset ID: {dataset_id}. Must be in the format: username/datasetname or orgname/datasetname."189 )190 try:191 huggingface_hub.repo_info(space_id, repo_type="space")192 print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")193 if dataset_id is not None:194 huggingface_hub.add_space_variable(195 space_id, "TRACKIO_DATASET_ID", dataset_id196 )197 if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):198 huggingface_hub.add_space_variable(199 space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url200 )201 if logo_dark_url := os.environ.get("TRACKIO_LOGO_DARK_URL"):202 huggingface_hub.add_space_variable(203 space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url204 )205 206 if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):207 huggingface_hub.add_space_variable(208 space_id, "TRACKIO_PLOT_ORDER", plot_order209 )210 211 if theme := os.environ.get("TRACKIO_THEME"):212 huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)213 return214 except RepositoryNotFoundError:215 pass216 except HfHubHTTPError as e:217 if e.response.status_code in [401, 403]: # unauthorized or forbidden218 print("Need 'write' access token to create a Spaces repo.")219 huggingface_hub.login(add_to_git_credential=False)220 huggingface_hub.add_space_variable(221 space_id, "TRACKIO_DATASET_ID", dataset_id222 )223 else:224 raise ValueError(f"Failed to create Space: {e}")225 226 print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")227 deploy_as_space(space_id, space_storage, dataset_id, private)228 229 230def wait_until_space_exists(231 space_id: str,232) -> None:233 """234 Blocks the current thread until the Space exists.235 236 Args:237 space_id (`str`):238 The ID of the Space to wait for.239 240 Raises:241 `TimeoutError`: If waiting for the Space takes longer than expected.242 """243 hf_api = huggingface_hub.HfApi()244 delay = 1245 for _ in range(30):246 try:247 hf_api.space_info(space_id)248 return249 except (huggingface_hub.utils.HfHubHTTPError, ReadTimeout):250 time.sleep(delay)251 delay = min(delay * 2, 60)252 raise TimeoutError("Waiting for space to exist took longer than expected")253 254 255def upload_db_to_space(project: str, space_id: str, force: bool = False) -> None:256 """257 Uploads the database of a local Trackio project to a Hugging Face Space.258 259 This uses the Gradio Client to upload since we do not want to trigger a new build of260 the Space, which would happen if we used `huggingface_hub.upload_file`.261 262 Args:263 project (`str`):264 The name of the project to upload.265 space_id (`str`):266 The ID of the Space to upload to.267 force (`bool`, *optional*, defaults to `False`):268 If `True`, overwrites the existing database without prompting. If `False`,269 prompts for confirmation.270 """271 db_path = SQLiteStorage.get_project_db_path(project)272 client = Client(space_id, verbose=False, httpx_kwargs={"timeout": 90})273 274 if not force:275 try:276 existing_projects = client.predict(api_name="/get_all_projects")277 if project in existing_projects:278 response = input(279 f"Database for project '{project}' already exists on Space '{space_id}'. "280 f"Overwrite it? (y/N): "281 )282 if response.lower() not in ["y", "yes"]:283 print("* Upload cancelled.")284 return285 except Exception as e:286 print(f"* Warning: Could not check if project exists on Space: {e}")287 print("* Proceeding with upload...")288 289 client.predict(290 api_name="/upload_db_to_space",291 project=project,292 uploaded_db=handle_file(db_path),293 hf_token=huggingface_hub.utils.get_token(),294 )295 296 297def sync(298 project: str,299 space_id: str | None = None,300 private: bool | None = None,301 force: bool = False,302 run_in_background: bool = False,303) -> str:304 """305 Syncs a local Trackio project's database to a Hugging Face Space.306 If the Space does not exist, it will be created.307 308 Args:309 project (`str`): The name of the project to upload.310 space_id (`str`, *optional*): The ID of the Space to upload to (e.g., `"username/space_id"`).311 If not provided, a random space_id (e.g. "username/project-2ac3z2aA") will be used.312 private (`bool`, *optional*):313 Whether to make the Space private. If None (default), the repo will be314 public unless the organization's default is private. This value is ignored315 if the repo already exists.316 force (`bool`, *optional*, defaults to `False`):317 If `True`, overwrite the existing database without prompting for confirmation.318 If `False`, prompt the user before overwriting an existing database.319 run_in_background (`bool`, *optional*, defaults to `False`):320 If `True`, the Space creation and database upload will be run in a background thread.321 If `False`, all the steps will be run synchronously.322 Returns:323 `str`: The Space ID of the synced project.324 """325 if space_id is None:326 space_id = f"{project}-{get_or_create_project_hash(project)}"327 space_id, _ = preprocess_space_and_dataset_ids(space_id, None)328 329 def space_creation_and_upload(330 space_id: str, private: bool | None = None, force: bool = False331 ):332 print(333 f"* Syncing local Trackio project to: {SPACE_URL.format(space_id=space_id)} (please wait...)"334 )335 create_space_if_not_exists(space_id, private=private)336 wait_until_space_exists(space_id)337 upload_db_to_space(project, space_id, force=force)338 print(f"* Synced successfully to space: {SPACE_URL.format(space_id=space_id)}")339 340 if run_in_background:341 threading.Thread(342 target=space_creation_and_upload, args=(space_id, private, force)343 ).start()344 else:345 space_creation_and_upload(space_id, private, force)346 return space_id347