From 00d7a5ce24f435f409013bc63c74d6ba698c16c4 Mon Sep 17 00:00:00 2001
From: Freddy Boulton <alfonsoboulton@gmail.com>
Date: Tue, 11 Apr 2023 19:38:17 -0700
Subject: [PATCH] Python Client: Handle invalid space state (#3822)

* Raise error on status check

* Add BUILDING

* Remove SLEEPING
---
 client/python/gradio_client/client.py | 12 ++++++++++++
 client/python/gradio_client/utils.py  |  8 ++++++++
 client/python/test/test_client.py     |  5 +++++
 3 files changed, 25 insertions(+)

diff --git a/client/python/gradio_client/client.py b/client/python/gradio_client/client.py
index 21155707b4..8f62384b1e 100644
--- a/client/python/gradio_client/client.py
+++ b/client/python/gradio_client/client.py
@@ -83,6 +83,12 @@ class Client:
         self.src = _src
         if self.verbose:
             print(f"Loaded as API: {self.src} ✔")
+        state = self._get_space_state()
+        if state in utils.INVALID_RUNTIME:
+            raise ValueError(
+                f"The current space is in the invalid state: {state}. "
+                "Please contact the owner to fix this."
+            )
 
         self.api_url = urllib.parse.urljoin(self.src, utils.API_URL)
         self.ws_url = urllib.parse.urljoin(
@@ -104,6 +110,12 @@ class Client:
         # Disable telemetry by setting the env variable HF_HUB_DISABLE_TELEMETRY=1
         threading.Thread(target=self._telemetry_thread).start()
 
+    def _get_space_state(self):
+        if not self.space_id:
+            return None
+        api = huggingface_hub.HfApi(token=self.hf_token)
+        return api.get_space_runtime(self.space_id).stage
+
     def predict(
         self,
         *args,
diff --git a/client/python/gradio_client/utils.py b/client/python/gradio_client/utils.py
index 5a8d164c70..80efcf6b1f 100644
--- a/client/python/gradio_client/utils.py
+++ b/client/python/gradio_client/utils.py
@@ -27,6 +27,14 @@ UPLOAD_URL = "/upload"
 RESET_URL = "/reset"
 DUPLICATE_URL = "https://huggingface.co/spaces/{}?duplicate=true"
 STATE_COMPONENT = "state"
+INVALID_RUNTIME = [
+    "NO_APP_FILE",
+    "CONFIG_ERROR",
+    "BUILDING",
+    "BUILD_ERROR",
+    "RUNTIME_ERROR",
+    "PAUSED",
+]
 
 __version__ = (pkgutil.get_data(__name__, "version.txt") or b"").decode("ascii").strip()
 
diff --git a/client/python/test/test_client.py b/client/python/test/test_client.py
index eb79ed7d66..dd9ea4bb22 100644
--- a/client/python/test/test_client.py
+++ b/client/python/test/test_client.py
@@ -19,6 +19,11 @@ HF_TOKEN = "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"  # Intentionally reveali
 
 
 class TestPredictionsFromSpaces:
+    @pytest.mark.flaky
+    def test_raise_error_invalid_state(self):
+        with pytest.raises(ValueError, match="invalid state"):
+            Client("gradio-tests/paused-space")
+
     @pytest.mark.flaky
     def test_numerical_to_label_space(self):
         client = Client("gradio-tests/titanic-survival")