Support factory reboot when restarting a Space (#1586)

* Implement factory reboot for Spaces * fix
huggingface · Aug 11, 2023 · 55f1bcd · 55f1bcd
1 parent dce3839
commit 55f1bcd
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/src/huggingface_hub/hf_api.py b/src/huggingface_hub/hf_api.py
@@ -5172,7 +5172,9 @@ def pause_space(self, repo_id: str, *, token: Optional[str] = None) -> SpaceRunt
         return SpaceRuntime(r.json())
 
     @validate_hf_hub_args
-    def restart_space(self, repo_id: str, *, token: Optional[str] = None) -> SpaceRuntime:
+    def restart_space(
+        self, repo_id: str, *, token: Optional[str] = None, factory_reboot: bool = False
+    ) -> SpaceRuntime:
         """Restart your Space.
 
         This is the only way to programmatically restart a Space if you've put it on Pause (see [`pause_space`]). You
@@ -5186,6 +5188,8 @@ def restart_space(self, repo_id: str, *, token: Optional[str] = None) -> SpaceRu
                 ID of the Space to restart. Example: `"Salesforce/BLIP2"`.
             token (`str`, *optional*):
                 Hugging Face token. Will default to the locally saved token if not provided.
+            factory_reboot (`bool`, *optional*):
+                If `True`, the Space will be rebuilt from scratch without caching any requirements.
 
         Returns:
             [`SpaceRuntime`]: Runtime information about your Space.
@@ -5201,8 +5205,11 @@ def restart_space(self, repo_id: str, *, token: Optional[str] = None) -> SpaceRu
                 If your Space is a static Space. Static Spaces are always running and never billed. If you want to hide
                 a static Space, you can set it to private.
         """
+        params = {}
+        if factory_reboot:
+            params["factory"] = "true"
         r = get_session().post(
-            f"{self.endpoint}/api/spaces/{repo_id}/restart", headers=self._build_hf_headers(token=token)
+            f"{self.endpoint}/api/spaces/{repo_id}/restart", headers=self._build_hf_headers(token=token), params=params
         )
         hf_raise_for_status(r)
         return SpaceRuntime(r.json())

diff --git a/tests/test_hf_api.py b/tests/test_hf_api.py
@@ -2691,6 +2691,14 @@ def test_delete_space_storage(self) -> None:
         )
         assert runtime.storage is None
 
+    def test_restart_space_factory_reboot(self) -> None:
+        self.api.restart_space(self.repo_id, factory_reboot=True)
+        self.post_mock.assert_called_once_with(
+            f"{self.api.endpoint}/api/spaces/{self.repo_id}/restart",
+            headers=self.api._build_hf_headers(),
+            params={"factory": "true"},
+        )
+
 
 class ListGitRefsTest(unittest.TestCase):
     @classmethod