Refactor restart_space script to accept space and token as command line arguments

This commit is contained in:
ogabrielluiz 2024-06-05 15:40:30 -03:00
commit bcdc329d1d
2 changed files with 20 additions and 6 deletions

View file

@ -100,6 +100,4 @@ jobs:
- name: Restart HuggingFace Spaces Build
run: |
poetry run python ./scripts/factory_restart_space.py
env:
HUGGINGFACE_API_TOKEN: ${{ secrets.HUGGINGFACE_API_TOKEN }}
poetry run python ./scripts/factory_restart_space.py --space "Langflow/Langflow-Preview" --token ${{ secrets.HUGGINGFACE_API_TOKEN }}

View file

@ -1,4 +1,4 @@
import os
import argparse
from huggingface_hub import HfApi, list_models
from rich import print
@ -6,11 +6,27 @@ from rich import print
# Use root method
models = list_models()
args = argparse.ArgumentParser(description="Restart a space in the Hugging Face Hub.")
args.add_argument("--space", type=str, help="The space to restart.")
args.add_argument("--token", type=str, help="The Hugging Face API token.")
parsed_args = args.parse_args()
space = parsed_args.space
if not space:
print("Please provide a space to restart.")
exit()
if not parsed_args.api_token:
print("Please provide an API token.")
exit()
# Or configure a HfApi client
hf_api = HfApi(
endpoint="https://huggingface.co", # Can be a Private Hub endpoint.
token=os.getenv("HUGGINFACE_API_TOKEN"),
token=parsed_args.token,
)
space_runtime = hf_api.restart_space("Langflow/Langflow-Preview", factory_reboot=True)
space_runtime = hf_api.restart_space(space, factory_reboot=True)
print(space_runtime)