File size: 4,282 Bytes
64d0d8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
import io
import os
import time
from importlib.resources import files
from pathlib import Path
import gradio
import huggingface_hub
from gradio_client import Client, handle_file
from httpx import ReadTimeout
from huggingface_hub.errors import RepositoryNotFoundError
from trackio.sqlite_storage import SQLiteStorage
SPACE_URL = "https://huggingface.co/spaces/{space_id}"
def deploy_as_space(
space_id: str,
dataset_id: str | None = None,
):
if (
os.getenv("SYSTEM") == "spaces"
): # in case a repo with this function is uploaded to spaces
return
trackio_path = files("trackio")
hf_api = huggingface_hub.HfApi()
whoami = None
login = False
try:
whoami = hf_api.whoami()
if whoami["auth"]["accessToken"]["role"] != "write":
login = True
except OSError:
login = True
if login:
print("Need 'write' access token to create a Spaces repo.")
huggingface_hub.login(add_to_git_credential=False)
whoami = hf_api.whoami()
huggingface_hub.create_repo(
space_id,
space_sdk="gradio",
repo_type="space",
exist_ok=True,
)
with open(Path(trackio_path, "README.md"), "r") as f:
readme_content = f.read()
readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)
readme_buffer = io.BytesIO(readme_content.encode("utf-8"))
hf_api.upload_file(
path_or_fileobj=readme_buffer,
path_in_repo="README.md",
repo_id=space_id,
repo_type="space",
)
huggingface_hub.utils.disable_progress_bars()
hf_api.upload_folder(
repo_id=space_id,
repo_type="space",
folder_path=trackio_path,
ignore_patterns=["README.md"],
)
hf_token = huggingface_hub.utils.get_token()
if hf_token is not None:
huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)
if dataset_id is not None:
huggingface_hub.add_space_variable(space_id, "TRACKIO_DATASET_ID", dataset_id)
def create_space_if_not_exists(
space_id: str,
dataset_id: str | None = None,
) -> None:
"""
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.
Args:
space_id: The ID of the Space to create.
dataset_id: The ID of the Dataset to add to the Space.
"""
if "/" not in space_id:
raise ValueError(
f"Invalid space ID: {space_id}. Must be in the format: username/reponame or orgname/reponame."
)
if dataset_id is not None and "/" not in dataset_id:
raise ValueError(
f"Invalid dataset ID: {dataset_id}. Must be in the format: username/datasetname or orgname/datasetname."
)
try:
huggingface_hub.repo_info(space_id, repo_type="space")
print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")
if dataset_id is not None:
huggingface_hub.add_space_variable(
space_id, "TRACKIO_DATASET_ID", dataset_id
)
return
except RepositoryNotFoundError:
pass
print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")
deploy_as_space(space_id, dataset_id)
client = None
for _ in range(30):
try:
client = Client(space_id, verbose=False)
if client:
break
except ReadTimeout:
print("* Space is not yet ready. Waiting 5 seconds...")
time.sleep(5)
except ValueError as e:
print(f"* Space gave error {e}. Trying again in 5 seconds...")
time.sleep(5)
def upload_db_to_space(project: str, space_id: str) -> None:
"""
Uploads the database of a local Trackio project to a Hugging Face Space.
Args:
project: The name of the project to upload.
space_id: The ID of the Space to upload to.
"""
db_path = SQLiteStorage.get_project_db_path(project)
client = Client(space_id, verbose=False)
client.predict(
api_name="/upload_db_to_space",
project=project,
uploaded_db=handle_file(db_path),
hf_token=huggingface_hub.utils.get_token(),
)
|