ping and try to wake up agent

app.py

@@ -43,12 +43,42 @@ def initialize_agent(endpoint_uri: str):
         additional_authorized_imports=["pytz"]
     ) 
 
+def wake_up_endpoint(endpoint_uri, max_wait=30):
+    """Poll the endpoint until it responds OK or timeout."""
+    payload = {"inputs": "ping"}  # Adjust to minimal valid input
+    start = time.time()
+
+    while time.time() - start < max_wait:
+        try:
+            print("Pinging endpoint...")
+            response = requests.post(endpoint_uri, json=payload, timeout=5)
+
+            if response.ok:
+                print("✅ Endpoint is awake.")
+                return True
+
+            # If it's a 503 or 504, it's likely still warming up
+            print(f"Response: {response.status_code}, retrying...")
+        except requests.exceptions.RequestException as e:
+            print(f"Exception: {e}, retrying...")
+
+        time.sleep(3)  # fixed short backoff; can be exponential if needed
+
+    print("❌ Timed out waiting for endpoint to wake up.")
+    return False
 
 def main():
     # Load endpoint and check health
 
     # Prechecks 
     endpoint_uri = load_huggingface_model() # Get the URI for the endpoint 
+   
+
+    # Wake it up before health check
+    wake_up_successful = wake_up_endpoint(endpoint_uri)
+    if not wake_up_successful:
+    print("Warning: Could not wake up the endpoint. Continuing anyway...")
+
     is_healthy, status_info = check_model_endpoint(endpoint_uri) # Test the endpoint
     
     if not is_healthy: