Merge pull request #5 from therealadityashankar/main

added RUNPOD_REQUEST_TIMEOUT, to env - to allow for more control over the timeout for the user

Author: deanq

Dockerfile

@@ -4,6 +4,8 @@ ARG COG_VERSION
 
 FROM r8.im/${COG_REPO}/${COG_MODEL}@sha256:${COG_VERSION}
 
+ENV RUNPOD_REQUEST_TIMEOUT=600
+
 # Install necessary packages and Python 3.10
 RUN apt-get update && apt-get upgrade -y && \
     apt-get install -y --no-install-recommends software-properties-common curl git openssh-server && \
@@ -21,4 +23,4 @@ RUN /opt/venv/bin/pip install runpod
 
 ADD src/handler.py /rp_handler.py
 
-CMD ["/opt/venv/bin/python3", "-u", "/rp_handler.py"]
+CMD ["/opt/venv/bin/python3", "-u", "/rp_handler.py"]

src/handler.py

@@ -1,10 +1,13 @@
 import time
 import subprocess
+import os
 
 import runpod
 import requests
 from requests.adapters import HTTPAdapter, Retry
 
+TIMEOUT = int(os.environ.get("RUNPOD_REQUEST_TIMEOUT", "600"))
+
 LOCAL_URL = "http://127.0.0.1:5000"
 
 cog_session = requests.Session()
@@ -46,13 +49,14 @@ def run_inference(inference_request):
     Run inference on a request.
     '''
     response = cog_session.post(url=f'{LOCAL_URL}/predictions',
-                                json=inference_request, timeout=600)
+                                json=inference_request, timeout=TIMEOUT)
+
+    if response.status_code != 200:
+        print("Request failed - reason :", response.status_code, response.text)
+
     return response.json()
 
 
-# ---------------------------------------------------------------------------- #
-#                                RunPod Handler                                #
-# ---------------------------------------------------------------------------- #
 def handler(event):
     '''
     This is the handler function that will be called by the serverless.
@@ -68,4 +72,4 @@ def handler(event):
 
     print("Cog API Service is ready. Starting RunPod serverless handler...")
 
-    runpod.serverless.start({"handler": handler})
+    runpod.serverless.start({"handler": handler})