AGENTS.md
This file provides guidance to coding agents when working with code in this repository.
Project Overview
Cog is a tool that packages machine learning models in production-ready containers.
It consists of:
- Cog CLI (
cmd/cog/) - Command-line interface for building, running, and deploying models, written in Go
- Python SDK (
python/cog/) - Python library for defining model predictors and training in Python
- Coglet (
crates/) - Rust-based prediction server that runs inside containers, with Python bindings via PyO3
Documentation for the CLI and SDK is available by reading ./docs/llms.txt.
Development Commands
Development tasks are managed with mise. Run mise tasks to see all available tasks.
Quick Reference
| Task | Description |
|---|
mise run fmt | Check formatting (all languages) |
mise run fmt:fix | Fix formatting (all languages) |
mise run lint | Run linters (all languages) |
mise run lint:fix | Fix lint issues (all languages) |
mise run test:go | Run Go tests |
mise run test:rust | Run Rust tests |
mise run test:python | Run Python tests |
mise run test:integration | Run integration tests |
mise run build:cog | Build cog CLI binary |
mise run build:coglet | Build coglet wheel (dev) |
mise run build:sdk | Build SDK wheel |
mise run install | Build and symlink cog to /usr/local/bin |
mise run docs:llm | IMPORTANT: Regenerate docs/llms.txt after editing docs |
mise run docs:cli | Generate CLI reference docs from Go source code |
Task Naming Convention
Tasks follow a consistent naming pattern:
- Language-based tasks for fmt/lint/test/typecheck:
task:go, task:rust, task:python
- Component-based tasks for build:
build:cog, build:coglet, build:sdk
- Check vs Fix:
fmt and lint default to check mode (non-destructive); use :fix suffix to auto-fix
All Tasks by Category
Format:
mise run fmt / mise run fmt:check - Check all (alias)mise run fmt:fix - Fix allmise run fmt:go / mise run fmt:rust / mise run fmt:python - Per-language
Lint:
mise run lint / mise run lint:check - Check all (alias)mise run lint:fix - Fix allmise run lint:go / mise run lint:rust / mise run lint:python - Per-languagemise run lint:rust:deny - Check Rust licenses/advisories
Test:
mise run test:go - Go unit testsmise run test:rust - Rust unit testsmise run test:python - Python unit tests (via tox)mise run test:coglet:python - Coglet Python binding testsmise run test:integration - Integration tests
Build:
mise run build:cog - Build cog CLI (development)mise run build:cog:release - Build cog CLI (release)mise run build:coglet - Build coglet wheel (dev install)mise run build:coglet:wheel - Build coglet wheel (native platform)mise run build:coglet:wheel:linux-x64 - Build for Linux x86_64mise run build:coglet:wheel:linux-arm64 - Build for Linux ARM64mise run build:sdk - Build SDK wheel
Install:
mise run install - Symlink cog CLI to /usr/local/bin (requires build:cog first)PREFIX=/custom/path mise run install - Symlink to custom location
Other:
mise run typecheck - Type check all languagesmise run generate - Run code generationmise run clean - Clean all build artifactsmise run docs - Build documentationmise run docs:serve - Serve docs locally
Code Style Guidelines
Go
- Imports: Organize in three groups separated by blank lines: (1) Standard library, (2) Third-party packages, (3) Internal packages (
github.com/replicate/cog/pkg/...)
- Formatting: Use
mise run fmt:go:fix
- Linting: Must pass golangci-lint with: errcheck, gocritic, gosec, govet, ineffassign, misspell, revive, staticcheck, unused
- Error Handling: Return errors as values; use
pkg/errors.CodedError for user-facing errors with error codes
- Naming: CamelCase for exported, camelCase for unexported
- Testing: Use
testify/require for assertions; prefer table-driven tests
Example import block:
import (
"fmt"
"github.com/spf13/cobra"
"github.com/replicate/cog/pkg/config"
)
Python
- Imports: Automatically organized by ruff/isort (stdlib → third-party → local)
- Formatting: Use
mise run fmt:python:fix
- Linting: Must pass ruff checks: E (pycodestyle), F (Pyflakes), I (isort), W (warnings), S (bandit), B (bugbear), ANN (annotations)
- Type Annotations: Required on all function signatures; use
typing_extensions for compatibility; avoid Any where possible
- Error Handling: Raise exceptions with descriptive messages; avoid generic exception catching
- Naming: snake_case for functions/variables/modules, PascalCase for classes
- Testing: Use pytest with fixtures; async tests with pytest-asyncio
- Compatibility: Must support Python 3.10-3.13
Rust
- Formatting: Use
mise run fmt:rust:fix
- Linting: Must pass
mise run lint:rust (clippy)
- Dependencies: Audited with
cargo-deny (see crates/deny.toml); run mise run lint:rust:deny
- Error Handling: Use
thiserror for typed errors, anyhow for application errors
- Naming: snake_case for functions/variables, PascalCase for types
- Testing: Use
cargo test; snapshot tests use insta
- Async: tokio runtime; async/await patterns
The CLI code is in the cmd/cog/ and pkg/ directories. Support tooling is in the tools/ directory.
The main commands for working on the CLI are:
go run ./cmd/cog - Runs the Cog CLI directly from source (requires wheel to be built first)mise run build:cog - Builds the Cog CLI binarymise run install - Symlinks the built binary to /usr/local/bin (run build:cog first), or to a custom path with PREFIX=/custom/path mise run installmise run test:go - Runs all Go unit testsgo test ./pkg/... - Runs tests directly with go test
Working on the Python SDK
The Python SDK is developed in the python/cog/ directory. It uses uv for virtual environments and tox for testing across multiple Python versions.
The main commands for working on the SDK are:
mise run build:sdk - Builds the Python wheelmise run test:python - Runs Python tests across all supported versions
Working on Coglet (Rust)
Coglet is the Rust-based prediction server that runs inside Cog containers, handling HTTP requests, worker process management, and prediction execution.
The code is in the crates/ directory:
crates/coglet/ - Core Rust library (HTTP server, worker orchestration, IPC)crates/coglet-python/ - PyO3 bindings for Python predictor integration (requires Python 3.10+)
For detailed architecture documentation, see crates/README.md and crates/coglet/README.md.
The main commands for working on Coglet are:
mise run build:coglet - Build and install coglet wheel for development (macOS, for local Rust/Python tests)mise run build:coglet:wheel:linux-x64 - Build Linux x86_64 wheel (required to test Rust changes in Docker containers via cog predict/cog train)mise run test:rust - Run Rust unit testsmise run lint:rust - Run clippy lintermise run fmt:rust:fix - Format code
Testing
Go code is tested using the built-in go test framework:
go test ./pkg/... -run <name> - Runs specific Go tests by namemise run test:go - Runs all Go unit tests
Python code is tested using tox, which allows testing across multiple Python versions and configurations:
mise run test:python - Runs all Python unit testsuv run tox -e py312-tests -- python/tests/server/test_http.py::test_openapi_specification_with_yield - Runs a specific Python test
The integration test suite in integration-tests/ tests the end-to-end functionality of the Cog CLI and Python SDK using Go's testscript framework:
mise run test:integration - Runs the integration testsmise run test:integration string_predictor - Runs a specific integration test
The integration tests require a built Cog binary, which defaults to the first cog in PATH. Run tests against a specific binary with the COG_BINARY environment variable:
mise run build:cog
COG_BINARY=dist/go/*/cog mise run test:integration
Development Workflow
- Run
mise install to set up the development environment
- Run
mise run build:sdk after making changes to the ./python directory
- Run
mise run build:coglet:wheel:linux-x64 after making changes to the ./crates directory (needed for Docker testing)
- Run
mise run build:cog to build the CLI (wheels are picked up from dist/ at Docker build time, not embedded in the binary)
- Run
mise run fmt:fix to format code
- Run
mise run lint to check code quality
- Run
mise run docs:llm to regenerate docs/llms.txt after changing README.md or any docs/*.md file
- Read the
./docs directory and make sure the documentation is up to date
Architecture
CLI Architecture (Go)
The CLI follows a command pattern with subcommands. The main components are:
pkg/cli/ - Command definitions (build, run, predict, serve, etc.)pkg/docker/ - Docker client and container managementpkg/dockerfile/ - Dockerfile generation and templatingpkg/config/ - cog.yaml parsing and validationpkg/image/ - Image building and pushing logic
Python SDK Architecture
python/cog/ - Core SDK
base_predictor.py - Base class for model predictorstypes.py - Input/output type definitionsserver/ - HTTP/queue server implementationcommand/ - Runner implementations for predict/train
Coglet Architecture (Rust)
The prediction server that runs inside Cog containers. Uses a two-process architecture: a parent process (HTTP server + orchestrator) and a worker subprocess (Python predictor execution).
See crates/README.md for detailed architecture documentation.
crates/coglet/ - Core Rust library (HTTP server, worker orchestration, IPC bridge)crates/coglet-python/ - PyO3 bindings for Python predictor integration
Key Design Patterns
- Local Wheel Resolution: The CLI discovers SDK and coglet wheels from
dist/ at Docker build time (not embedded in the binary)
- Docker SDK Integration: Uses Docker Go SDK for container operations
- Type Safety: Dataclasses for Python type validation, strongly typed Go interfaces
- Compatibility Matrix: Automated CUDA/PyTorch/TensorFlow compatibility management
For comprehensive architecture documentation, see architecture/.
Common Tasks
Adding a new CLI command
- Create command file in
pkg/cli/
- Add command to
pkg/cli/root.go
- Implement business logic in appropriate
pkg/ subdirectory
- Add tests
Modifying Python SDK behavior
- Edit files in
python/cog/
- Run
mise run build:sdk to rebuild wheel
- Test with
mise run test:python
- Integration test with
mise run test:integration
Updating ML framework compatibility
- See
tools/compatgen/ for compatibility matrix generation
- Update framework versions in relevant Dockerfile templates
- Test with various framework combinations
Updating the docs
- Documentation is in the
docs/ directory, written in Markdown and generated into HTML using mkdocs.
- IMPORTANT: After editing any file in
docs/ or README.md, you MUST run mise run docs:llm to regenerate docs/llms.txt. This file is used by coding agents and should be kept in sync with the documentation.
- IMPORTANT: CLI reference docs (
docs/cli.md) are auto-generated from Go source code. After modifying CLI commands in cmd/ or pkg/cli/, run mise run docs:cli to regenerate, and ensure mise run docs:cli:check passes before committing.
Development tools are managed in two places that must be kept in sync:
mise.toml — Tool versions for local development (uses aqua backend for prebuilt binaries).github/workflows/ci.yaml — Tool installation for CI (uses dedicated GitHub Actions)
CI deliberately avoids aqua downloads from GitHub Releases to prevent transient 502 failures. Instead, it uses:
| Tool | CI installation method | Why |
|---|
| gotestsum | go install | Uses Go module proxy, not GitHub Releases |
| cargo-deny | taiki-e/install-action | Prebuilt with checksum verification |
| cargo-nextest | taiki-e/install-action | Prebuilt with checksum verification |
| coglet wheel (maturin+zig) | PyO3/maturin-action | Bundles maturin and zig |
| golangci-lint | golangci/golangci-lint-action | Built-in caching |
| Rust toolchain | dtolnay/rust-toolchain | Guaranteed ordering |
Tools disabled in CI are listed in MISE_DISABLE_TOOLS in ci.yaml.
When updating a tool version, update both:
- The version in
mise.toml (for local dev)
- The corresponding version pin in
.github/workflows/ci.yaml (for CI)
Important Files
cog.yaml - User-facing model configurationpkg/config/config.go - Go code for parsing and validating cog.yamlpkg/config/data/config_schema_v1.0.json - JSON schema for cog.yamlpython/cog/base_predictor.py - Predictor interfacecrates/Cargo.toml - Rust workspace configurationcrates/README.md - Coglet architecture overviewmise.toml - Task definitions for development workflow
Testing Philosophy
- Unit tests for individual components (Go and Python)
- Integration tests for end-to-end workflows
- Tests use real Docker operations (no mocking Docker API)
- Always run
mise run build:sdk after making Python changes before testing Go code
- Python 3.10-3.13 compatibility is required