adding reties for space to wake up any downstream space

app.py

@@ -41,7 +41,91 @@ lock = threading.Lock()
 # Create an OrderedDict to store clients, limited to 15 entries
 client_cache = OrderedDict()
 MAX_CACHE_SIZE = 15
-default_client=Client("pi19404/ai-worker", hf_token=API_TOKEN)
+
+
+def my_inference_function(client,input_data, output_data,mode, max_length, max_new_tokens, model_size):
+    """
+    The main inference function to process input data and return results.
+    
+    Args:
+        input_data (str or dict): The input data in JSON format.
+        mode (str): The mode of operation ("scoring" or "generative").
+        max_length (int): The maximum length of the input prompt.
+        max_new_tokens (int): The maximum number of new tokens to generate.
+        model_size (str): The size of the model to be used.
+    
+    Returns:
+        str: The output data in JSON format.
+    """
+    with lock:
+        try:
+        
+
+            
+            result = client[0].predict(
+                    input_data=input_data,
+                    output_data=output_data,
+                    mode=mode,
+                    max_length=max_length,
+                    max_new_tokens=max_new_tokens,
+                    model_size=model_size,
+                    api_name="/my_inference_function"
+            )
+            print(result)
+            print("entering return",result)
+            return result  # Pretty-print the JSON
+        except json.JSONDecodeError:
+            return json.dumps({"error": "Invalid JSON input"})
+        except KeyError:
+            return json.dumps({"error": "Missing 'input' key in JSON"})
+        except ValueError as e:
+            return json.dumps({"error": str(e)})
+
+
+def wake_up_space_with_retries(space_url, token, retries=5, wait_time=10):
+    """
+    Attempt to wake up the Hugging Face Space with retries.
+    Retries a number of times in case of a delay due to the Space waking up.
+    
+    :param space_url: The URL of the Hugging Face Space.
+    :param token: The Hugging Face API token.
+    :param retries: Number of retries if the Space is sleeping.
+    :param wait_time: Time to wait between retries (in seconds).
+    """
+    for attempt in range(retries):
+        try:
+            print(f"Attempt {attempt + 1} to wake up the Space...")
+            
+            # Initialize the Gradio Client
+            client = Client(space_url, hf_token=token, timeout=httpx.Timeout(30.0))  # 30-second timeout
+            
+
+            my_inference_function(client,"test input","",scoring,10,10,"2B")
+
+            # Make a prediction or call to wake the Space
+            #result = client.predict("<your_input>")  # Replace with actual inputs
+            print("Space is awake and ready!")
+            return client
+
+        except httpx.ReadTimeout:
+            print(f"Request timed out on attempt {attempt + 1}. Retrying in {wait_time} seconds...")
+            time.sleep(wait_time)
+        
+        except Exception as e:
+            print(f"An error occurred on attempt {attempt + 1}: {e}")
+        
+        # Wait before retrying
+        if attempt < retries - 1:
+            print(f"Waiting for {wait_time} seconds before retrying...")
+
+    print("Space is still not active after multiple attempts.")
+    return None
+
+
+#default_client=Client("pi19404/ai-worker", hf_token=API_TOKEN)
+
+default_client=wake_up_space_with_retries("pi19404/ai-worker",API_TOKEN)
+
 def get_client_for_ip(ip_address,x_ip_token):
     """
     Retrieve or create a client for the given IP address.
@@ -118,43 +202,7 @@ def set_client_for_session(request: gr.Request):
     # The "gradio/text-to-image" space is a ZeroGPU space
     
     
-def my_inference_function(client,input_data, output_data,mode, max_length, max_new_tokens, model_size):
-    """
-    The main inference function to process input data and return results.
-    
-    Args:
-        input_data (str or dict): The input data in JSON format.
-        mode (str): The mode of operation ("scoring" or "generative").
-        max_length (int): The maximum length of the input prompt.
-        max_new_tokens (int): The maximum number of new tokens to generate.
-        model_size (str): The size of the model to be used.
-    
-    Returns:
-        str: The output data in JSON format.
-    """
-    with lock:
-        try:
-        
 
-            
-            result = client.predict(
-                    input_data=input_data,
-                    output_data=output_data,
-                    mode=mode,
-                    max_length=max_length,
-                    max_new_tokens=max_new_tokens,
-                    model_size=model_size,
-                    api_name="/my_inference_function"
-            )
-            print(result)
-            print("entering return",result)
-            return result  # Pretty-print the JSON
-        except json.JSONDecodeError:
-            return json.dumps({"error": "Invalid JSON input"})
-        except KeyError:
-            return json.dumps({"error": "Missing 'input' key in JSON"})
-        except ValueError as e:
-            return json.dumps({"error": str(e)})
 
 with gr.Blocks() as demo:
     """
@@ -166,7 +214,7 @@ with gr.Blocks() as demo:
     - Other UI components (not shown in this snippet).
     - A load event that calls set_client_for_session when the interface is loaded.
     """
-    
+
     gr.Markdown("## LLM Safety Evaluation")
     client = gr.State()
     with gr.Tab("ShieldGemma2"):