cog: The Death of the Hand-Written ML Dockerfile

Software engineers can take these models and run them with one line of code, without having to understand all the internals about how the model works, and without having to set up GPUs

Ben Firshman, Co-founder of Replicate · Interview

The Containerization Trap

Deploying machine learning models is notoriously painful. A standard Dockerfile for a web application might be ten lines long, but an ML container is a sprawling, fragile beast. Engineers must navigate a minefield of mismatched CUDA versions, specific cuDNN binaries, and PyTorch releases that only play nice under exact, undocumented conditions.

This is "CUDA Hell." When a single minor driver update can render a production container unbootable, Docker ceases to be a helpful abstraction and becomes a low-level primitive that ML engineers are forced to micromanage.

Infrastructure from Code

Cog was built to solve this exact problem. Created by Ben Firshman—who previously co-created Fig, the tool that eventually became `docker-compose`—Cog applies the same philosophy of declarative simplicity to the uniquely hostile environment of machine learning.

Instead of writing a Dockerfile, users define their environment in a straightforward `cog.yaml` file. But this isn't just syntactic sugar; it's a completely different paradigm.

The Expert System in the CLI

The magic of Cog lies in its Go-based orchestration layer. When a user runs `cog build`, the CLI doesn't just parse the YAML; it consults an internal expert system. The file `pkg/config/config.go` contains a hardcoded matrix of dependency rules, essentially acting as an automated DevOps engineer.

If you request PyTorch 2.0, Cog knows exactly which Nvidia base image, CUDA driver, and cuDNN version are required. It automatically generates the complex Dockerfile instructions needed to build a stable, reproducible image, entirely shielding the user from the underlying dependency graph.

Schema by Introspection

Cog also automates the tedious process of building a REST API. Instead of requiring users to write Flask or FastAPI wrappers, Cog uses static analysis to inspect the user's `predict.py` file. By leveraging `tree-sitter` within the Go CLI, Cog can parse the Python AST (Abstract Syntax Tree) without actually executing the code.

It reads the type hints on the `predict` function and automatically generates a comprehensive OpenAPI schema. In this model, infrastructure and API definitions are treated as a direct side effect of the code itself.

# predict.py
from cog import BasePredictor, Input, Path

class Predictor(BasePredictor):
    def predict(
        self,
        image: Path = Input(description="Input image"),
        scale: float = Input(default=1.5, description="Scaling factor")
    ) -> Path:
        # Model inference logic here
        return output_path

The Rust-to-Python Bridge

While the CLI is written in Go, the runtime engine inside the container—known as `coglet`—is built in Rust. This architectural shift addresses the inherent limitations of serving ML models with pure Python web servers, which struggle with concurrency due to the Global Interpreter Lock (GIL).

The Rust orchestrator manages a `PermitPool` to handle incoming HTTP requests. Instead of relying on standard HTTP internal proxies, `coglet` uses Unix Domain Sockets to stream binary data and logs directly to the Python subprocesses. This "Slot Socket" architecture provides high-throughput I/O and precise lifecycle control, ensuring that heavy ML workloads don't crash the serving layer.

The Deployment Landscape

Cog occupies a specific niche in the ML deployment ecosystem. While BentoML offers a broader, more framework-heavy approach with extensive integrations, and raw Docker provides ultimate flexibility at the cost of immense manual effort, Cog is highly opinionated.

It is optimized for the Replicate platform but outputs standard Docker containers that can run anywhere. By codifying years of painful ML DevOps experience into a single tool, Cog allows engineers to focus on their models rather than their infrastructure.