fix: improve docker build process and error handling

- Replace r8.im base image with standard Python 3.10 image
- Update runpod package to version 1.7.13
- Add proper error handling for Cog server startup
- Improve documentation and testing setup
- Clean up Dockerfile and build process

Author: TimPietrusky

CHANGES.md

@@ -0,0 +1,23 @@
+# Changes Made
+
+## Dockerfile
+- Replaced r8.im base image with standard Python 3.10 image
+- Updated runpod package to version 1.7.13
+- Added explicit COPY for test_input.json
+- Simplified build process by removing build arguments
+- Added proper cleanup after apt-get operations
+
+## README.md
+- Updated documentation to reflect new build process
+- Removed build arguments section
+- Added section about customizing for specific models
+- Added testing information
+
+## src/handler.py
+- Enhanced error handling for Cog server startup
+- Added alternative method to start Cog server if primary method fails
+- Improved logging
+- Added input validation
+
+## Added Files
+- Ensured test_input.json is properly copied into the Docker image for RunPod SDK automated testing

Dockerfile

@@ -1,26 +1,34 @@
-ARG COG_REPO
-ARG COG_MODEL
-ARG COG_VERSION
-
-FROM r8.im/${COG_REPO}/${COG_MODEL}@sha256:${COG_VERSION}
+# Use a standard Python image as base
+FROM python:3.10-slim
 
 ENV RUNPOD_REQUEST_TIMEOUT=600
 
-# Install necessary packages and Python 3.10
+# Install necessary packages
 RUN apt-get update && apt-get upgrade -y && \
-    apt-get install -y --no-install-recommends software-properties-common curl git openssh-server && \
-    add-apt-repository ppa:deadsnakes/ppa -y && \
-    apt-get update && apt-get install -y --no-install-recommends python3.10 python3.10-dev python3.10-distutils && \
-    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 &&\
-    curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && \
-    python3 get-pip.py
+    apt-get install -y --no-install-recommends \
+    software-properties-common \
+    curl \
+    git \
+    openssh-server \
+    build-essential \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
 
 # Create a virtual environment
 RUN python3 -m venv /opt/venv
 
-# Install runpod within the virtual environment
-RUN /opt/venv/bin/pip install runpod
+# Install Cog and RunPod
+RUN /opt/venv/bin/pip install --upgrade pip && \
+    /opt/venv/bin/pip install cog runpod==1.7.13
+
+# Copy the test input file
+COPY test_input.json /
+
+# Add the handler script
+COPY src/handler.py /rp_handler.py
 
-ADD src/handler.py /rp_handler.py
+# Set the working directory
+WORKDIR /
 
+# Start the handler
 CMD ["/opt/venv/bin/python3", "-u", "/rp_handler.py"]

README.md

@@ -8,12 +8,43 @@ Easily convert cog based images to a runpod serverless worker.
 
 ## Getting Started
 
+This repository provides a template for creating RunPod serverless workers from Cog-based models.
+
+### Building the Docker Image
+
 ```bash
 git clone https://github.com/runpod-workers/cog-worker.git
 
 cd cog-worker/
 
-docker build --tag user/repo:tag --build-arg COG_REPO=user --build-arg COG_MODEL=model_name --build-arg COG_VERSION=model_version .
+# Build the Docker image
+docker build --tag user/repo:tag .
 
 docker push user/repo:tag
 ```
+
+### Customizing for Your Model
+
+To use this worker with your specific Cog model:
+
+1. Install your model's dependencies in the Dockerfile
+2. Create a `predict.py` file with your model's prediction logic
+3. Update the handler.py file if needed to work with your specific model
+
+## RunPod Deployment
+
+After pushing your Docker image, you can deploy it as a RunPod Serverless endpoint through the RunPod platform.
+
+## Testing
+
+The repository includes a test_input.json file that can be used to test your worker:
+
+```json
+{
+  "input": {
+    "prompt": "A beautiful landscape with mountains and a lake"
+  }
+}
+```
+
+You can customize this file to match your model's expected input format.

src/handler.py

@@ -1,6 +1,7 @@
 import time
 import subprocess
 import os
+import sys
 
 import runpod
 import requests
@@ -17,7 +18,17 @@
 
 # ----------------------------- Start API Service ---------------------------- #
 # Call "python -m cog.server.http" in a subprocess to start the API service.
-subprocess.Popen(["python", "-m", "cog.server.http"])
+try:
+    subprocess.Popen([sys.executable, "-m", "cog.server.http"])
+    print("Started Cog server")
+except Exception as e:
+    print(f"Error starting Cog server: {e}")
+    # Try alternative method if the first one fails
+    try:
+        subprocess.Popen(["cog", "server"])
+        print("Started Cog server using alternative method")
+    except Exception as e2:
+        print(f"Error starting Cog server with alternative method: {e2}")
 
 
 # ---------------------------------------------------------------------------- #
@@ -27,6 +38,7 @@ def wait_for_service(url):
     '''
     Check if the service is ready to receive requests.
     '''
+    print(f"Waiting for service at {url}...")
     while True:
         try:
             health = requests.get(url, timeout=120)
@@ -61,6 +73,9 @@ def handler(event):
     '''
     This is the handler function that will be called by the serverless.
     '''
+    # Check if input is provided
+    if not event or "input" not in event:
+        return {"error": "No input provided"}
 
     json = run_inference({"input": event["input"]})
 

test_input.json

@@ -0,0 +1,5 @@
+{
+  "input": {
+    "prompt": "A beautiful landscape with mountains and a lake"
+  }
+}