Cog: The Docker Abstraction That Solves ML Dependency Hell

Hook

If you've ever spent three days debugging why CUDA 11.7 won't work with PyTorch 1.13 and cuDNN 8.5 on Ubuntu 20.04, only to discover you needed CUDA 11.6 all along, Cog was built for you.

Context

Machine learning models are notoriously difficult to deploy. A researcher can train a breakthrough model in a Jupyter notebook, but shipping it to production requires navigating a minefield of incompatible dependencies. CUDA versions must align with cuDNN versions, which must align with PyTorch or TensorFlow versions, which must align with Python versions, which must align with system libraries. Get one combination wrong and you'll spend hours deciphering cryptic error messages about missing shared objects or version mismatches.

Beyond dependencies, there's the HTTP API problem. Most ML models need to be wrapped in a web server for production use, which means writing Flask or FastAPI boilerplate, handling request validation, managing input/output serialization, and maintaining OpenAPI documentation separately from your code. This infrastructure work often takes longer than training the model itself. Cog emerged from this pain point, built by Ben Firshman (creator of Docker Compose) and Andreas Jansson (former Spotify ML infrastructure engineer) who witnessed these problems at scale. It's an opinionated build tool that takes a declarative configuration file and Python prediction code, then generates optimized Docker containers with production-ready HTTP APIs—no Dockerfile or server code required.

Technical Insight

Cog's architecture centers on three components: a Go-based CLI, a YAML-to-Docker transpiler, and a Rust-powered HTTP server. The workflow starts with a cog.yaml file where you declare your model's dependencies at a higher level than raw pip packages. Instead of specifying exact CUDA and cuDNN versions, you declare what you actually care about—your Python version and ML framework.

Here's a minimal configuration:

build:
  gpu: true
  python_version: "3.10"
  python_packages:
    - "torch==2.0.0"
    - "transformers==4.28.0"
    - "Pillow==9.5.0"
predict: "predict.py:Predictor"

From this, Cog automatically determines that PyTorch 2.0.0 requires CUDA 11.7 or 11.8, selects compatible cuDNN versions, chooses an appropriate base image (nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04), and generates a Dockerfile with optimized layer caching. It handles the dependency resolution that typically requires consulting compatibility matrices across multiple documentation sites.

The prediction interface uses a Python class with type-annotated methods. Cog introspects these type hints to generate OpenAPI schemas and handle validation automatically:

from cog import BasePredictor, Input, Path
import torch

class Predictor(BasePredictor):
    def setup(self):
        """Load model weights once at container startup"""
        self.model = torch.load("./weights.pth")
        self.model.eval()
    
    def predict(
        self,
        image: Path = Input(description="Input image"),
        scale: float = Input(
            description="Scaling factor",
            default=1.0,
            ge=0.1,
            le=10.0
        )
    ) -> Path:
        """Run inference and return processed image"""
        img = Image.open(image)
        output = self.model(img, scale=scale)
        output_path = Path("/tmp/output.png")
        output.save(output_path)
        return output_path

The Input type provides declarative validation—ge and le constraints become OpenAPI schema properties and are enforced at the HTTP layer before your prediction code runs. Supported types include primitive values, files (Path), and structured data, all automatically serialized and deserialized.

When you run cog build, the Go CLI parses your configuration, queries an internal compatibility database, and generates a Dockerfile with aggressive layer caching strategies. Dependencies are installed in a specific order to maximize Docker layer reuse across different models. The generated image includes a Rust-based HTTP server (built with Axum) that's significantly faster than typical Python WSGI servers. This server dynamically creates REST endpoints by introspecting your Predictor class at runtime.

For local development, cog predict executes your model without building a container, using your local Python environment but validating that your code matches the containerized interface. This tight feedback loop helps catch issues before committing to a multi-gigabyte Docker build. Once containerized, the same image can run locally (docker run), on Replicate's platform, or any container orchestrator, because it's just a standard Docker image with an HTTP server listening on port 5000.

The separation of concerns is clever: your prediction code knows nothing about HTTP, Docker, or serialization. Cog handles the infrastructure layer, letting you focus on model logic. This is particularly valuable for research teams where Docker expertise is rare, but it also benefits experienced engineers by eliminating repetitive server scaffolding across multiple models.

Gotcha

Cog's opinionated nature is both its strength and limitation. The abstraction works beautifully for the 80% use case—single Python models with standard dependencies—but breaks down at the edges. If you need to customize the Dockerfile beyond what cog.yaml supports (like installing system packages with complex configurations, using multi-stage builds, or building from alternative base images), you'll find yourself fighting the tool. While you can add run commands for arbitrary shell execution during build, this escape hatch undermines the declarative benefits and makes debugging harder since you're now reasoning about both Cog's generated Dockerfile and your custom modifications.

The Replicate platform integration also creates subtle coupling. While Docker images are theoretically portable, certain features are optimized for Replicate's infrastructure. The HTTP server expects specific environment variables and file paths that may require translation for other deployment targets. Documentation skews toward the happy path of deploying to Replicate, leaving self-hosted users to reverse-engineer implementation details. Multi-model serving isn't supported—each container runs one Predictor class, so if you need to ensemble multiple models or serve variants, you're managing multiple containers and a routing layer yourself. For teams already invested in Kubernetes-native tools like KServe or Seldon, Cog adds another abstraction layer that may conflict with existing patterns.

Verdict

Use if: You're shipping Python ML models to production and spending more time on Docker/dependency issues than model development, your team lacks deep DevOps expertise but needs containerized deployment, you're building multiple models and want standardized interfaces, or you're deploying to Replicate and want the tightest integration. Cog dramatically reduces time-to-deployment for typical PyTorch/TensorFlow workloads. Skip if: You need fine-grained Dockerfile control for complex system dependencies, you're working with non-Python ML runtimes (C++, Rust, Java), you require multi-model containers or advanced serving features like A/B testing and canary deployments, or you've already built robust ML infrastructure and don't want another abstraction layer. For teams starting fresh or researchers crossing into production, Cog's tradeoff of flexibility for simplicity is usually worth it.