Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions CHANGES.md
Original file line number Diff line number Diff line change
@@ -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
38 changes: 23 additions & 15 deletions Dockerfile
Original file line number Diff line number Diff line change
@@ -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"]
33 changes: 32 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
17 changes: 16 additions & 1 deletion src/handler.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import time
import subprocess
import os
import sys

import runpod
import requests
Expand All @@ -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}")


# ---------------------------------------------------------------------------- #
Expand All @@ -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)
Expand Down Expand Up @@ -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"]})

Expand Down
5 changes: 5 additions & 0 deletions test_input.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
{
"input": {
"prompt": "A beautiful landscape with mountains and a lake"
}
}
Loading