DOCSEARCHbyAlgolia
Submit Docs

Cog Documentation

repository·main·Indexed 27 days ago

https://github.com/replicate/cog

An open-source tool for packaging machine learning models into production-ready Docker containers. It includes Coglet, a Rust-based prediction server that provides process isolation, concurrent slot management, and high-performance IPC via Unix domain sockets. The system features coglet-python PyO3 bindings to bridge Python predictor classes to the Rust runtime, supporting both synchronous and asynchronous prediction handlers with native cancellation and routed logging.

Tokens
54.5K
Snippets
156
Records
302
Agent score
94%

What's inside Cog

  1. Overview of Coglet Rust Runtime

    main
    Coglet is the Rust-based prediction server that powers Cog's subprocess isolation model. It provides process isolation, concurrent slot management, and high-performance Inter-Process Communication (IPC) for running machine learning predictions. It utilizes a parent process (containing an HTTP server and an orchestrator) and a worker subprocess (containing the Python runtime and a Tokio runtime).

  3. Overview of coglet

    main
    coglet is a core Rust library designed for the coglet prediction server. It is written in pure Rust and contains no Python dependencies. Python bindings for this library are provided separately in the coglet-python package. It provides the architecture for a parent-side orchestrator and a child-side worker to manage machine learning model predictions via IPC (Unix sockets) and HTTP.

  4. Understand Managed Weights OCI Format

    main

    Managed weights are named sets of files (model weights, configs, tokenizers, etc.) stored as OCI artifacts in a container registry. They are delivered independently from the model image and map to a target directory in the running container.

    Key characteristics:

    • Immutable Layers: Once pushed, layers are fixed and identified by their digest.
    • Order-Invariant Extraction: Layers are independent and disjoint. They can be downloaded and extracted in any order to produce identical results. Unlike Docker layers, they do not use overlay/union filesystem semantics.
    • No File Splitting: Each file is contained entirely within a single layer. Files are never split across multiple layers.

  5. Understand the Cog Build System Flow

    main

    The Cog build system transforms your model source (consisting of cog.yaml, run.py, and model weights) into a production-ready OCI container image.

    The build process follows these stages:

    1. Config Parsing & Validation: Reads cog.yaml, validates Python versions (3.10-3.13), and auto-detects CUDA/cuDNN versions based on your PyTorch or TensorFlow requirements.
    2. Dockerfile Generation: Generates a Dockerfile including a base image, system packages, Python packages, Cog SDK/coglet wheels, user run commands, and your source code.
    3. Docker Build: Executes the build using Buildkit.
    4. Post-Build: Generates an OpenAPI schema, runs pip freeze to capture dependencies, and applies metadata labels to the image.

  6. Understand the Cog Model Schema

    main

    The Cog schema is an OpenAPI 3.0.2 specification that defines the contract for a model's interface. While every Cog model uses a fixed prediction API envelope, the schema specifically defines the model-specific input and output fields.

    Consumers use this schema to:

    • Replicate platform: Generate web UI input forms and validate requests.
    • HTTP server (coglet): Validate incoming JSON and reject malformed requests.
    • CLI (cog run): Parse -i key=value flags into correctly-typed Python objects.
    • Docker labels: Extract the model interface without running the container.
    • API clients: Determine required inputs, types, constraints (min/max, required fields), and output structures.

  7. Understand the Cog Container Runtime Architecture

    main

    Cog uses a two-process architecture to run machine learning models. This design isolates the user's model code from the HTTP server to improve stability, resource management, and shutdown handling.

    1. Parent Process (Rust HTTP Server): Acts as the orchestrator. It handles the HTTP API (on port 5000), request validation, input/output file management, webhook delivery, and the lifecycle of the worker subprocess.
    2. Worker Subprocess (Python): Executes the actual model code. It loads the user's predictor, runs the setup() method, and executes the run() (or legacy predict()) method.

    Communication between the processes occurs via a Control Channel (stdin/stdout) for lifecycle messages and a Slot Channel (Unix sockets) for per-prediction data to avoid head-of-line blocking.

  8. Understand the Coglet Architecture

    main

    Coglet is the Rust runtime for Cog, consisting of two primary components:

    1. coglet (core library): A pure Rust library that manages the lifecycle of workers, routes messages, and handles concurrency via a PermitPool. It includes an Axum-based HTTP server for prediction transport.
    2. coglet-python (PyO3 bindings): A bridge that allows Python predictors to run within the Coglet runtime. It handles the wrapping of Python predictor classes, input processing, and log routing.

  9. Core Cog Components

    main

    Cog consists of several key components that facilitate the model lifecycle:

    • Model Source: The user-provided cog.yaml (environment config), a Runner class (containing setup() and run() methods), and model weights.
    • Python SDK: The cog package used to build models. It provides BaseRunner, a type system (Input, Path, Secret, ConcatenateIterator), and the entry point for the server.
    • Schema: An OpenAPI specification automatically generated from the predictor's type hints, defining inputs and outputs.
    • Prediction API: An HTTP interface using a fixed PredictionRequest/PredictionResponse envelope format.
    • Container Runtime: A Rust HTTP server (Axum) that manages worker processes via subprocess isolation and executes predictions using PyO3 bindings.
    • Build System: Converts cog.yaml and code into a Docker image with appropriate Python, CUDA, and dependency configurations.
    • CLI: The primary command-line tool for building, testing, and deploying models.

  10. Understand Cog Architecture Overview

    main
    Cog packages machine learning models into production-ready OCI images. The workflow involves providing model code and a cog.yaml configuration, which the Cog CLI and Python SDK use to build a container image. This image contains a Rust-based server (Coglet) that serves a Prediction API via an HTTP interface.

  11. Initialize a Cog project

    main

    To start using Cog with your own machine learning model, you need to generate the necessary configuration files: cog.yaml (for system and Python dependencies) and run.py (to define the model's execution interface). Use the cog init command within your model's directory to create these files.

    $ cd path/to/your/model
$ cog init