[ { "path": ".all-contributorsrc", "content": "{\n \"projectName\": \"cog\",\n \"projectOwner\": \"replicate\",\n \"repoType\": \"github\",\n \"repoHost\": \"https://github.com\",\n \"files\": [\n \"README.md\"\n ],\n \"imageSize\": 100,\n \"commit\": false,\n \"commitConvention\": \"none\",\n \"contributors\": [\n {\n \"login\": \"bfirsh\",\n \"name\": \"Ben Firshman\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/40906?v=4\",\n \"profile\": \"https://fir.sh/\",\n \"contributions\": [\n \"code\",\n \"doc\"\n ]\n },\n {\n \"login\": \"andreasjansson\",\n \"name\": \"Andreas Jansson\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/713993?v=4\",\n \"profile\": \"https://replicate.ai/\",\n \"contributions\": [\n \"code\",\n \"doc\",\n \"maintenance\"\n ]\n },\n {\n \"login\": \"zeke\",\n \"name\": \"Zeke Sikelianos\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/2289?v=4\",\n \"profile\": \"http://zeke.sikelianos.com/\",\n \"contributions\": [\n \"code\",\n \"doc\",\n \"tool\"\n ]\n },\n {\n \"login\": \"synek\",\n \"name\": \"Rory Byrne\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/9436784?v=4\",\n \"profile\": \"https://rory.bio/\",\n \"contributions\": [\n \"code\",\n \"doc\",\n \"test\"\n ]\n },\n {\n \"login\": \"hangtwenty\",\n \"name\": \"Michael Floering\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/2420688?v=4\",\n \"profile\": \"https://github.com/hangtwenty\",\n \"contributions\": [\n \"code\",\n \"doc\",\n \"ideas\"\n ]\n },\n {\n \"login\": \"bencevans\",\n \"name\": \"Ben Evans\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/638535?v=4\",\n \"profile\": \"https://bencevans.io/\",\n \"contributions\": [\n \"doc\"\n ]\n },\n {\n \"login\": \"imshashank\",\n \"name\": \"shashank agarwal\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/778870?v=4\",\n \"profile\": \"https://shashank.pw/\",\n \"contributions\": [\n \"code\",\n \"doc\"\n ]\n },\n {\n \"login\": \"VictorXLR\",\n \"name\": \"VictorXLR\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/22397950?v=4\",\n \"profile\": \"https://victorxlr.me/\",\n \"contributions\": [\n \"code\",\n \"doc\",\n \"test\"\n ]\n },\n {\n \"login\": \"annahung31\",\n \"name\": \"hung anna\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/39179888?v=4\",\n \"profile\": \"https://annahung31.github.io/\",\n \"contributions\": [\n \"bug\"\n ]\n },\n {\n \"login\": \"bwhitman\",\n \"name\": \"Brian Whitman\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/76612?v=4\",\n \"profile\": \"http://notes.variogr.am/\",\n \"contributions\": [\n \"bug\"\n ]\n },\n {\n \"login\": \"JimothyJohn\",\n \"name\": \"JimothyJohn\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/24216724?v=4\",\n \"profile\": \"https://github.com/JimothyJohn\",\n \"contributions\": [\n \"bug\"\n ]\n },\n {\n \"login\": \"ericguizzo\",\n \"name\": \"ericguizzo\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/26746670?v=4\",\n \"profile\": \"https://github.com/ericguizzo\",\n \"contributions\": [\n \"bug\"\n ]\n },\n {\n \"login\": \"evilstreak\",\n \"name\": \"Dominic Baggott\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/74812?v=4\",\n \"profile\": \"http://www.dominicbaggott.com\",\n \"contributions\": [\n \"code\",\n \"test\"\n ]\n },\n {\n \"login\": \"dashstander\",\n \"name\": \"Dashiell Stander\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/7449128?v=4\",\n \"profile\": \"https://github.com/dashstander\",\n \"contributions\": [\n \"bug\",\n \"code\",\n \"test\"\n ]\n },\n {\n \"login\": \"Hurricane-eye\",\n \"name\": \"Shuwei Liang\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/31437546?v=4\",\n \"profile\": \"https://github.com/Hurricane-eye\",\n \"contributions\": [\n \"bug\",\n \"question\"\n ]\n },\n {\n \"login\": \"ericallam\",\n \"name\": \"Eric Allam\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/534?v=4\",\n \"profile\": \"https://github.com/ericallam\",\n \"contributions\": [\n \"ideas\"\n ]\n },\n {\n \"login\": \"iperdomo\",\n \"name\": \"Iván Perdomo\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/178474?v=4\",\n \"profile\": \"https://perdomo.me\",\n \"contributions\": [\n \"bug\"\n ]\n },\n {\n \"login\": \"charlesfrye\",\n \"name\": \"Charles Frye\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/10442975?v=4\",\n \"profile\": \"http://charlesfrye.github.io\",\n \"contributions\": [\n \"doc\"\n ]\n },\n {\n \"login\": \"phamquiluan\",\n \"name\": \"Luan Pham\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/24642166?v=4\",\n \"profile\": \"https://github.com/phamquiluan\",\n \"contributions\": [\n \"bug\",\n \"doc\"\n ]\n },\n {\n \"login\": \"TommyDew42\",\n \"name\": \"TommyDew\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/46992350?v=4\",\n \"profile\": \"https://github.com/TommyDew42\",\n \"contributions\": [\n \"code\"\n ]\n },\n {\n \"login\": \"anotherjesse\",\n \"name\": \"Jesse Andrews\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/27?v=4\",\n \"profile\": \"https://m4ke.org\",\n \"contributions\": [\n \"code\",\n \"doc\",\n \"test\"\n ]\n },\n {\n \"login\": \"nickstenning\",\n \"name\": \"Nick Stenning\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/3602?v=4\",\n \"profile\": \"https://whiteink.com\",\n \"contributions\": [\n \"code\",\n \"doc\",\n \"design\",\n \"infra\",\n \"test\"\n ]\n },\n {\n \"login\": \"justinmerrell\",\n \"name\": \"Justin Merrell\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/14996837?v=4\",\n \"profile\": \"https://merrell.io/\",\n \"contributions\": [\n \"doc\"\n ]\n },\n {\n \"login\": \"ruriky\",\n \"name\": \"Rurik Ylä-Onnenvuori\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/19946546?v=4\",\n \"profile\": \"https://github.com/ruriky\",\n \"contributions\": [\n \"bug\"\n ]\n },\n {\n \"login\": \"youkaclub\",\n \"name\": \"Youka\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/59315275?v=4\",\n \"profile\": \"https://www.youka.club/\",\n \"contributions\": [\n \"bug\"\n ]\n },\n {\n \"login\": \"afiaka87\",\n \"name\": \"Clay Mullis\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/3994972?v=4\",\n \"profile\": \"https://github.com/afiaka87\",\n \"contributions\": [\n \"doc\"\n ]\n },\n {\n \"login\": \"mattt\",\n \"name\": \"Mattt\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/7659?v=4\",\n \"profile\": \"https://github.com/mattt\",\n \"contributions\": [\n \"code\",\n \"doc\",\n \"infra\"\n ]\n },\n {\n \"login\": \"Juneezee\",\n \"name\": \"Eng Zer Jun\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/20135478?v=4\",\n \"profile\": \"https://github.com/Juneezee\",\n \"contributions\": [\n \"test\"\n ]\n },\n {\n \"login\": \"bbedward\",\n \"name\": \"BB\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/550752?v=4\",\n \"profile\": \"https://github.com/bbedward\",\n \"contributions\": [\n \"code\"\n ]\n },\n {\n \"login\": \"williamluer\",\n \"name\": \"williamluer\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/85975676?v=4\",\n \"profile\": \"https://github.com/williamluer\",\n \"contributions\": [\n \"doc\"\n ]\n },\n {\n \"login\": \"sirupsen\",\n \"name\": \"Simon Eskildsen\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/97400?v=4\",\n \"profile\": \"http://sirupsen.com\",\n \"contributions\": [\n \"code\"\n ]\n },\n {\n \"login\": \"erbridge\",\n \"name\": \"F\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/1027364?v=4\",\n \"profile\": \"https://erbridge.co.uk\",\n \"contributions\": [\n \"bug\",\n \"code\"\n ]\n },\n {\n \"login\": \"philandstuff\",\n \"name\": \"Philip Potter\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/581269?v=4\",\n \"profile\": \"https://github.com/philandstuff\",\n \"contributions\": [\n \"bug\",\n \"code\"\n ]\n },\n {\n \"login\": \"joannejchen\",\n \"name\": \"Joanne Chen\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/33409024?v=4\",\n \"profile\": \"https://github.com/joannejchen\",\n \"contributions\": [\n \"doc\"\n ]\n },\n {\n \"login\": \"technillogue\",\n \"name\": \"technillogue\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/945691?v=4\",\n \"profile\": \"http://technillogue.github.io\",\n \"contributions\": [\n \"code\"\n ]\n },\n {\n \"login\": \"aron\",\n \"name\": \"Aron Carroll\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/47144?v=4\",\n \"profile\": \"http://aroncarroll.com\",\n \"contributions\": [\n \"doc\",\n \"code\",\n \"ideas\"\n ]\n },\n {\n \"login\": \"Theodotus1243\",\n \"name\": \"Bohdan Mykhailenko\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/32220358?v=4\",\n \"profile\": \"https://github.com/Theodotus1243\",\n \"contributions\": [\n \"doc\",\n \"bug\"\n ]\n },\n {\n \"login\": \"one1zero1one\",\n \"name\": \"Daniel Radu\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/724604?v=4\",\n \"profile\": \"https://github.com/one1zero1one\",\n \"contributions\": [\n \"doc\",\n \"bug\"\n ]\n },\n {\n \"login\": \"Etelis\",\n \"name\": \"Itay Etelis\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/92247226?v=4\",\n \"profile\": \"https://github.com/Etelis\",\n \"contributions\": [\n \"code\"\n ]\n },\n {\n \"login\": \"gschian0\",\n \"name\": \"Gennaro Schiano\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/54407820?v=4\",\n \"profile\": \"http://www.wavefunction.dev\",\n \"contributions\": [\n \"doc\"\n ]\n },\n {\n \"login\": \"aknoerig\",\n \"name\": \"André Knörig\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/481350?v=4\",\n \"profile\": \"http://andreknoerig.de\",\n \"contributions\": [\n \"doc\"\n ]\n },\n {\n \"login\": \"danfairs\",\n \"name\": \"Dan Fairs\",\n \"avatar_url\": \"https://avatars.githubusercontent.com/u/24726?v=4\",\n \"profile\": \"https://condense.live\",\n \"contributions\": [\n \"code\"\n ]\n }\n ],\n \"contributorsPerLine\": 7,\n \"skipCi\": true,\n \"commitType\": \"docs\"\n}\n" }, { "path": ".git_archival.txt", "content": "node: $Format:%H$\nnode-date: $Format:%cI$\ndescribe-name: $Format:%(describe:tags=true,match=*[0-9]*)$\n" }, { "path": ".gitattributes", "content": ".git_archival.txt export-subst\nMakefile -linguist-detectable\ndocs/llms.txt linguist-generated=true\n" }, { "path": ".github/CODEOWNERS", "content": "# Default code owners for the entire repository\n* @replicate/cog\n" }, { "path": ".github/dependabot.yml", "content": "version: 2\nupdates:\n - package-ecosystem: \"gomod\"\n directory: \"/\"\n schedule:\n interval: \"weekly\"\n - package-ecosystem: \"pip\"\n directory: \"/\"\n schedule:\n interval: \"weekly\"\n allow:\n - dependency-type: \"direct\"\n - package-ecosystem: \"cargo\"\n directory: \"/crates\"\n schedule:\n interval: \"weekly\"\n - package-ecosystem: \"github-actions\"\n directory: \"/\"\n schedule:\n interval: \"weekly\"\n" }, { "path": ".github/workflows/README.md", "content": "# CI Architecture\n\nThis document describes the CI/CD architecture for the Cog repository.\n\n## Design Principles\n\n1. **Single gate job** - Branch protection uses one required check (`ci-complete`) that depends on all other jobs\n2. **Path-based filtering** - Jobs skip when irrelevant files change (Go changes don't trigger Rust tests)\n3. **Build once, test many** - Artifacts built once and reused across test jobs\n4. **Parallel execution** - Independent jobs run concurrently\n5. **Skipped = passing** - Jobs that skip due to path filtering count as passing for the gate\n\n## Workflows\n\n### `ci.yaml` - Main CI Pipeline\n\nThe primary CI workflow that runs on all PRs and pushes to main.\n\n```\n┌─────────────────────────────────────────────────────────────────────────────┐\n│ CHANGES DETECTION │\n│ Determines which components changed: go, rust, python, integration-tests │\n└─────────────────────────────────────────────────────────────────────────────┘\n │\n ┌─────────────────┼─────────────────┐\n ▼ ▼ ▼\n ┌──────────┐ ┌──────────┐ ┌──────────┐\n │build-rust│ │ build-sdk│ │ (none) │\n │ (wheel) │ │ (wheel) │ │ │\n └────┬─────┘ └────┬─────┘ └──────────┘\n │ │\n ┌─────────────┼────────────────┼─────────────────────┐\n │ │ │ │\n ▼ ▼ ▼ ▼\n┌─────────┐ ┌──────────┐ ┌───────────┐ ┌───────────┐\n│fmt-rust │ │test-rust │ │ fmt-go │ │fmt-python │\n│lint-rust│ │coglet-py │ │ lint-go │ │lint-python│\n│ deny │ │ (matrix) │ │ test-go │ │test-python│\n└─────────┘ └────┬─────┘ └───────────┘ └───────────┘\n │ │ │\n └────────────────┼─────────────────────┘\n ▼\n ┌────────────────┐\n │test-integration│\n │ (matrix) │\n └───────┬────────┘\n ▼\n ┌───────────────┐\n │ ci-complete │ ← Branch protection requires this\n └───────────────┘\n```\n\n#### Jobs\n\n| Job | Runs when | Depends on | Purpose |\n|-----|-----------|------------|---------|\n| `changes` | Always | - | Detect which components changed |\n| `build-sdk` | python changed | changes | Build cog SDK wheel |\n| `build-rust` | rust changed | changes | Build coglet ABI3 wheel |\n| `fmt-go` | go changed | changes | Check Go formatting |\n| `fmt-rust` | rust changed | changes | Check Rust formatting |\n| `fmt-python` | python changed | changes | Check Python formatting |\n| `lint-go` | go changed | changes | Lint Go code |\n| `lint-rust` | rust changed | changes | Run clippy |\n| `lint-rust-deny` | rust changed | changes | Check licenses/advisories |\n| `lint-python` | python changed | build-sdk | Lint Python code |\n| `test-go` | go changed | build-sdk | Run Go tests (matrix: ubuntu, macos) |\n| `test-rust` | rust changed | changes | Run Rust tests |\n| `test-python` | python changed | build-sdk | Run Python tests (matrix: 3.10-3.13) |\n| `test-coglet-python` | rust or python changed | build-rust | Test coglet bindings (matrix: 3.10-3.13) |\n| `test-integration` | any changed | build-sdk, build-rust | Integration tests (matrix: cog, cog-rust) |\n| `ci-complete` | Always | all jobs | Gate job for branch protection |\n\n#### Python Version Matrix\n\nPython versions are defined once at the workflow level:\n\n```yaml\nenv:\n SUPPORTED_PYTHONS: '[\"3.10\", \"3.11\", \"3.12\", \"3.13\"]'\n```\n\nJobs that need the matrix reference it via `fromJson(env.SUPPORTED_PYTHONS)`.\n\n### `codeql.yml` - Security Analysis\n\nRuns CodeQL security scanning for Go, Python, and Rust.\n\n- **Triggers**: Push to main, PRs to main, weekly schedule\n- **Languages**: go, python, rust\n\n### Deleted Workflows\n\n- `rust.yaml` - Consolidated into `ci.yaml`. The separate workflow was redundant.\n- `pypi-package.yaml` - Replaced by `release-build.yaml` + `release-publish.yaml`.\n- `version-bump.yaml` - Removed. Just edit `crates/Cargo.toml` directly.\n\n## Caching Strategy\n\n### Rust Cache\n- **Save**: Only on `main` branch pushes (to avoid PR cache pollution)\n- **Restore**: On all runs (PRs restore from main's cache)\n- Uses `Swatinem/rust-cache@v2` with workspace path `crates -> target`\n\n### Go Cache\n- Built into `actions/setup-go` via `cache-dependency-path`\n\n### Python/uv Cache\n- Built into `jdx/mise-action` and `astral-sh/setup-uv`\n\n## Artifacts\n\n| Artifact | Contents | Retention |\n|----------|----------|-----------|\n| `CogPackage` | cog-*.whl, cog-*.tar.gz | Default (90 days) |\n| `CogletRustWheel` | coglet-*-cp310-abi3-*.whl | Default (90 days) |\n\nThe ABI3 wheel is built with Python 3.10 minimum but works on all 3.10+ versions.\n\n## Local Development\n\nUse mise tasks to run the same checks locally:\n\n```bash\n# Format (check)\nmise run fmt\n\n# Format (fix)\nmise run fmt:fix\n\n# Lint\nmise run lint\n\n# Test\nmise run test:go\nmise run test:rust\nmise run test:python\n\n# Build\nmise run build:cog\nmise run build:coglet\nmise run build:sdk\n```\n\n## Adding New Checks\n\n1. Add a mise task in `mise.toml`\n2. Add a job in `ci.yaml` with appropriate `needs` and path filtering\n3. Add the job to `ci-complete`'s needs list\n4. Update this README\n\n## Branch Protection\n\nConfigure branch protection to require only `ci-complete`:\n\n```\nSettings > Branches > main > Require status checks:\n ✓ ci-complete\n```\n\nSkipped jobs (from path filtering) are treated as passing by the gate job.\n\n## Release Workflow\n\nReleases use a two-workflow system. There are three release types:\n\n| Type | Example tag | Branch rule | Draft? | PyPI/crates.io? |\n|------|-------------|-------------|--------|-----------------|\n| **Stable** | `v0.17.0` | Must be on main | Yes (manual publish) | Yes |\n| **Pre-release** | `v0.17.0-alpha3` | Must be on main | Yes (manual publish) | Yes |\n| **Dev** | `v0.17.0-dev1` | Any branch | No (immediate) | No |\n\n### Stable / Pre-release Flow\n\n```\n Developer pushes tag on main (e.g. v0.17.0, v0.17.0-rc1)\n │\n ▼\n release-build.yaml (automatic)\n ┌──────────────────────────────────────────────┐\n │ verify-tag ──▶ build-sdk ──┐ │\n │ (must be build-coglet ┼──▶ create- │\n │ main) build-CLI ──┘ release │\n │ (DRAFT) │\n └──────────────────────────────────────────────┘\n │\n Maintainer publishes draft in GitHub UI\n │\n ▼\n release-publish.yaml (automatic)\n ┌──────────────────────────────────────────────┐\n │ coglet → PyPI ──▶ SDK → PyPI │\n │ coglet → crates.io │\n └──────────────────────────────────────────────┘\n```\n\n### Dev Release Flow\n\n```\n Developer pushes tag from any branch (e.g. v0.17.0-dev1)\n │\n ▼\n release-build.yaml (automatic)\n ┌──────────────────────────────────────────────┐\n │ verify-tag ──▶ build-sdk ──┐ │\n │ (no branch build-coglet ┼──▶ create- │\n │ restriction) build-CLI ──┘ release │\n │ (PRE- │\n │ RELEASE) │\n └──────────────────────────────────────────────┘\n │\n Done. No PyPI/crates.io.\n Wheels + CLI binaries on GH release.\n```\n\n### Workflows\n\n#### `release-build.yaml`\n\nTriggered by version tags (`v*.*.*`). Builds all artifacts and creates a GitHub release.\n\n| Job | Purpose |\n|-----|---------|\n| `verify-tag` | Cargo.toml version match + branch rules (main for stable/pre-release, any for dev) |\n| `build-sdk` | Build cog SDK wheel and sdist |\n| `build-coglet-wheels` | Build coglet wheels (3 platforms via zig cross-compile) |\n| `create-release` | Goreleaser builds CLI + creates release, then appends wheels. Dev releases are immediately published as pre-release; stable/pre-release remain as draft. |\n\n**Security**: No secrets needed for dev. Stable/pre-release require maintainer to publish draft.\n\n#### `release-publish.yaml`\n\nTriggered when a release is published. Publishes to PyPI and crates.io.\n**Skips entirely for dev releases** (all jobs gated on `is_dev != true`).\n\n| Job | Depends on | Purpose |\n|-----|------------|---------|\n| `verify-release` | - | Validate tag format, classify release type |\n| `publish-pypi-coglet` | verify-release | Publish coglet to PyPI (trusted publishing) |\n| `publish-pypi-sdk` | publish-pypi-coglet | Publish SDK to PyPI (waits for coglet) |\n| `publish-crates-io` | verify-release | Publish coglet crate (OIDC) |\n| `update-homebrew-tap` | publish-pypi-sdk, publish-crates-io | Update `replicate/homebrew-tap` cask (stable only, macOS, via GH App) |\n\n### Package Versioning\n\nAll packages use **lockstep versioning** from `crates/Cargo.toml`.\n\n| Package | Registry | Version format | Example |\n|---------|----------|----------------|---------|\n| cog SDK | PyPI | PEP 440 | `cog==0.17.0`, `cog==0.17.0a3`, `cog==0.17.0.dev1` |\n| coglet | PyPI | PEP 440 | `coglet==0.17.0`, `coglet==0.17.0a3` |\n| coglet | crates.io | semver | `coglet@0.17.0`, `coglet@0.17.0-alpha3` |\n| CLI | GitHub Release | semver | `cog v0.17.0`, `cog v0.17.0-dev1` |\n\n**Version conversion** (semver -> PEP 440):\n- `0.17.0-alpha3` -> `0.17.0a3`\n- `0.17.0-beta1` -> `0.17.0b1`\n- `0.17.0-rc1` -> `0.17.0rc1`\n- `0.17.0-dev1` -> `0.17.0.dev1`\n- `0.17.0` -> `0.17.0`\n\n### SDK Wheel Sourcing\n\nThe CLI installs the cog SDK from PyPI at container build time:\n\n| Scenario | COG_SDK_WHEEL env var | Behavior |\n|----------|-----------------------|----------|\n| Released CLI | (unset) | Install latest `cog` from PyPI |\n| Dev CLI (in repo) | (unset) | Auto-detect `dist/cog-*.whl` if present, else PyPI |\n| Force PyPI | `pypi` | Install latest from PyPI |\n| Specific version | `pypi:0.12.0` | Install `cog==0.12.0` from PyPI |\n| Local wheel | `/path/to/cog.whl` | Install from local file |\n| Force dist | `dist` | Install from `dist/` (error if missing) |\n\nSame pattern for `COGLET_WHEEL` (but coglet is optional by default).\n\n### GitHub Environment Setup\n\n1. Create environments in **Settings -> Environments**:\n - `pypi` - For PyPI publishing (trusted publishing, no secrets)\n - `crates-io` - For crates.io publishing (trusted publishing, no secrets)\n\n2. Configure protection rules for each environment:\n - **Deployment branches**: \"Selected branches and tags\"\n - **Add pattern**: `v*` (restricts to version tags)\n - **Required reviewers**: Add maintainers\n\n3. Configure trusted publishers:\n - **PyPI** (both `cog` and `coglet`): workflow `release-publish.yaml`, environment `pypi`\n - **crates.io** (`coglet`): workflow `release-publish.yaml`, environment `crates-io`\n\n4. Configure the Homebrew tap GitHub App:\n - App: `cog-homebrew-tapbot` (ID: 1232932405)\n - Create environment `homebrew` with secret `COG_HOMEBREW_TAP_PRIVATE_KEY` (app private key)\n - App must have write access to `replicate/homebrew-tap`\n\n### Performing a Stable / Pre-release\n\n```bash\n# 1. Update crates/Cargo.toml version (e.g. \"0.17.0\" or \"0.17.0-alpha3\")\n# 2. Merge to main\n\n# 3. Tag and push\ngit tag v0.17.0\ngit push origin v0.17.0\n\n# 4. Wait for release-build.yaml to complete (creates draft release)\n# 5. Review the draft release in GitHub UI\n# 6. Click \"Publish release\" -> triggers release-publish.yaml -> PyPI + crates.io\n```\n\n### Performing a Dev Release\n\n```bash\n# From any branch:\n# 1. Update crates/Cargo.toml version (e.g. \"0.17.0-dev1\")\n# 2. Commit and push\n\n# 3. Tag and push\ngit tag v0.17.0-dev1\ngit push origin v0.17.0-dev1\n\n# 4. Done. release-build.yaml creates a pre-release with all artifacts.\n# No PyPI/crates.io publishing. No manual approval needed.\n```\n" }, { "path": ".github/workflows/ci.yaml", "content": "name: CI\n\non:\n merge_group:\n push:\n branches: [main]\n pull_request:\n workflow_dispatch:\n\n# Cancel in-progress runs for PRs, queue for merge group and main\nconcurrency:\n group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}-v2\n cancel-in-progress: ${{ github.event_name == 'pull_request' }}\n\nenv:\n # Single source of truth for supported Python versions\n SUPPORTED_PYTHONS: '[\"3.10\", \"3.11\", \"3.12\", \"3.13\"]'\n # Default Python version for non-matrix jobs\n PYTHON_VERSION: \"3.13\"\n # Minimum supported Python — used for ABI3 wheel builds and glob patterns.\n # Must match the lowest entry in SUPPORTED_PYTHONS.\n MINIMUM_PYTHON: \"3.10\"\n # Number of runners to shard integration tests across (per runtime)\n # Slow tests ([short] skip) are distributed round-robin first, then fast tests fill in\n NUM_IT_RUNNER_SHARDS: \"4\"\n # Standard environment\n HYPOTHESIS_PROFILE: ci\n FORCE_COLOR: \"1\"\n PIP_DISABLE_PIP_VERSION_CHECK: \"1\"\n PIP_NO_PYTHON_VERSION_WARNING: \"1\"\n CARGO_TERM_COLOR: always\n # CGo required for go-tree-sitter (static Python schema parser)\n CGO_ENABLED: \"1\"\n # Disable tools in mise that CI installs via dedicated GitHub Actions for\n # better reliability (avoids transient GitHub Releases 502s from aqua downloads),\n # better caching, and guaranteed tool ordering.\n # - Rust toolchain: dtolnay/rust-toolchain\n # - cargo-binstall: taiki-e/install-action\n # - Python: astral-sh/setup-uv\n # - golangci-lint: golangci/golangci-lint-action\n # - gotestsum: go install (uses Go module proxy, not GitHub Releases)\n # - cargo-deny, cargo-nextest: taiki-e/install-action\n # - zig, cargo-zigbuild, maturin, cargo-insta: not needed in CI (maturin-action bundles zig)\n MISE_DISABLE_TOOLS: rust,rustup,rustup-init,cargo-binstall,python,golangci-lint,gotestsum,cargo-deny,cargo-insta,cargo-nextest,cargo:cargo-nextest,zig,cargo-zigbuild,maturin,cargo:maturin\n\npermissions: {}\n\n# =============================================================================\n# Change Detection\n# =============================================================================\n\njobs:\n changes:\n name: Detect changes\n runs-on: ubuntu-latest\n timeout-minutes: 5\n outputs:\n go: ${{ steps.filter.outputs.go }}\n rust: ${{ steps.filter.outputs.rust }}\n python: ${{ steps.filter.outputs.python }}\n integration: ${{ steps.filter.outputs.integration }}\n docs: ${{ steps.filter.outputs.docs }}\n version_only: ${{ steps.filter.outputs.version_only }}\n version_changed: ${{ steps.filter.outputs.version_changed }}\n # Pass through for matrix jobs (env context unavailable in strategy)\n supported_pythons: ${{ env.SUPPORTED_PYTHONS }}\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n\n - name: Detect changed paths\n id: filter\n run: |\n # For PRs, compare against base; for pushes, compare against previous commit;\n # for merge_group, compare against the merge group base.\n if [ \"${{ github.event_name }}\" = \"pull_request\" ]; then\n BASE=\"${{ github.event.pull_request.base.sha }}\"\n elif [ \"${{ github.event_name }}\" = \"merge_group\" ]; then\n BASE=\"${{ github.event.merge_group.base_sha }}\"\n else\n BASE=\"${{ github.event.before }}\"\n # Handle initial push (no before)\n if [ \"$BASE\" = \"0000000000000000000000000000000000000000\" ]; then\n BASE=\"HEAD~1\"\n fi\n fi\n\n echo \"Comparing $BASE..HEAD\"\n\n # Get changed files\n CHANGED=$(git diff --name-only \"$BASE\" HEAD 2>/dev/null || echo \"\")\n \n # Check if coglet version changed\n VERSION_CHANGED=\"false\"\n if echo \"$CHANGED\" | grep -qE '^crates/Cargo\\.toml$'; then\n # Check if only the version line changed\n if git diff \"$BASE\" HEAD -- crates/Cargo.toml | grep -qE '^\\+version = '; then\n VERSION_CHANGED=\"true\"\n echo \"Coglet version changed\"\n fi\n fi\n echo \"version_changed=$VERSION_CHANGED\" >> $GITHUB_OUTPUT\n \n # Check if ONLY the version changed (version bump PR)\n # This is true if crates/Cargo.toml is the only file and only version line changed\n VERSION_ONLY=\"false\"\n if [ \"$VERSION_CHANGED\" = \"true\" ]; then\n FILE_COUNT=$(echo \"$CHANGED\" | grep -c . || echo \"0\")\n if [ \"$FILE_COUNT\" = \"1\" ]; then\n # Only crates/Cargo.toml changed, check if only version line changed\n # Get actual diff lines (excluding +++ and --- headers)\n DIFF_CONTENT=$(git diff \"$BASE\" HEAD -- crates/Cargo.toml | grep -E '^[+-]' | grep -v '^[+-]{3}')\n # Should be exactly: -version = \"old\" and +version = \"new\"\n MINUS_LINES=$(echo \"$DIFF_CONTENT\" | grep -c '^-' || echo \"0\")\n PLUS_LINES=$(echo \"$DIFF_CONTENT\" | grep -c '^\\+' || echo \"0\")\n VERSION_MINUS=$(echo \"$DIFF_CONTENT\" | grep -c '^-version = ' || echo \"0\")\n VERSION_PLUS=$(echo \"$DIFF_CONTENT\" | grep -c '^\\+version = ' || echo \"0\")\n \n if [ \"$MINUS_LINES\" = \"1\" ] && [ \"$PLUS_LINES\" = \"1\" ] && \\\n [ \"$VERSION_MINUS\" = \"1\" ] && [ \"$VERSION_PLUS\" = \"1\" ]; then\n VERSION_ONLY=\"true\"\n echo \"Version-only change detected - skipping heavy CI\"\n fi\n fi\n fi\n echo \"version_only=$VERSION_ONLY\" >> $GITHUB_OUTPUT\n\n # CI/tooling changes should run everything (unless version-only)\n if [ \"$VERSION_ONLY\" = \"true\" ]; then\n echo \"go=false\" >> $GITHUB_OUTPUT\n echo \"rust=false\" >> $GITHUB_OUTPUT\n echo \"python=false\" >> $GITHUB_OUTPUT\n echo \"integration=false\" >> $GITHUB_OUTPUT\n echo \"docs=false\" >> $GITHUB_OUTPUT\n elif echo \"$CHANGED\" | grep -qE '^(\\.github/workflows/|mise\\.toml)'; then\n echo \"CI/tooling changed - running all jobs\"\n echo \"go=true\" >> $GITHUB_OUTPUT\n echo \"rust=true\" >> $GITHUB_OUTPUT\n echo \"python=true\" >> $GITHUB_OUTPUT\n echo \"integration=true\" >> $GITHUB_OUTPUT\n echo \"docs=true\" >> $GITHUB_OUTPUT\n else\n # Detect Go changes\n if echo \"$CHANGED\" | grep -qE '^(cmd/|pkg/|go\\.(mod|sum)|\\.golangci\\.yml|Makefile)'; then\n echo \"go=true\" >> $GITHUB_OUTPUT\n else\n echo \"go=false\" >> $GITHUB_OUTPUT\n fi\n\n # Detect Rust changes\n if echo \"$CHANGED\" | grep -qE '^(crates/|Cargo\\.(toml|lock))'; then\n echo \"rust=true\" >> $GITHUB_OUTPUT\n else\n echo \"rust=false\" >> $GITHUB_OUTPUT\n fi\n\n # Detect Python changes\n if echo \"$CHANGED\" | grep -qE '^(python/|pyproject\\.toml|uv\\.lock|noxfile\\.py|\\.ruff\\.toml)'; then\n echo \"python=true\" >> $GITHUB_OUTPUT\n else\n echo \"python=false\" >> $GITHUB_OUTPUT\n fi\n\n # Detect integration test changes (or if any code changed)\n if echo \"$CHANGED\" | grep -qE '^(integration-tests/|cmd/|pkg/|python/|crates/|go\\.(mod|sum)|uv\\.lock|pyproject\\.toml)'; then\n echo \"integration=true\" >> $GITHUB_OUTPUT\n else\n echo \"integration=false\" >> $GITHUB_OUTPUT\n fi\n\n # Detect docs changes (includes CLI source which generates docs/cli.md)\n if echo \"$CHANGED\" | grep -qE '^(docs/|README\\.md|cmd/|pkg/cli/)'; then\n echo \"docs=true\" >> $GITHUB_OUTPUT\n else\n echo \"docs=false\" >> $GITHUB_OUTPUT\n fi\n fi\n\n # Debug output\n echo \"Changed files:\"\n echo \"$CHANGED\" | head -50\n\n# =============================================================================\n# Version Check - Validates coglet version changes\n# =============================================================================\n\n version-check:\n name: Validate coglet version\n needs: changes\n if: needs.changes.outputs.version_changed == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 5\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n\n - name: Validate version\n run: |\n # Get version from Cargo.toml\n VERSION=$(grep '^version = ' crates/Cargo.toml | head -1 | sed 's/version = \"\\(.*\\)\"/\\1/')\n echo \"Coglet version: $VERSION\"\n \n # Validate semver format\n if [[ ! \"$VERSION\" =~ ^[0-9]+\\.[0-9]+\\.[0-9]+(-[a-zA-Z0-9.]+)?$ ]]; then\n echo \"::error::Invalid version format: $VERSION\"\n echo \"::error::Expected semver format: MAJOR.MINOR.PATCH or MAJOR.MINOR.PATCH-prerelease\"\n exit 1\n fi\n echo \"✓ Valid semver format\"\n \n # Check version doesn't already exist as a tag\n if git tag -l \"v$VERSION\" | grep -q .; then\n echo \"::error::Tag v$VERSION already exists!\"\n echo \"::error::Cannot set version to an already-released version.\"\n exit 1\n fi\n echo \"✓ Version not yet released\"\n \n # Get the highest existing stable version tag\n HIGHEST_TAG=$(git tag -l 'v[0-9]*.[0-9]*.[0-9]*' | grep -v '-' | sed 's/^v//' | sort -V | tail -1)\n if [ -n \"$HIGHEST_TAG\" ]; then\n echo \"Highest released version: $HIGHEST_TAG\"\n \n # Check it's not a downgrade (using sort -V for proper semver comparison)\n BASE_VERSION=\"${VERSION%%-*}\"\n SORTED_HIGHEST=$(printf '%s\\n%s' \"$HIGHEST_TAG\" \"$BASE_VERSION\" | sort -V | tail -1)\n if [ \"$SORTED_HIGHEST\" = \"$HIGHEST_TAG\" ] && [ \"$HIGHEST_TAG\" != \"$BASE_VERSION\" ]; then\n echo \"::error::Cannot downgrade version from $HIGHEST_TAG to $VERSION\"\n echo \"::error::New version must be greater than the highest released version.\"\n exit 1\n fi\n echo \"✓ Version is not a downgrade\"\n else\n echo \"No existing version tags found\"\n fi\n \n echo \"\"\n echo \"✓ Version $VERSION is valid for release\"\n\n# =============================================================================\n# Build Stage - Produces artifacts for downstream jobs\n# =============================================================================\n\n build-sdk:\n name: Build SDK\n needs: changes\n if: needs.changes.outputs.python == 'true' || needs.changes.outputs.go == 'true' || needs.changes.outputs.integration == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 15\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - uses: astral-sh/setup-uv@v6\n with:\n python-version: ${{ env.PYTHON_VERSION }}\n - uses: dtolnay/rust-toolchain@stable\n - run: rustup default stable\n - uses: taiki-e/install-action@cargo-binstall\n - uses: jdx/mise-action@v4\n with:\n cache: false\n - name: Build SDK\n run: mise run ci:build:sdk\n - name: Upload SDK package\n uses: actions/upload-artifact@v6\n with:\n name: CogPackage\n path: dist/cog-*\n\n build-rust:\n name: Build coglet wheel\n needs: changes\n if: needs.changes.outputs.rust == 'true' || needs.changes.outputs.integration == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 15\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - uses: astral-sh/setup-uv@v6\n with:\n python-version: ${{ env.PYTHON_VERSION }}\n - uses: dtolnay/rust-toolchain@stable\n - uses: Swatinem/rust-cache@v2\n with:\n workspaces: crates -> target\n save-if: ${{ github.ref == 'refs/heads/main' }}\n # No mise needed - maturin-action bundles maturin and zig\n # Explicitly request MINIMUM_PYTHON inside the manylinux container so\n # maturin produces an ABI3 wheel (cp310-abi3). Without this, maturin\n # picks up the container's default Python (3.8), which doesn't support\n # ABI3, producing a cp38-cp38 wheel that the upload glob won't match.\n - name: Build coglet wheel (ABI3)\n uses: PyO3/maturin-action@v1\n with:\n target: x86_64-unknown-linux-gnu\n args: --release --out dist -m crates/coglet-python/Cargo.toml --interpreter python${{ env.MINIMUM_PYTHON }}\n manylinux: auto\n - name: Verify ABI3 wheel exists\n run: |\n CPVER=\"cp${MINIMUM_PYTHON//.}\"\n ls -la dist/coglet-*-${CPVER}-abi3-*.whl\n - name: Upload coglet wheel\n uses: actions/upload-artifact@v6\n with:\n name: CogletRustWheel\n # ABI3 wheels use cpXYZ-abi3 naming; just match any abi3 wheel\n path: dist/coglet-*-abi3-*.whl\n\n build-cog:\n name: Build cog CLI\n needs: changes\n if: needs.changes.outputs.go == 'true' || needs.changes.outputs.integration == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 15\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - uses: actions/setup-go@v6\n with:\n go-version-file: go.mod\n cache-dependency-path: go.sum\n - uses: mlugg/setup-zig@v2\n with:\n version: 0.15.2\n - name: Get version from Cargo.toml\n id: version\n run: echo \"version=$(grep '^version' crates/Cargo.toml | head -1 | sed 's/.*\"\\(.*\\)\"/\\1/')\" >> \"$GITHUB_OUTPUT\"\n - name: Build cog binary\n uses: goreleaser/goreleaser-action@v7\n with:\n version: '~> v2'\n args: build --clean --snapshot --single-target --id cog --output cog\n env:\n GOFLAGS: -buildvcs=false\n # Use Cargo.toml as version source so snapshot builds match the wheel version\n COG_VERSION: ${{ steps.version.outputs.version }}\n - name: Upload cog binary\n uses: actions/upload-artifact@v6\n with:\n name: CogBinary\n path: cog\n\n# =============================================================================\n# Format Checks - Fast, parallel\n# =============================================================================\n\n fmt-go:\n name: Format Go\n needs: changes\n if: needs.changes.outputs.go == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 10\n steps:\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n - run: rustup default stable\n - uses: taiki-e/install-action@cargo-binstall\n - uses: jdx/mise-action@v4\n with:\n cache: false\n - name: Check Go formatting\n run: mise run fmt:go\n\n fmt-rust:\n name: Format Rust\n needs: changes\n if: needs.changes.outputs.rust == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 10\n steps:\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n with:\n components: rustfmt\n # No mise needed - rustfmt comes with toolchain\n - name: Check Rust formatting\n run: cargo fmt --manifest-path crates/Cargo.toml --all -- --check\n\n fmt-python:\n name: Format Python\n needs: changes\n if: needs.changes.outputs.python == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 10\n steps:\n - uses: actions/checkout@v6\n - uses: astral-sh/setup-uv@v6\n with:\n python-version: ${{ env.PYTHON_VERSION }}\n - uses: dtolnay/rust-toolchain@stable\n - run: rustup default stable\n - uses: taiki-e/install-action@cargo-binstall\n - uses: jdx/mise-action@v4\n with:\n cache: false\n - name: Check Python formatting\n run: mise run fmt:python\n\n check-llm-docs:\n name: Check LLM docs\n needs: changes\n if: needs.changes.outputs.docs == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 5\n steps:\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n - run: rustup default stable\n - uses: taiki-e/install-action@cargo-binstall\n - uses: jdx/mise-action@v4\n with:\n cache: false\n - name: Check llms.txt is up to date\n run: mise run docs:llm:check\n - name: Check CLI docs are up to date\n run: mise run docs:cli:check\n\n# =============================================================================\n# Lint Checks - Parallel\n# =============================================================================\n\n lint-go:\n name: Lint Go\n needs: changes\n if: needs.changes.outputs.go == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 15\n steps:\n - uses: actions/checkout@v6\n - uses: actions/setup-go@v6\n with:\n go-version-file: go.mod\n - uses: golangci/golangci-lint-action@v9\n with:\n version: v2.10.1\n\n lint-rust:\n name: Lint Rust\n needs: changes\n if: needs.changes.outputs.rust == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 15\n steps:\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n with:\n components: clippy\n - uses: Swatinem/rust-cache@v2\n with:\n workspaces: crates -> target\n save-if: ${{ github.ref == 'refs/heads/main' }}\n # No mise needed - clippy comes with toolchain\n - name: Lint Rust (clippy)\n run: cargo clippy --manifest-path crates/Cargo.toml --workspace -- -D warnings\n\n lint-rust-deny:\n name: Lint Rust (deny)\n needs: changes\n if: needs.changes.outputs.rust == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 10\n steps:\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n - uses: taiki-e/install-action@v2\n with:\n tool: cargo-deny@0.19.0\n # No mise needed - cargo-deny installed via taiki-e\n - name: Check licenses and advisories\n run: cargo deny --manifest-path crates/Cargo.toml check\n\n lint-python:\n name: Lint Python\n needs: [changes, build-sdk, build-rust]\n if: |\n needs.changes.outputs.python == 'true' &&\n (needs.build-rust.result == 'success' || needs.build-rust.result == 'skipped')\n runs-on: ubuntu-latest-8-cores\n timeout-minutes: 15\n steps:\n - name: Download SDK\n uses: actions/download-artifact@v8\n with:\n name: CogPackage\n path: dist\n - name: Download coglet wheel\n uses: actions/download-artifact@v8\n with:\n name: CogletRustWheel\n path: dist\n if: needs.build-rust.result == 'success'\n - name: Extract source distribution\n run: tar xf dist/*.tar.gz --strip-components=1\n - uses: astral-sh/setup-uv@v6\n with:\n python-version: ${{ env.PYTHON_VERSION }}\n - uses: dtolnay/rust-toolchain@stable\n - run: rustup default stable\n - uses: taiki-e/install-action@cargo-binstall\n - uses: jdx/mise-action@v4\n with:\n cache: false\n - name: Lint Python\n run: mise run lint:python\n\n# =============================================================================\n# Test Jobs\n# =============================================================================\n\n test-go:\n name: \"Test Go (${{ matrix.platform }})\"\n needs: changes\n if: needs.changes.outputs.go == 'true'\n timeout-minutes: 30\n strategy:\n fail-fast: false\n matrix:\n platform: [ubuntu-latest, macos-latest]\n runs-on: ${{ matrix.platform }}\n steps:\n - uses: actions/checkout@v6\n - uses: actions/setup-go@v6\n with:\n go-version-file: go.mod\n # gotestsum via Go module proxy (not GitHub Releases) for reliability\n - name: Install gotestsum\n run: go install gotest.tools/gotestsum@v1.13.0\n - name: Test Go\n shell: bash\n run: |\n set -euo pipefail\n set -m # job control, ensures script is in its own process group\n cleanup() {\n echo \"::warning::Cancelling...\"\n kill -TERM -- -$$ 2>/dev/null || true\n sleep 5\n kill -KILL -- -$$ 2>/dev/null || true\n }\n trap cleanup INT TERM\n gotestsum -- -short -timeout 1200s -parallel 5 ./... &\n wait $!\n\n fuzz-go:\n name: Fuzz Go\n needs: changes\n if: needs.changes.outputs.go == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 10\n env:\n CGO_ENABLED: \"1\"\n steps:\n - uses: actions/checkout@v6\n - uses: actions/setup-go@v6\n with:\n go-version-file: go.mod\n - name: Fuzz schema type resolution\n run: go test ./pkg/schema/ -run='^$' -fuzz=FuzzResolveSchemaType -fuzztime=30s\n - name: Fuzz JSON schema generation\n run: go test ./pkg/schema/ -run='^$' -fuzz=FuzzJSONSchema -fuzztime=30s\n - name: Fuzz Python parser\n run: go test ./pkg/schema/python/ -run='^$' -fuzz=FuzzParsePredictor -fuzztime=30s\n - name: Fuzz type annotation parsing\n run: go test ./pkg/schema/python/ -run='^$' -fuzz=FuzzParseTypeAnnotation -fuzztime=30s\n\n test-rust:\n name: Test Rust\n needs: changes\n if: needs.changes.outputs.rust == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 15\n steps:\n - uses: actions/checkout@v6\n - uses: dtolnay/rust-toolchain@stable\n - uses: taiki-e/install-action@v2\n with:\n tool: cargo-nextest@0.9.120\n - uses: Swatinem/rust-cache@v2\n with:\n workspaces: crates -> target\n save-if: ${{ github.ref == 'refs/heads/main' }}\n # No mise needed - cargo-nextest installed via taiki-e\n - name: Test Rust\n run: cargo nextest run --manifest-path crates/Cargo.toml --workspace --exclude coglet-python --no-tests=pass\n\n test-python:\n name: \"Test Python ${{ matrix.python-version }}\"\n needs: [changes, build-sdk, build-rust]\n if: |\n needs.changes.outputs.python == 'true' &&\n needs.build-sdk.result == 'success' &&\n (needs.build-rust.result == 'success' || needs.build-rust.result == 'skipped')\n runs-on: ubuntu-latest-8-cores\n timeout-minutes: 30\n strategy:\n fail-fast: false\n matrix:\n python-version: ${{ fromJSON(needs.changes.outputs.supported_pythons) }}\n steps:\n - name: Download artifacts\n uses: actions/download-artifact@v8\n with:\n path: dist\n merge-multiple: true\n - name: Extract source distribution\n run: tar xf dist/*.tar.gz --strip-components=1\n - uses: astral-sh/setup-uv@v6\n with:\n python-version: ${{ matrix.python-version }}\n - uses: dtolnay/rust-toolchain@stable\n - run: rustup default stable\n - uses: taiki-e/install-action@cargo-binstall\n - uses: jdx/mise-action@v4\n - name: Remove src to ensure tests run against wheel\n run: rm -rf python/cog\n - name: Test Python\n run: uvx nox -s tests -p ${{ matrix.python-version }}\n\n test-coglet-python:\n name: \"Test coglet-python (${{ matrix.python-version }})\"\n needs: [changes, build-rust]\n if: |\n always() &&\n (needs.changes.outputs.rust == 'true' || needs.changes.outputs.python == 'true') &&\n (needs.build-rust.result == 'success' || needs.build-rust.result == 'skipped')\n runs-on: ubuntu-latest\n timeout-minutes: 15\n strategy:\n fail-fast: false\n matrix:\n python-version: ${{ fromJSON(needs.changes.outputs.supported_pythons) }}\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - name: Download coglet wheel\n uses: actions/download-artifact@v8\n with:\n name: CogletRustWheel\n path: dist\n if: needs.build-rust.result == 'success'\n - uses: dtolnay/rust-toolchain@stable\n - run: rustup default stable # Required for cargo-binstall to find cargo\n - uses: taiki-e/install-action@cargo-binstall\n - uses: Swatinem/rust-cache@v2\n with:\n workspaces: crates -> target\n save-if: ${{ github.ref == 'refs/heads/main' }}\n - uses: astral-sh/setup-uv@v6\n with:\n python-version: ${{ matrix.python-version }}\n - name: Test coglet-python bindings\n run: uvx nox -s coglet -p ${{ matrix.python-version }}\n\n # Compute integration test shards dynamically.\n # Slow tests (tagged with [short] skip) are distributed round-robin first,\n # then remaining tests fill in. This ensures slow tests don't pile up on one runner.\n integration-shards:\n name: Compute test shards\n needs: changes\n if: needs.changes.outputs.integration == 'true'\n runs-on: ubuntu-latest\n timeout-minutes: 5\n outputs:\n shards: ${{ steps.shard.outputs.shards }}\n steps:\n - uses: actions/checkout@v6\n - name: Compute shards\n id: shard\n run: |\n NUM_SHARDS=${{ env.NUM_IT_RUNNER_SHARDS }}\n \n # Find unconditionally skipped tests (bare \"skip\" without condition brackets)\n # These are disabled tests that shouldn't affect shard distribution\n SKIPPED_TESTS=$(grep -rl '^skip ' integration-tests/tests/*.txtar | \\\n xargs -I{} basename {} .txtar | sort || echo \"\")\n \n # Identify slow tests (have [short] skip marker), excluding unconditionally skipped\n SLOW_TESTS=$(grep -rl '\\[short\\] skip' integration-tests/tests/*.txtar | \\\n xargs -I{} basename {} .txtar | sort)\n if [ -n \"$SKIPPED_TESTS\" ]; then\n SLOW_TESTS=$(comm -23 <(echo \"$SLOW_TESTS\") <(echo \"$SKIPPED_TESTS\"))\n fi\n \n # All tests\n ALL_TESTS=$(ls integration-tests/tests/*.txtar | \\\n xargs -I{} basename {} .txtar | sort)\n \n # Fast tests = all - slow (skipped tests end up here but run instantly)\n FAST_TESTS=$(comm -23 <(echo \"$ALL_TESTS\") <(echo \"$SLOW_TESTS\"))\n \n # Distribute slow tests round-robin across shards\n declare -a SHARDS\n for i in $(seq 0 $((NUM_SHARDS - 1))); do\n SHARDS[$i]=\"\"\n done\n \n idx=0\n while IFS= read -r test; do\n [ -z \"$test\" ] && continue\n if [ -n \"${SHARDS[$idx]}\" ]; then\n SHARDS[$idx]=\"${SHARDS[$idx]}|${test}\"\n else\n SHARDS[$idx]=\"$test\"\n fi\n idx=$(( (idx + 1) % NUM_SHARDS ))\n done <<< \"$SLOW_TESTS\"\n \n # Distribute fast tests round-robin across shards\n while IFS= read -r test; do\n [ -z \"$test\" ] && continue\n if [ -n \"${SHARDS[$idx]}\" ]; then\n SHARDS[$idx]=\"${SHARDS[$idx]}|${test}\"\n else\n SHARDS[$idx]=\"$test\"\n fi\n idx=$(( (idx + 1) % NUM_SHARDS ))\n done <<< \"$FAST_TESTS\"\n \n # Build JSON array of shard objects\n JSON=\"[\"\n for i in $(seq 0 $((NUM_SHARDS - 1))); do\n PATTERN=\"${SHARDS[$i]}\"\n COUNT=$(echo \"$PATTERN\" | tr '|' '\\n' | wc -l | tr -d ' ')\n [ $i -gt 0 ] && JSON=\"${JSON},\"\n JSON=\"${JSON}{\\\"index\\\":$i,\\\"pattern\\\":\\\"${PATTERN}\\\",\\\"count\\\":$COUNT}\"\n done\n JSON=\"${JSON}]\"\n \n echo \"shards=$JSON\" >> \"$GITHUB_OUTPUT\"\n \n # Debug output\n echo \"Shard distribution:\"\n for i in $(seq 0 $((NUM_SHARDS - 1))); do\n COUNT=$(echo \"${SHARDS[$i]}\" | tr '|' '\\n' | wc -l | tr -d ' ')\n SLOW_COUNT=$(echo \"${SHARDS[$i]}\" | tr '|' '\\n' | while read t; do\n echo \"$SLOW_TESTS\" | grep -q \"^${t}$\" && echo \"$t\"\n done | wc -l | tr -d ' ')\n echo \" Shard $i: $COUNT tests ($SLOW_COUNT slow)\"\n done\n\n test-integration:\n name: \"Test integration (shard ${{ matrix.shard.index }})\"\n needs: [changes, build-cog, build-sdk, build-rust, integration-shards]\n if: |\n !cancelled() &&\n needs.changes.outputs.integration == 'true' &&\n needs.integration-shards.result == 'success' &&\n (needs.build-cog.result == 'success' || needs.build-cog.result == 'skipped') &&\n (needs.build-sdk.result == 'success' || needs.build-sdk.result == 'skipped') &&\n (needs.build-rust.result == 'success' || needs.build-rust.result == 'skipped')\n runs-on: ubuntu-latest-16-cores\n timeout-minutes: 30\n strategy:\n fail-fast: false\n matrix:\n shard: ${{ fromJSON(needs.integration-shards.outputs.shards) }}\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - name: Login to Docker Hub\n uses: docker/login-action@v4\n if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request'\n with:\n registry: index.docker.io\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_TOKEN }}\n - name: Download artifacts\n uses: actions/download-artifact@v8\n with:\n path: dist\n merge-multiple: true\n - name: Install cog binary\n run: |\n cp dist/cog ./cog\n chmod +x ./cog\n - uses: actions/setup-go@v6\n with:\n go-version-file: go.mod\n cache-dependency-path: go.sum\n # gotestsum via Go module proxy (not GitHub Releases) for reliability\n - name: Install gotestsum\n run: go install gotest.tools/gotestsum@v1.13.0\n - name: Set wheel environment\n run: |\n # Use locally-built wheels, not PyPI (version may not be published yet)\n # Must use absolute paths — cog subprocess runs from txtar workdir, not checkout root\n echo \"COG_SDK_WHEEL=${{ github.workspace }}/dist\" >> $GITHUB_ENV\n echo \"COGLET_WHEEL=${{ github.workspace }}/dist\" >> $GITHUB_ENV\n - name: Run integration tests (shard ${{ matrix.shard.index }}, ${{ matrix.shard.count }} tests)\n env:\n COG_BINARY: ./cog\n TEST_PARALLEL: 4\n BUILDKIT_PROGRESS: 'quiet'\n shell: bash\n run: |\n set -euo pipefail\n set -m # job control, ensures script is in its own process group\n cleanup() {\n echo \"::warning::Cancelling...\"\n kill -TERM -- -$$ 2>/dev/null || true\n sleep 5\n kill -KILL -- -$$ 2>/dev/null || true\n }\n trap cleanup INT TERM\n \n # Build -run regex from shard pattern\n # Pattern is \"test1|test2|test3\" - wrap each in TestIntegration//\n RUN_PATTERN=\"${{ matrix.shard.pattern }}\"\n echo \"Running tests matching: $RUN_PATTERN\"\n \n gotestsum --format github-actions -- \\\n -tags integration \\\n -parallel $TEST_PARALLEL \\\n -timeout 30m \\\n -run \"TestIntegration/($RUN_PATTERN)/\" \\\n ./integration-tests/... &\n wait $!\n\n# =============================================================================\n# Gate Job - Single required check for branch protection\n# =============================================================================\n\n ci-complete:\n name: CI Complete\n needs:\n - changes\n - version-check\n - build-cog\n - build-sdk\n - build-rust\n - fmt-go\n - fmt-rust\n - fmt-python\n - check-llm-docs\n - lint-go\n - lint-rust\n - lint-rust-deny\n - lint-python\n - test-go\n - test-rust\n - test-python\n - test-coglet-python\n - integration-shards\n - test-integration\n if: always()\n runs-on: ubuntu-latest\n timeout-minutes: 5\n steps:\n - name: Check job results\n run: |\n echo \"Job results:\"\n echo \" changes: ${{ needs.changes.result }}\"\n echo \" build-sdk: ${{ needs.build-sdk.result }}\"\n echo \" build-rust: ${{ needs.build-rust.result }}\"\n echo \" fmt-go: ${{ needs.fmt-go.result }}\"\n echo \" fmt-rust: ${{ needs.fmt-rust.result }}\"\n echo \" fmt-python: ${{ needs.fmt-python.result }}\"\n echo \" check-llm-docs: ${{ needs.check-llm-docs.result }}\"\n echo \" lint-go: ${{ needs.lint-go.result }}\"\n echo \" lint-rust: ${{ needs.lint-rust.result }}\"\n echo \" lint-rust-deny: ${{ needs.lint-rust-deny.result }}\"\n echo \" lint-python: ${{ needs.lint-python.result }}\"\n echo \" test-go: ${{ needs.test-go.result }}\"\n echo \" test-rust: ${{ needs.test-rust.result }}\"\n echo \" test-python: ${{ needs.test-python.result }}\"\n echo \" test-coglet-python: ${{ needs.test-coglet-python.result }}\"\n echo \" integration-shards: ${{ needs.integration-shards.result }}\"\n echo \" test-integration: ${{ needs.test-integration.result }}\"\n\n # Fail if any job failed (skipped is OK)\n FAILED=false\n for result in \\\n \"${{ needs.changes.result }}\" \\\n \"${{ needs.build-sdk.result }}\" \\\n \"${{ needs.build-rust.result }}\" \\\n \"${{ needs.fmt-go.result }}\" \\\n \"${{ needs.fmt-rust.result }}\" \\\n \"${{ needs.fmt-python.result }}\" \\\n \"${{ needs.check-llm-docs.result }}\" \\\n \"${{ needs.lint-go.result }}\" \\\n \"${{ needs.lint-rust.result }}\" \\\n \"${{ needs.lint-rust-deny.result }}\" \\\n \"${{ needs.lint-python.result }}\" \\\n \"${{ needs.test-go.result }}\" \\\n \"${{ needs.test-rust.result }}\" \\\n \"${{ needs.test-python.result }}\" \\\n \"${{ needs.test-coglet-python.result }}\" \\\n \"${{ needs.integration-shards.result }}\" \\\n \"${{ needs.test-integration.result }}\"\n do\n if [ \"$result\" = \"failure\" ] || [ \"$result\" = \"cancelled\" ]; then\n FAILED=true\n fi\n done\n\n if [ \"$FAILED\" = \"true\" ]; then\n echo \"::error::Some jobs failed or were cancelled\"\n exit 1\n fi\n\n echo \"All CI checks passed!\"\n\n# =============================================================================\n# Release Validation - Dry-run checks (PRs and main)\n# =============================================================================\n\n release-dry-run:\n name: Release Dry Run\n needs: ci-complete\n if: \"!startsWith(github.ref, 'refs/tags/')\"\n runs-on: ubuntu-latest\n timeout-minutes: 15\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - uses: dtolnay/rust-toolchain@stable\n - uses: Swatinem/rust-cache@v2\n with:\n workspaces: crates -> target\n save-if: ${{ github.ref == 'refs/heads/main' }}\n - name: Check coglet crates.io publish\n run: cargo publish --dry-run -p coglet --manifest-path crates/Cargo.toml\n - uses: mlugg/setup-zig@v2\n with:\n version: 0.15.2\n - uses: goreleaser/goreleaser-action@v7\n with:\n version: '~> v2'\n args: check\n" }, { "path": ".github/workflows/codeql.yml", "content": "# For most projects, this workflow file will not need changing; you simply need\n# to commit it to your repository.\n#\n# You may wish to alter this file to override the set of languages analyzed,\n# or to provide custom queries or build logic.\n#\n# ******** NOTE ********\n# We have attempted to detect the languages in your repository. Please check\n# the `language` matrix defined below to confirm you have the correct set of\n# supported CodeQL languages.\n#\nname: \"CodeQL\"\n\non:\n push:\n branches: [ \"main\" ]\n pull_request:\n # The branches below must be a subset of the branches above\n branches: [ \"main\" ]\n schedule:\n - cron: '37 18 * * 5'\n\njobs:\n analyze:\n name: Analyze\n runs-on: ubuntu-latest\n permissions:\n actions: read\n contents: read\n security-events: write\n\n strategy:\n fail-fast: false\n matrix:\n # CodeQL supports: cpp, csharp, go, java, javascript, python, ruby, rust\n # https://aka.ms/codeql-docs/language-support\n language: ['go', 'python', 'rust']\n\n steps:\n - name: Checkout repository\n uses: actions/checkout@v6\n\n # Initializes the CodeQL tools for scanning.\n - name: Initialize CodeQL\n uses: github/codeql-action/init@v4\n with:\n languages: ${{ matrix.language }}\n # If you wish to specify custom queries, you can do so here or in a config file.\n # By default, queries listed here will override any specified in a config file.\n # Prefix the list here with \"+\" to use these queries and those in the config file.\n\n # Details on CodeQL's query packs refer to : https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs\n # queries: security-extended,security-and-quality\n\n\n # Autobuild attempts to build any compiled languages (C/C++, C#, Go, or Java).\n # If this step fails, then you should remove it and run the build manually (see below)\n - name: Autobuild\n uses: github/codeql-action/autobuild@v4\n\n # ℹ️ Command-line programs to run using the OS shell.\n # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun\n\n # If the Autobuild fails above, remove it and uncomment the following three lines.\n # modify them (or add more) to build your code if your project, please refer to the EXAMPLE below for guidance.\n\n # - run: |\n # echo \"Run, Build Application using script\"\n # ./location_of_script_within_repo/buildscript.sh\n\n - name: Perform CodeQL Analysis\n uses: github/codeql-action/analyze@v4\n with:\n category: \"/language:${{matrix.language}}\"\n" }, { "path": ".github/workflows/docs.yaml", "content": "name: Deploy docs\n\non:\n push:\n branches:\n - main\n\njobs:\n deploy:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v6\n - uses: actions/setup-go@v5\n with:\n go-version: '1.23'\n - uses: actions/setup-python@v6\n with:\n python-version: '3.13'\n\n - name: Generate CLI docs\n run: go run ./tools/gendocs/main.go -o docs/cli.md\n\n - name: Copy top-level docs like README and CONTRIBUTING\n run: |\n sed 's/docs\\///g' README.md > ./docs/README.md\n cp CONTRIBUTING.md ./docs/\n\n - name: Deploy\n run: |\n pip install mkdocs-material\n mkdocs gh-deploy --force\n" }, { "path": ".github/workflows/release-build.yaml", "content": "---\nname: Release Build\n\n# Triggered on version tags to build release artifacts and create a GitHub release.\n#\n# THREE RELEASE TYPES:\n#\n# 1. Stable (v0.17.0) - must be on main\n# - Creates DRAFT release → maintainer publishes → release-publish.yaml → PyPI/crates.io\n#\n# 2. Pre-release (v0.17.0-alpha3, v0.17.0-rc1) - must be on main\n# - Same flow as stable, but marked as pre-release\n#\n# 3. Dev (v0.17.0-dev1) - can be tagged from ANY branch\n# - Creates a published pre-release immediately (no draft, no human approval)\n# - Does NOT publish to PyPI or crates.io\n# - Artifacts (CLI binaries, wheels) attached to the GH release\n#\n# SECURITY:\n# - Stable/pre-release tags verified on main branch\n# - Dev releases have no branch restriction but no registry publishing\n\non:\n push:\n tags: [\"v[0-9]+.[0-9]+.[0-9]+*\"]\n\npermissions:\n contents: write\n\nenv:\n CARGO_TERM_COLOR: always\n\njobs:\n verify-tag:\n name: Verify tag and version\n runs-on: ubuntu-latest\n timeout-minutes: 5\n outputs:\n is_dev: ${{ steps.check.outputs.is_dev }}\n version: ${{ steps.check.outputs.version }}\n pep440: ${{ steps.check.outputs.pep440 }}\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n\n - name: Verify tag and version\n id: check\n run: |\n TAG=\"${{ github.ref_name }}\"\n VERSION=\"${TAG#v}\"\n TAG_COMMIT=\"${{ github.sha }}\"\n\n # Get version from Cargo.toml\n CARGO_VERSION=$(grep '^version = ' crates/Cargo.toml | head -1 | sed 's/version = \"\\(.*\\)\"/\\1/')\n\n echo \"Tag: $TAG\"\n echo \"Tag version: $VERSION\"\n echo \"Cargo.toml version: $CARGO_VERSION\"\n echo \"Commit: $TAG_COMMIT\"\n echo \"\"\n\n # Check Cargo.toml matches tag\n if [[ \"$CARGO_VERSION\" != \"$VERSION\" ]]; then\n echo \"::error::Version mismatch! crates/Cargo.toml has version $CARGO_VERSION but tag is $TAG\"\n echo \"::error::\"\n echo \"::error::To fix: Update crates/Cargo.toml to match, merge to main, then delete this tag and re-tag.\"\n exit 1\n fi\n echo \"✓ Cargo.toml version matches tag\"\n\n # Determine release type\n IS_DEV=\"false\"\n if [[ \"$VERSION\" == *-dev* ]]; then\n IS_DEV=\"true\"\n fi\n echo \"is_dev=$IS_DEV\" >> \"$GITHUB_OUTPUT\"\n echo \"version=$VERSION\" >> \"$GITHUB_OUTPUT\"\n\n # Compute PEP 440 version: v0.17.0-alpha1 -> 0.17.0a1, v0.17.0-dev1 -> 0.17.0.dev1\n PEP440=$(echo \"$VERSION\" | sed -E 's/-alpha\\.?/a/; s/-beta\\.?/b/; s/-rc\\.?/rc/; s/-dev\\.?/.dev/')\n echo \"pep440=$PEP440\" >> \"$GITHUB_OUTPUT\"\n echo \"PEP 440 version: $PEP440\"\n\n # Branch rules\n if [[ \"$IS_DEV\" == \"true\" ]]; then\n echo \"Dev release - no branch restriction\"\n echo \"✓ Dev release, skipping branch check\"\n else\n # Stable and pre-release tags must be on main\n echo \"Stable/pre-release detected, verifying main branch...\"\n git fetch origin main\n if ! git merge-base --is-ancestor \"$TAG_COMMIT\" origin/main; then\n echo \"::error::Release tags must be on the main branch\"\n echo \"::error::Tag commit $TAG_COMMIT is not reachable from origin/main\"\n echo \"::error::\"\n echo \"::error::To fix: Merge to main first, then tag\"\n exit 1\n fi\n echo \"✓ Tag is on main branch\"\n fi\n\n build-sdk:\n name: Build SDK wheel\n needs: verify-tag\n runs-on: ubuntu-latest\n timeout-minutes: 10\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n\n - uses: astral-sh/setup-uv@v6\n with:\n python-version: \"3.13\"\n\n - name: Update coglet version constraint\n run: |\n VERSION=\"${{ needs.verify-tag.outputs.pep440 }}\"\n echo \"Setting coglet constraint to >=$VERSION,<1.0\"\n\n # Update pyproject.toml with lockstep version constraint\n sed -i \"s/coglet>=0\\.1\\.0,<1\\.0/coglet>=$VERSION,<1.0/\" pyproject.toml\n\n # Verify the change took effect\n grep \"coglet>=$VERSION\" pyproject.toml\n\n - name: Build SDK wheel\n run: |\n echo \"Building SDK with version: $SETUPTOOLS_SCM_PRETEND_VERSION\"\n uv build --out-dir dist .\n env:\n SETUPTOOLS_SCM_PRETEND_VERSION: ${{ needs.verify-tag.outputs.pep440 }}\n\n - name: Upload SDK artifacts\n uses: actions/upload-artifact@v6\n with:\n name: sdk-dist\n path: dist/*\n\n build-coglet-wheels:\n name: Build coglet wheel (${{ matrix.target }})\n needs: verify-tag\n runs-on: ${{ matrix.os }}\n timeout-minutes: 20\n strategy:\n fail-fast: false\n matrix:\n include:\n - os: ubuntu-latest\n target: x86_64-unknown-linux-gnu\n artifact-suffix: linux-x64\n manylinux: auto\n zig: true\n - os: ubuntu-latest\n target: aarch64-unknown-linux-gnu\n artifact-suffix: linux-arm64\n manylinux: auto\n zig: true\n - os: macos-14\n target: aarch64-apple-darwin\n artifact-suffix: macos-arm64\n manylinux: \"off\"\n zig: false\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n\n - uses: dtolnay/rust-toolchain@stable\n with:\n targets: ${{ matrix.target }}\n\n - uses: Swatinem/rust-cache@v2\n with:\n workspaces: crates -> target\n key: release-${{ matrix.target }}\n\n - uses: astral-sh/setup-uv@v6\n with:\n python-version: \"3.13\"\n enable-cache: false\n\n - name: Build coglet wheel\n uses: PyO3/maturin-action@v1\n with:\n target: ${{ matrix.target }}\n manylinux: ${{ matrix.manylinux }}\n args: --release --out dist -m crates/coglet-python/Cargo.toml ${{ matrix.zig && '--zig' || '' }}\n\n - name: Upload coglet wheel\n uses: actions/upload-artifact@v6\n with:\n name: coglet-wheel-${{ matrix.artifact-suffix }}\n path: dist/*.whl\n\n create-release:\n name: Create release\n needs: [verify-tag, build-sdk, build-coglet-wheels]\n # macOS arm64 runner: native clang for darwin targets, zig for linux targets.\n # CGo required for go-tree-sitter (static Python schema parser).\n runs-on: macos-14\n timeout-minutes: 30\n env:\n CGO_ENABLED: \"1\"\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n\n - uses: actions/setup-go@v6\n with:\n go-version-file: go.mod\n\n - uses: mlugg/setup-zig@v2\n with:\n version: 0.15.2\n\n - name: Check for existing release\n run: |\n TAG=\"${{ github.ref_name }}\"\n EXISTING=$(gh release view \"$TAG\" --json isDraft,isPrerelease --jq '.' 2>/dev/null || echo \"\")\n if [ -n \"$EXISTING\" ]; then\n echo \"::error::Release for $TAG already exists. Delete it before re-running.\"\n echo \"::error::Run: gh release delete $TAG --yes\"\n exit 1\n fi\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n\n # Goreleaser builds CLI binaries and creates draft release\n - name: Build CLI and create draft release\n uses: goreleaser/goreleaser-action@v7\n with:\n version: '~> v2'\n args: release --clean\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n\n # Append Python wheels to the release\n - name: Download all wheel artifacts\n uses: actions/download-artifact@v8\n with:\n path: artifacts\n\n - name: Upload wheels to release\n run: |\n TAG=\"${{ github.ref_name }}\"\n\n # Collect all wheels\n mkdir -p release-wheels\n cp artifacts/sdk-dist/* release-wheels/\n cp artifacts/coglet-wheel-*/*.whl release-wheels/\n\n # Download goreleaser's checksums.txt and append wheel checksums\n gh release download \"$TAG\" -p checksums.txt -D release-wheels\n cd release-wheels\n shasum -a 256 *.whl *.tar.gz >> checksums.txt\n echo \"Checksums:\"\n cat checksums.txt\n cd ..\n\n echo \"Uploading wheels and updated checksums...\"\n gh release upload \"$TAG\" release-wheels/* --clobber\n\n echo \"Release assets:\"\n gh release view \"$TAG\" --json assets --jq '.assets[].name'\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n\n # For dev releases: immediately publish as pre-release (no draft review)\n # For stable/pre-release: leave as draft for maintainer review\n - name: Publish dev release\n if: needs.verify-tag.outputs.is_dev == 'true'\n run: |\n TAG=\"${{ github.ref_name }}\"\n echo \"Publishing dev release $TAG as pre-release (no draft phase)...\"\n gh release edit \"$TAG\" --draft=false --prerelease\n echo \"✓ Dev release $TAG published as pre-release\"\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n" }, { "path": ".github/workflows/release-publish.yaml", "content": "---\nname: Release Publish\n\n# Publishes packages to PyPI and crates.io when a release is published.\n#\n# For stable releases: publishes to PyPI, crates.io, and updates Homebrew tap.\n# For pre-releases: publishes to PyPI and crates.io only (no Homebrew tap).\n# For dev releases: ALL jobs are skipped. (Dev releases do trigger this workflow\n# when release-build.yaml publishes them, but every job gates on is_dev.)\n#\n# PUBLISH ORDER:\n# 1. coglet -> PyPI (must be first, SDK depends on it)\n# 2. coglet -> crates.io (parallel with coglet PyPI)\n# 3. SDK -> PyPI (after coglet is on PyPI)\n# 4. Homebrew cask (stable only, after all publishing completes)\n#\n# REQUIRED GITHUB CONFIGURATION:\n# 1. Create environments in Settings -> Environments:\n# - \"pypi\": For PyPI publishing (Trusted Publisher)\n# - \"crates-io\": For crates.io publishing\n# - \"homebrew\": For Homebrew tap updates\n#\n# 2. Configure environment protection rules:\n# - Deployment branches: \"Selected branches and tags\"\n# - Add pattern: v* (to restrict to version tags only)\n#\n# 3. crates-io uses Trusted Publishing (OIDC via rust-lang/crates-io-auth-action)\n#\n# 4. Homebrew tap uses the cog-homebrew-tapbot GitHub App (ID: 1232932405)\n# - Secret: COG_HOMEBREW_TAP_PRIVATE_KEY (app's private key)\n\non:\n release:\n types: [published]\n\npermissions:\n contents: read\n id-token: write\n\nenv:\n CARGO_TERM_COLOR: always\n\njobs:\n verify-release:\n name: Verify release tag\n runs-on: ubuntu-latest\n timeout-minutes: 5\n outputs:\n is_dev: ${{ steps.check.outputs.is_dev }}\n is_prerelease: ${{ steps.check.outputs.is_prerelease }}\n version: ${{ steps.check.outputs.version }}\n steps:\n - name: Verify valid release tag\n id: check\n run: |\n TAG=\"${{ github.event.release.tag_name }}\"\n\n # Release must be from a valid version tag\n if [[ ! \"$TAG\" =~ ^v[0-9]+\\.[0-9]+\\.[0-9]+ ]]; then\n echo \"::error::Invalid tag format: $TAG\"\n echo \"::error::Tags must match pattern v*.*.* (e.g., v1.0.0)\"\n exit 1\n fi\n echo \"✓ Valid release tag: $TAG\"\n\n VERSION=\"${TAG#v}\"\n echo \"version=$VERSION\" >> \"$GITHUB_OUTPUT\"\n\n # Classify release type\n IS_DEV=\"false\"\n IS_PRERELEASE=\"false\"\n if [[ \"$VERSION\" == *-dev* ]]; then\n IS_DEV=\"true\"\n echo \"Dev release detected - skipping all publishing\"\n elif [[ \"$VERSION\" == *-* ]]; then\n IS_PRERELEASE=\"true\"\n echo \"Pre-release detected - publishing to PyPI/crates.io (no Homebrew tap)\"\n else\n echo \"Stable release detected - full publishing including Homebrew tap\"\n fi\n echo \"is_dev=$IS_DEV\" >> \"$GITHUB_OUTPUT\"\n echo \"is_prerelease=$IS_PRERELEASE\" >> \"$GITHUB_OUTPUT\"\n\n publish-pypi-coglet:\n name: Publish coglet to PyPI\n needs: verify-release\n if: needs.verify-release.outputs.is_dev != 'true'\n runs-on: ubuntu-latest\n environment: pypi\n timeout-minutes: 10\n steps:\n - name: Download coglet wheels from release\n run: |\n mkdir -p dist\n gh release download \"$TAG\" -p \"coglet-*.whl\" -D dist -R \"${{ github.repository }}\"\n env:\n TAG: ${{ github.event.release.tag_name }}\n GH_TOKEN: ${{ github.token }}\n\n - name: Publish to PyPI\n uses: pypa/gh-action-pypi-publish@release/v1\n\n publish-crates-io:\n name: Publish coglet to crates.io\n needs: verify-release\n if: needs.verify-release.outputs.is_dev != 'true'\n runs-on: ubuntu-latest\n environment: crates-io\n timeout-minutes: 15\n permissions:\n contents: read\n id-token: write\n steps:\n - uses: actions/checkout@v6\n with:\n ref: ${{ github.event.release.tag_name }}\n\n - uses: dtolnay/rust-toolchain@stable\n\n - uses: rust-lang/crates-io-auth-action@v1\n id: auth\n\n - name: Publish to crates.io\n run: cargo publish -p coglet --manifest-path crates/Cargo.toml\n env:\n CARGO_REGISTRY_TOKEN: ${{ steps.auth.outputs.token }}\n\n publish-pypi-sdk:\n name: Publish SDK to PyPI\n needs: [verify-release, publish-pypi-coglet]\n if: needs.verify-release.outputs.is_dev != 'true'\n runs-on: ubuntu-latest\n environment: pypi\n timeout-minutes: 10\n steps:\n - name: Download SDK artifacts from release\n run: |\n mkdir -p dist\n gh release download \"$TAG\" -p \"cog-*.whl\" -D dist -R \"${{ github.repository }}\"\n gh release download \"$TAG\" -p \"cog-*.tar.gz\" -D dist -R \"${{ github.repository }}\"\n env:\n TAG: ${{ github.event.release.tag_name }}\n GH_TOKEN: ${{ github.token }}\n\n - name: Publish to PyPI\n uses: pypa/gh-action-pypi-publish@release/v1\n\n update-homebrew-tap:\n name: Update Homebrew cask\n needs: [verify-release, publish-pypi-sdk, publish-crates-io]\n # Stable releases only — no dev, no pre-release\n if: >-\n needs.verify-release.outputs.is_dev != 'true' &&\n needs.verify-release.outputs.is_prerelease != 'true'\n runs-on: ubuntu-latest\n environment: homebrew\n timeout-minutes: 10\n steps:\n - name: Generate GitHub App token\n id: app-token\n uses: actions/create-github-app-token@v2\n with:\n app-id: 1232932405\n private-key: ${{ secrets.COG_HOMEBREW_TAP_PRIVATE_KEY }}\n owner: replicate\n repositories: homebrew-tap\n\n - name: Download checksums from release\n run: gh release download \"$TAG\" -p checksums.txt -R \"${{ github.repository }}\"\n env:\n TAG: ${{ github.event.release.tag_name }}\n GH_TOKEN: ${{ github.token }}\n\n - name: Generate and push cask\n env:\n GH_TOKEN: ${{ steps.app-token.outputs.token }}\n TAG: ${{ github.event.release.tag_name }}\n VERSION: ${{ needs.verify-release.outputs.version }}\n run: |\n # Extract SHA256s for Darwin binaries from checksums.txt\n SHA_X86=$(grep 'cog_Darwin_x86_64' checksums.txt | awk '{print $1}')\n SHA_ARM=$(grep 'cog_Darwin_arm64' checksums.txt | awk '{print $1}')\n\n if [ -z \"$SHA_X86\" ] || [ -z \"$SHA_ARM\" ]; then\n echo \"::error::Missing Darwin binary checksums in checksums.txt\"\n echo \"Darwin x86_64: ${SHA_X86:-MISSING}\"\n echo \"Darwin arm64: ${SHA_ARM:-MISSING}\"\n cat checksums.txt\n exit 1\n fi\n\n echo \"Checksums:\"\n echo \" Darwin x86_64: $SHA_X86\"\n echo \" Darwin arm64: $SHA_ARM\"\n\n BASE_URL=\"https://github.com/replicate/cog/releases/download/${TAG}\"\n\n # Generate cask file (no indentation to avoid stripping issues)\n cat > cog.rb <-\n {{ .Binary }}_\n {{- title .Os }}_\n {{- if eq .Arch \"amd64\" }}x86_64\n {{- else if eq .Arch \"386\" }}i386\n {{- else }}{{ .Arch }}{{ end -}}\nchecksum:\n name_template: \"checksums.txt\"\nsnapshot:\n version_template: '{{ envOrDefault \"COG_VERSION\" (printf \"%s-dev+g%s\" (incpatch .Version) .ShortCommit) }}'\nchangelog:\n sort: asc\n filters:\n exclude:\n - \"^docs:\"\n - \"^test:\"\nrelease:\n draft: true\n # If set to auto, will mark the release as not ready for production\n # in case there is an indicator for this in the tag e.g. v1.0.0-alpha\n # If set to true, will mark the release as not ready for production.\n # Default is false.\n prerelease: auto\n" }, { "path": ".mockery.yml", "content": "all: false\ndir: '{{.InterfaceDir}}'\nfilename: mocks_test.go\nforce-file-write: true\nformatter: goimports\nlog-level: info\nstructname: '{{.Mock}}{{.InterfaceName}}'\npkgname: '{{.SrcPackageName}}'\nrecursive: false\nrequire-template-schema-exists: true\ntemplate: testify\ntemplate-schema: '{{.Template}}.schema.json'\npackages:\n github.com/replicate/cog/pkg/docker/command:\n config:\n all: true\n dir: \"pkg/docker/dockertest\"\n filename: \"command_mocks.go\"\n pkgname: \"dockertest\"\n structname: \"{{.Mock}}{{.InterfaceName}}2\"\n" }, { "path": ".vscode/extensions.json", "content": "{\n \"recommendations\": [\n \"charliermarsh.ruff\",\n \"golang.go\",\n \"ms-python.python\",\n \"ms-python.vscode-pylance\"\n ]\n}\n" }, { "path": ".vscode/settings.json", "content": "{\n \"editor.formatOnSave\": true,\n \"editor.formatOnType\": true,\n \"editor.formatOnPaste\": true,\n \"editor.renderControlCharacters\": true,\n \"editor.suggest.localityBonus\": true,\n \"files.insertFinalNewline\": true,\n \"files.trimFinalNewlines\": true,\n \"[go]\": {\n \"editor.defaultFormatter\": \"golang.go\"\n },\n \"go.coverOnTestPackage\": false,\n \"go.lintTool\": \"golangci-lint\",\n \"go.formatTool\": \"goimports\",\n \"go.testOnSave\": true,\n \"gopls\": {\n \"formatting.local\": \"github.com/replicate/cog\"\n },\n \"[json]\": {\n \"editor.defaultFormatter\": \"vscode.json-language-features\"\n },\n \"[jsonc]\": {\n \"editor.defaultFormatter\": \"vscode.json-language-features\"\n },\n \"[python]\": {\n \"editor.formatOnSave\": true,\n \"editor.codeActionsOnSave\": {\n \"source.fixAll\": \"explicit\",\n \"source.organizeImports\": \"explicit\"\n },\n \"editor.defaultFormatter\": \"charliermarsh.ruff\"\n },\n \"python.languageServer\": \"Pylance\",\n \"python.testing.pytestArgs\": [\n \"-vvv\",\n \"python\"\n ],\n \"python.testing.unittestEnabled\": false,\n \"python.testing.pytestEnabled\": true,\n}\n" }, { "path": "AGENTS.md", "content": "# AGENTS.md\n\nThis file provides guidance to coding agents when working with code in this repository.\n\n## Project Overview\n\nCog is a tool that packages machine learning models in production-ready containers. \n\nIt consists of:\n- **Cog CLI** (`cmd/cog/`) - Command-line interface for building, running, and deploying models, written in Go\n- **Python SDK** (`python/cog/`) - Python library for defining model predictors and training in Python\n- **Coglet** (`crates/`) - Rust-based prediction server that runs inside containers, with Python bindings via PyO3\n\nDocumentation for the CLI and SDK is available by reading ./docs/llms.txt.\n\n## Development Commands\n\nDevelopment tasks are managed with [mise](https://mise.jdx.dev/). Run `mise tasks` to see all available tasks.\n\n### Quick Reference\n\n| Task | Description |\n|------|-------------|\n| `mise run fmt` | Check formatting (all languages) |\n| `mise run fmt:fix` | Fix formatting (all languages) |\n| `mise run lint` | Run linters (all languages) |\n| `mise run lint:fix` | Fix lint issues (all languages) |\n| `mise run test:go` | Run Go tests |\n| `mise run test:rust` | Run Rust tests |\n| `mise run test:python` | Run Python tests |\n| `mise run test:integration` | Run integration tests |\n| `mise run build:cog` | Build cog CLI binary |\n| `mise run build:coglet` | Build coglet wheel (dev) |\n| `mise run build:sdk` | Build SDK wheel |\n| `mise run install` | Build and symlink cog to /usr/local/bin |\n| `mise run docs:llm` | **IMPORTANT:** Regenerate `docs/llms.txt` after editing docs |\n| `mise run docs:cli` | Generate CLI reference docs from Go source code |\n\n### Task Naming Convention\n\nTasks follow a consistent naming pattern:\n\n- **Language-based tasks** for fmt/lint/test/typecheck: `task:go`, `task:rust`, `task:python`\n- **Component-based tasks** for build: `build:cog`, `build:coglet`, `build:sdk`\n- **Check vs Fix**: `fmt` and `lint` default to check mode (non-destructive); use `:fix` suffix to auto-fix\n\n### All Tasks by Category\n\n**Format:**\n- `mise run fmt` / `mise run fmt:check` - Check all (alias)\n- `mise run fmt:fix` - Fix all\n- `mise run fmt:go` / `mise run fmt:rust` / `mise run fmt:python` - Per-language\n\n**Lint:**\n- `mise run lint` / `mise run lint:check` - Check all (alias)\n- `mise run lint:fix` - Fix all\n- `mise run lint:go` / `mise run lint:rust` / `mise run lint:python` - Per-language\n- `mise run lint:rust:deny` - Check Rust licenses/advisories\n\n**Test:**\n- `mise run test:go` - Go unit tests\n- `mise run test:rust` - Rust unit tests\n- `mise run test:python` - Python unit tests (via tox)\n- `mise run test:coglet:python` - Coglet Python binding tests\n- `mise run test:integration` - Integration tests\n\n**Build:**\n- `mise run build:cog` - Build cog CLI (development)\n- `mise run build:cog:release` - Build cog CLI (release)\n- `mise run build:coglet` - Build coglet wheel (dev install)\n- `mise run build:coglet:wheel` - Build coglet wheel (native platform)\n- `mise run build:coglet:wheel:linux-x64` - Build for Linux x86_64\n- `mise run build:coglet:wheel:linux-arm64` - Build for Linux ARM64\n- `mise run build:sdk` - Build SDK wheel\n\n**Install:**\n- `mise run install` - Symlink cog CLI to `/usr/local/bin` (requires `build:cog` first)\n- `PREFIX=/custom/path mise run install` - Symlink to custom location\n\n**Other:**\n- `mise run typecheck` - Type check all languages\n- `mise run generate` - Run code generation\n- `mise run clean` - Clean all build artifacts\n- `mise run docs` - Build documentation\n- `mise run docs:serve` - Serve docs locally\n\n## Code Style Guidelines\n\n### Go\n- **Imports**: Organize in three groups separated by blank lines: (1) Standard library, (2) Third-party packages, (3) Internal packages (`github.com/replicate/cog/pkg/...`)\n- **Formatting**: Use `mise run fmt:go:fix`\n- **Linting**: Must pass golangci-lint with: errcheck, gocritic, gosec, govet, ineffassign, misspell, revive, staticcheck, unused\n- **Error Handling**: Return errors as values; use `pkg/errors.CodedError` for user-facing errors with error codes\n- **Naming**: CamelCase for exported, camelCase for unexported\n- **Testing**: Use `testify/require` for assertions; prefer table-driven tests\n\nExample import block:\n```go\nimport (\n \"fmt\"\n \n \"github.com/spf13/cobra\"\n \n \"github.com/replicate/cog/pkg/config\"\n)\n```\n\n### Python\n- **Imports**: Automatically organized by ruff/isort (stdlib → third-party → local)\n- **Formatting**: Use `mise run fmt:python:fix`\n- **Linting**: Must pass ruff checks: E (pycodestyle), F (Pyflakes), I (isort), W (warnings), S (bandit), B (bugbear), ANN (annotations)\n- **Type Annotations**: Required on all function signatures; use `typing_extensions` for compatibility; avoid `Any` where possible\n- **Error Handling**: Raise exceptions with descriptive messages; avoid generic exception catching\n- **Naming**: snake_case for functions/variables/modules, PascalCase for classes\n- **Testing**: Use pytest with fixtures; async tests with pytest-asyncio\n- **Compatibility**: Must support Python 3.10-3.13\n\n### Rust\n- **Formatting**: Use `mise run fmt:rust:fix`\n- **Linting**: Must pass `mise run lint:rust` (clippy)\n- **Dependencies**: Audited with `cargo-deny` (see `crates/deny.toml`); run `mise run lint:rust:deny`\n- **Error Handling**: Use `thiserror` for typed errors, `anyhow` for application errors\n- **Naming**: snake_case for functions/variables, PascalCase for types\n- **Testing**: Use `cargo test`; snapshot tests use `insta`\n- **Async**: tokio runtime; async/await patterns\n\n## Working on the CLI and support tooling\nThe CLI code is in the `cmd/cog/` and `pkg/` directories. Support tooling is in the `tools/` directory. \n\nThe main commands for working on the CLI are:\n- `go run ./cmd/cog` - Runs the Cog CLI directly from source (requires wheel to be built first)\n- `mise run build:cog` - Builds the Cog CLI binary\n- `mise 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 install`\n- `mise run test:go` - Runs all Go unit tests\n- `go test ./pkg/...` - Runs tests directly with `go test`\n\n## Working on the Python SDK\nThe Python SDK is developed in the `python/cog/` directory. It uses `uv` for virtual environments and `tox` for testing across multiple Python versions.\n\nThe main commands for working on the SDK are:\n- `mise run build:sdk` - Builds the Python wheel\n- `mise run test:python` - Runs Python tests across all supported versions\n\n## Working on Coglet (Rust)\nCoglet is the Rust-based prediction server that runs inside Cog containers, handling HTTP requests, worker process management, and prediction execution.\n\nThe code is in the `crates/` directory:\n- `crates/coglet/` - Core Rust library (HTTP server, worker orchestration, IPC)\n- `crates/coglet-python/` - PyO3 bindings for Python predictor integration (requires Python 3.10+)\n\nFor detailed architecture documentation, see `crates/README.md` and `crates/coglet/README.md`.\n\nThe main commands for working on Coglet are:\n- `mise run build:coglet` - Build and install coglet wheel for development (macOS, for local Rust/Python tests)\n- `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`)\n- `mise run test:rust` - Run Rust unit tests\n- `mise run lint:rust` - Run clippy linter\n- `mise run fmt:rust:fix` - Format code\n\n### Testing\nGo code is tested using the built-in `go test` framework:\n- `go test ./pkg/... -run ` - Runs specific Go tests by name\n- `mise run test:go` - Runs all Go unit tests\n\nPython code is tested using `tox`, which allows testing across multiple Python versions and configurations:\n- `mise run test:python` - Runs all Python unit tests\n- `uv run tox -e py312-tests -- python/tests/server/test_http.py::test_openapi_specification_with_yield` - Runs a specific Python test\n\nThe integration test suite in `integration-tests/` tests the end-to-end functionality of the Cog CLI and Python SDK using Go's testscript framework:\n- `mise run test:integration` - Runs the integration tests\n- `mise run test:integration string_predictor` - Runs a specific integration test\n\nThe 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:\n```bash\nmise run build:cog\nCOG_BINARY=dist/go/*/cog mise run test:integration\n```\n\n### Development Workflow\n1. Run `mise install` to set up the development environment\n2. Run `mise run build:sdk` after making changes to the `./python` directory\n3. Run `mise run build:coglet:wheel:linux-x64` after making changes to the `./crates` directory (needed for Docker testing)\n4. Run `mise run build:cog` to build the CLI (wheels are picked up from `dist/` at Docker build time, not embedded in the binary)\n5. Run `mise run fmt:fix` to format code\n6. Run `mise run lint` to check code quality\n7. Run `mise run docs:llm` to regenerate `docs/llms.txt` after changing `README.md` or any `docs/*.md` file\n8. Read the `./docs` directory and make sure the documentation is up to date\n\n**IMPORTANT:** Always run `mise run lint` (or the language-specific variant, e.g. `mise run lint:go`) before committing to catch linter errors early. CI will reject PRs that fail lint checks.\n\n## Architecture\n\n### CLI Architecture (Go)\nThe CLI follows a command pattern with subcommands. The main components are:\n- `pkg/cli/` - Command definitions (build, run, predict, serve, etc.)\n- `pkg/docker/` - Docker client and container management\n- `pkg/dockerfile/` - Dockerfile generation and templating\n- `pkg/config/` - cog.yaml parsing and validation\n- `pkg/image/` - Image building and pushing logic\n\n### Python SDK Architecture\n- `python/cog/` - Core SDK\n - `base_predictor.py` - Base class for model predictors\n - `types.py` - Input/output type definitions\n - `server/` - HTTP/queue server implementation\n - `command/` - Runner implementations for predict/train\n\n### Coglet Architecture (Rust)\nThe 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).\n\nSee `crates/README.md` for detailed architecture documentation.\n\n- `crates/coglet/` - Core Rust library (HTTP server, worker orchestration, IPC bridge)\n- `crates/coglet-python/` - PyO3 bindings for Python predictor integration\n\n### Key Design Patterns\n1. **Local Wheel Resolution**: The CLI discovers SDK and coglet wheels from `dist/` at Docker build time (not embedded in the binary)\n2. **Docker SDK Integration**: Uses Docker Go SDK for container operations\n3. **Type Safety**: Dataclasses for Python type validation, strongly typed Go interfaces\n4. **Compatibility Matrix**: Automated CUDA/PyTorch/TensorFlow compatibility management\n\nFor comprehensive architecture documentation, see [`architecture/`](./architecture/00-overview.md).\n\n## Common Tasks\n\n### Adding a new CLI command\n1. Create command file in `pkg/cli/`\n2. Add command to `pkg/cli/root.go`\n3. Implement business logic in appropriate `pkg/` subdirectory\n4. Add tests\n\n### Modifying Python SDK behavior\n1. Edit files in `python/cog/`\n2. Run `mise run build:sdk` to rebuild wheel\n3. Test with `mise run test:python`\n4. Integration test with `mise run test:integration`\n\n### Updating ML framework compatibility\n1. See `tools/compatgen/` for compatibility matrix generation\n2. Update framework versions in relevant Dockerfile templates\n3. Test with various framework combinations\n\n### Updating the docs\n- Documentation is in the `docs/` directory, written in Markdown and generated into HTML using `mkdocs`.\n- **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.\n- **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.\n\n## CI Tool Dependencies\n\nDevelopment tools are managed in **two places** that must be kept in sync:\n\n1. **`mise.toml`** — Tool versions for local development (uses aqua backend for prebuilt binaries)\n2. **`.github/workflows/ci.yaml`** — Tool installation for CI (uses dedicated GitHub Actions)\n\nCI deliberately avoids aqua downloads from GitHub Releases to prevent transient 502 failures. Instead, it uses:\n\n| Tool | CI installation method | Why |\n|------|----------------------|-----|\n| gotestsum | `go install` | Uses Go module proxy, not GitHub Releases |\n| cargo-deny | `taiki-e/install-action` | Prebuilt with checksum verification |\n| cargo-nextest | `taiki-e/install-action` | Prebuilt with checksum verification |\n| coglet wheel (maturin+zig) | `PyO3/maturin-action` | Bundles maturin and zig |\n| golangci-lint | `golangci/golangci-lint-action` | Built-in caching |\n| Rust toolchain | `dtolnay/rust-toolchain` | Guaranteed ordering |\n\nTools disabled in CI are listed in `MISE_DISABLE_TOOLS` in `ci.yaml`.\n\n**When updating a tool version**, update both:\n- The version in `mise.toml` (for local dev)\n- The corresponding version pin in `.github/workflows/ci.yaml` (for CI)\n\n## Important Files\n- `cog.yaml` - User-facing model configuration\n- `pkg/config/config.go` - Go code for parsing and validating `cog.yaml`\n- `pkg/config/data/config_schema_v1.0.json` - JSON schema for `cog.yaml`\n- `python/cog/base_predictor.py` - Predictor interface\n- `crates/Cargo.toml` - Rust workspace configuration\n- `crates/README.md` - Coglet architecture overview\n- `mise.toml` - Task definitions for development workflow\n\n## Testing Philosophy\n- Unit tests for individual components (Go and Python)\n- Integration tests for end-to-end workflows\n- Tests use real Docker operations (no mocking Docker API)\n- Always run `mise run build:sdk` after making Python changes before testing Go code\n- Python 3.10-3.13 compatibility is required\n\n### Go Test Conventions\nAll Go tests must use [testify](https://github.com/stretchr/testify) for assertions. Do **not** use raw `if` checks with `t.Fatal`/`t.Errorf` — use `require` and `assert` instead.\n\n- **`require`** — for fatal assertions that should stop the test (setup failures, preconditions):\n ```go\n require.NoError(t, err, \"failed to create client\")\n require.Equal(t, expected, actual)\n require.True(t, condition, \"server should be ready\")\n ```\n- **`assert`** — for non-fatal checks where the test should continue (e.g. validating multiple fields in a loop):\n ```go\n assert.Equal(t, http.StatusOK, resp.StatusCode)\n assert.Contains(t, output, \"expected substring\")\n assert.NoError(t, err, \"prediction %d failed\", i)\n ```\n- Use `require` for errors in setup/teardown and `assert` for the actual test expectations\n- Prefer specific assertions (`Equal`, `Contains`, `NoError`, `Len`, `Less`) over generic `True`/`False` — they produce better failure messages\n- Prefer table-driven tests for testing multiple similar cases\n" }, { "path": "CONTRIBUTING.md", "content": "# Contributing guide\n\n## Development environment\n\nDevelopment tasks are managed with [mise](https://mise.jdx.dev/). Run `mise tasks` to see all available tasks.\n\n### Prerequisites\n\n- [mise](https://mise.jdx.dev/getting-started.html): Manages Go, Rust, Python, and other tools\n- [Docker](https://docs.docker.com/desktop) or [OrbStack](https://orbstack.dev)\n\n### Setup\n\n```sh\n# Trust the mise configuration and install tools\nmise trust\nmise install\n\n# Create Python virtualenv and install dependencies\nuv venv\nuv sync --all-groups\n```\n\n### Build & install\n\n```sh\nmise run build\n\n# symlink the binary to /usr/local/bin\nsudo mise run install \n```\n\nAfter making changes, run `mise run build` to rebuild and it will get picked up by the symlink.\n\n### Common tasks\n\n```sh\n# Run all tests\nmise run test:go\nmise run test:python\nmise run test:rust\n\n# Run specific tests\nmise run test:go -- ./pkg/config\nuv run tox -e py312-tests -- python/tests/server/test_http.py -k test_name\n\n# Format code (all languages)\nmise run fmt:fix\n\n# Lint code (all languages)\nmise run lint\n```\n\nRun `mise tasks` for the complete list of available tasks.\n\nIf you encounter any errors, see the troubleshooting section below.\n\n## Project structure\n\nAs much as possible, this is attempting to follow the [Standard Go Project Layout](https://github.com/golang-standards/project-layout).\n\n- `cmd/` - The root `cog` command.\n- `pkg/cli/` - CLI commands.\n- `pkg/config` - Everything `cog.yaml` related.\n- `pkg/docker/` - Low-level interface for Docker commands.\n- `pkg/dockerfile/` - Creates Dockerfiles.\n- `pkg/image/` - Creates and manipulates Cog Docker images.\n- `pkg/predict/` - Runs predictions on models.\n- `pkg/util/` - Various packages that aren't part of Cog. They could reasonably be separate re-usable projects.\n- `python/` - The Cog Python library.\n- `integration-tests/` - Go-based integration tests using testscript.\n- `tools/compatgen/` - Tool for generating CUDA/PyTorch/TensorFlow compatibility matrices.\n\nFor deeper architectural understanding, see the [architecture documentation](./architecture/00-overview.md).\n\n## Updating compatibility matrices\n\nThe CUDA base images and framework compatibility matrices in `pkg/config/` are checked into source control and only need to be regenerated when adding support for new versions of CUDA, PyTorch, or TensorFlow.\n\nTo regenerate the compatibility matrices, run:\n\n```sh\n# Regenerate all matrices\nmise run generate:compat\n\n# Or regenerate specific matrices\nmise run generate:compat cuda\nmise run generate:compat torch\nmise run generate:compat tensorflow\n```\n\nThe generated files are:\n- `pkg/config/cuda_base_images.json` - Available NVIDIA CUDA base images\n- `pkg/config/torch_compatibility_matrix.json` - PyTorch/CUDA/Python compatibility\n- `pkg/config/tf_compatibility_matrix.json` - TensorFlow/CUDA/Python compatibility\n\n## CI tool dependencies\n\nDevelopment tools are managed in **two places** that must be kept in sync:\n\n1. **`mise.toml`** — Tool versions for local development (uses aqua backend for prebuilt binaries)\n2. **`.github/workflows/ci.yaml`** — Tool installation for CI (uses dedicated GitHub Actions)\n\nCI deliberately avoids aqua downloads from GitHub Releases to prevent transient 502 failures. Instead, it uses dedicated actions (`taiki-e/install-action`, `go install`, `PyO3/maturin-action`, etc.) that are more reliable.\n\nTools disabled in CI are listed in `MISE_DISABLE_TOOLS` in `ci.yaml`.\n\n**When updating a tool version**, update both:\n- The version in `mise.toml` (for local dev)\n- The corresponding version pin in `.github/workflows/ci.yaml` (for CI)\n\nSee the [CI Tool Dependencies section in AGENTS.md](./AGENTS.md#ci-tool-dependencies) for the full mapping of tools to their CI installation methods.\n\n## Concepts\n\nThere are a few concepts used throughout Cog that might be helpful to understand.\n\n- **Config**: The `cog.yaml` file.\n- **Image**: Represents a built Docker image that serves the Cog API, containing a **model**.\n- **Input**: Input from a **prediction**, as key/value JSON object.\n- **Model**: A user's machine learning model, consisting of code and weights.\n- **Output**: Output from a **prediction**, as arbitrarily complex JSON object.\n- **Prediction**: A single run of the model, that takes **input** and produces **output**.\n- **Predictor**: Defines how Cog runs **predictions** on a **model**.\n\n## Running tests\n\n**To run the entire test suite:**\n\n```sh\nmise run test:go\nmise run test:python\nmise run test:rust\n```\n\n**To run just the Go unit tests:**\n\n```sh\nmise run test:go\n```\n\n**To run just the Python tests:**\n\n```sh\nmise run test:python\n```\n\n> [!INFO]\n> This runs the Python test suite across all supported Python versions (3.10-3.13) using tox.\n\n### Integration Tests\n\nIntegration tests are in `integration-tests/` using [testscript](https://pkg.go.dev/github.com/rogpeppe/go-internal/testscript). Each test is a self-contained `.txtar` file in `integration-tests/tests/`, with some specialized tests as Go test functions in subpackages.\n\n```sh\n# Run all integration tests\nmise run test:integration\n\n# Run a specific test\nmise run test:integration string_predictor\n\n# Run fast tests only (skip slow GPU/framework tests)\ncd integration-tests && go test -short -v\n\n# Run with a custom cog binary\nCOG_BINARY=/path/to/cog mise run test:integration\n```\n\n### Writing Integration Tests\n\nWhen adding new functionality, add integration tests in `integration-tests/tests/`. They are:\n- Self-contained (embedded fixtures in `.txtar` files)\n- Faster to run (parallel execution with automatic cleanup)\n- Easier to read and write (simple command script format)\n\nExample test structure:\n\n```txtar\n# Test string predictor\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def predict(self, s: str) -> str:\n return \"hello \" + s\n```\n\nFor testing `cog serve`, use `cog serve` and the `curl` command:\n\n```txtar\ncog build -t $TEST_IMAGE\ncog serve\ncurl POST /predictions '{\"input\":{\"s\":\"test\"}}'\nstdout '\"output\":\"hello test\"'\n```\n\n#### Advanced Test Commands\n\nFor tests that require subprocess initialization or async operations, use `retry-curl`:\n\n**`retry-curl` - HTTP request with automatic retries:**\n\n```txtar\n# Make HTTP request with retry logic (useful for subprocess initialization delays)\n# retry-curl [method] [path] [body] [max-attempts] [retry-delay]\nretry-curl POST /predictions '{\"input\":{\"s\":\"test\"}}' 30 1s\nstdout '\"output\":\"hello test\"'\n```\n\n**Example: Testing predictor with subprocess in setup**\n\n```txtar\ncog build -t $TEST_IMAGE\ncog serve\n\n# Use generous retries since setup spawns a background process\nretry-curl POST /predictions '{\"input\":{\"s\":\"test\"}}' 30 1s\nstdout '\"output\":\"hello test\"'\n\n-- predict.py --\nclass Predictor(BasePredictor):\n def setup(self):\n self.process = subprocess.Popen([\"./background.sh\"])\n \n def predict(self, s: str) -> str:\n return \"hello \" + s\n```\n\n#### Test Conditions\n\nUse conditions to control when tests run based on environment:\n\n**`[short]` - Skip slow tests in short mode:**\n\n```txtar\n[short] skip 'requires GPU or long build time'\n\ncog build -t $TEST_IMAGE\n# ... rest of test\n```\n\nRun with `go test -short` to skip these tests.\n\n**`[linux]` / `[!linux]` - Platform-specific tests:**\n\n```txtar\n[!linux] skip 'requires Linux'\n\n# Linux-specific test\ncog build -t $TEST_IMAGE\n```\n\n**`[amd64]` / `[!amd64]` - Architecture-specific tests:**\n\n```txtar\n[!amd64] skip 'requires amd64 architecture'\n\n# amd64-specific test\ncog build -t $TEST_IMAGE\n```\n\n**`[linux_amd64]` - Combined platform and architecture:**\n\n```txtar\n[!linux_amd64] skip 'requires Linux on amd64'\n\n# Test that requires both Linux and amd64\ncog build -t $TEST_IMAGE\n```\n\n**Combining conditions:**\n\nConditions can be negated with `!`. Examples:\n- `[short]` - True when `go test -short` is used (skip this test in short mode)\n- `[!short]` - True when NOT running with `-short` flag (only run this in full test mode)\n- `[!linux]` - True when NOT on Linux\n- `[linux_amd64]` - True when on Linux AND amd64\n\nSee existing tests in `integration-tests/tests/`, especially `setup_subprocess_*.txtar`, for more examples.\n\n## Running the docs server\n\nTo run the docs website server locally:\n\n```sh\nmise run docs:serve\n```\n\n## Publishing a release\n\nReleases are managed by GitHub Actions workflows. See [`.github/workflows/README.md`](.github/workflows/README.md) for full details.\n\nAll packages use **lockstep versioning** from `crates/Cargo.toml`. There are three release types:\n\n| Type | Example tag | Branch rule | PyPI/crates.io? |\n|------|-------------|-------------|-----------------|\n| **Stable** | `v0.17.0` | Must be on main | Yes |\n| **Pre-release** | `v0.17.0-alpha3` | Must be on main | Yes |\n| **Dev** | `v0.17.0-dev1` | Any branch | No |\n\n### Stable / Pre-release\n\n```bash\n# 1. Update crates/Cargo.toml version (e.g. \"0.17.0\" or \"0.17.0-alpha3\")\n# 2. Merge to main\n# 3. Tag and push\ngit tag v0.17.0\ngit push origin v0.17.0\n# 4. Wait for release-build.yaml to create a draft release\n# 5. Review the draft in GitHub UI, then click \"Publish release\"\n# This triggers release-publish.yaml -> PyPI + crates.io\n```\n\n### Dev release\n\n```bash\n# From any branch:\n# 1. Update crates/Cargo.toml version (e.g. \"0.17.0-dev1\")\n# 2. Commit and push\n# 3. Tag and push\ngit tag v0.17.0-dev1\ngit push origin v0.17.0-dev1\n# 4. Done. Artifacts are built and published as a GH pre-release.\n# No PyPI/crates.io. No manual approval.\n```\n\n\n## Troubleshooting\n\n### `cog command not found`\n\nThe compiled `cog` binary will be installed in `$GOPATH/bin/cog`, e.g. `~/go/bin/cog`. Make sure that Golang's bin directory is present on your system PATH by adding it to your shell config (`.bashrc`, `.zshrc`, etc):\n\n export PATH=~/go/bin:$PATH\n\n---\n\nStill having trouble? Please [open an issue](https://github.com/replicate/cog/issues) on GitHub.\n" }, { "path": "DESIGN.md", "content": "# Design\n\n## Background\n\nCog came from Andreas's experience at Spotify and Ben's experience at Docker.\n\nAt Spotify, Andreas noticed a cluster of related problems:\n\n- **It was hard to run open-source machine learning models.** All the advances in machine learning were locked up inside prose in PDFs, scraps of code on GitHub, weights on Google Drive (if you were lucky!). If you wanted to build upon this research, or apply it to real-world problems, you had to implement it all from scratch.\n- **It was hard to deploy machine learning models to production.** Andreas was the only person on the research team who was also an infrastructure engineer. Typically a researcher would have to sit down with Andreas to decide on an API, get a server written, package up dependencies, battle CUDA, get it running efficiently, get it deployed on the cluster, and so on and so forth. It would take weeks to get something running in production.\n\nBen connected this back to his experience at Docker. What Docker did was define a standard box that software could go in. You could put any kind of server software in there – Python, Java, Ruby on Rails, whatever – and you could then know that you could run it on your local machine or on any cloud, as long as it supported Docker. We wanted to do the same thing for machine learning.\n\n## Vision\n\nWe want Cog to be a standard artifact for what a model is and how that model is run.\n\n(More detail...)\n\n## Design principles\n\nThere are a few things driving Cog's design:\n\n- **Reproducible artifact**: When you've put your model in Cog, it'll run anywhere, and _keep_ on running. This is why it's a Docker image, with all of the model's dependencies. Docker images have a content-addressable SHA256 ID, which is the identifier for that model's behavior, byte-for-byte. \n- **Weights inside the image**: We encourage users to put model weights in images. If the weights are on cloud storage somewhere, then they might change or disappear, and the image will produce different results. There's nothing magical about Docker images – they're just a bundle of files. Docker moves around that bundle of files quite slowly, though, but we can optimize that process so it's as fast as reading weights directly from blob storage, or wherever.\n- **Models are just functions**: Models can be lots of things. We are of the opinion that [machine learning is just software](https://replicate.com/blog/machine-learning-needs-better-tools), and a model is just a function. It often needs to be attached to a GPU, but apart from that it's just a normal function that has some input and some output. This is the core difference between Docker's abstraction and Cog's abstraction: Docker packages up an executable, whereas Cog packages up a _function_.\n- **Standard interface**: When you run the Docker container, it serves an HTTP server, that is a standard API for running that function. You can think of it like a remote procedure call.\n- **Self-describing artifact**: A Cog model has it's schema (or type signature, if you're thinking of it as a function) attached to the image as a label. This means systems that work with Cog models can know what the model is and what requests to send to it. This is what powers the forms on Replicate, for example.\n- **Not just the model**: Before Cog, the typical standard packaging formats for machine learning models were at the network level. A way of taking a Tensorflow or PyTorch network and packaging up in a way that would run on lots of different types of accelerators. Things like [ONNX](https://onnx.ai/) or [TVM's IR](https://tvm.apache.org/). We realized that \"models\" are not just the network, but they are also pre- and post-processing, and are so diverse and the field is so fast-moving that you can't possible squeeze it into some high-level abstraction. It just needs to be code running on a computer.\n- **It's just Docker**: Cog models need to run anywhere, and they'll only run anywhere if it's vanilla Docker. We might optimize how Docker images get shipped around to make it faster, but we're not going to invent our own image format.\n- **The API is for software developers**: In the olden days, you have to pass tensors to TFServing and know how to generate a tensor from a JPEG. Cog's API intentionally just speaks JSON, strings, files, etc. It's intended to be the interface between the software developer and the ML engineer. Sort of like Docker was intended to be the interface between the software developer and the infrastructure engineer.\n- **Cog is the APIs and interfaces, not just the software**: The most important thing about Cog is that it defines a standard for what a model is and how to run it. It doesn't necessarily need to involve Cog the piece of software itself. For example, Replicate could serve a model from OpenAI with a Cog API and schema, but it's not packaged or running with Cog under the hood at all – it's just calling the OpenAI API directly.\n" }, { "path": "LICENSE", "content": "\n Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright 2022, Replicate, Inc.\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "Makefile", "content": "# Makefile - Shim that delegates to mise tasks\n#\n# This Makefile provides backward compatibility for common targets.\n# All task definitions live in mise.toml. Run `mise tasks` to see available tasks.\n#\n# For new development, prefer using mise directly:\n# mise run build:cog instead of make cog\n# mise run test:go instead of make test-go\n# mise run fmt:fix instead of make fmt\n\nSHELL := bash\n\n# Show deprecation warning (set MAKE_NO_WARN=1 to suppress)\nifndef MAKE_NO_WARN\n$(info )\n$(info ┌────────────────────────────────────────────────────────────────────┐)\n$(info │ NOTE: This Makefile is a compatibility shim. Prefer using mise: │)\n$(info │ │)\n$(info │ mise run build:cog mise run test:go mise run fmt:fix │)\n$(info │ │)\n$(info │ Run 'mise tasks' to see all available tasks. │)\n$(info │ Set MAKE_NO_WARN=1 to suppress this message. │)\n$(info └────────────────────────────────────────────────────────────────────┘)\n$(info )\nendif\n\nPREFIX ?= /usr/local\n\nGO ?= go\n\nCOG_BINARIES := cog base-image\n\ndefault: cog\n\n# =============================================================================\n# Build targets\n# =============================================================================\n\n.PHONY: cog base-image\n$(COG_BINARIES):\n\tmise run build:cog\n\n.PHONY: wheel\nwheel:\n\tmise run build:sdk\n\n.PHONY: install\ninstall: cog\n\tPREFIX=$(PREFIX) mise run install\n\n# =============================================================================\n# Test targets\n# =============================================================================\n\n.PHONY: test\ntest:\n\tmise run test:go\n\tmise run test:python\n\n.PHONY: test-go\ntest-go:\n\tmise run test:go\n\n.PHONY: test-python\ntest-python:\n\tmise run test:python\n\n.PHONY: test-integration\ntest-integration:\n\tmise run test:integration\n\n.PHONY: test-coglet\ntest-coglet: test-coglet-rust\n\n.PHONY: test-coglet-rust\ntest-coglet-rust:\n\tmise run test:rust\n\n.PHONY: test-coglet-python\ntest-coglet-python:\n\tmise run test:coglet:python\n\n# =============================================================================\n# Format and lint targets\n# =============================================================================\n\n.PHONY: fmt\nfmt:\n\tmise run fmt:fix\n\n.PHONY: check-fmt\ncheck-fmt:\n\tmise run fmt\n\n.PHONY: lint\nlint:\n\tmise run lint\n\n.PHONY: vet\nvet:\n\t$(GO) vet ./...\n\n# =============================================================================\n# Code generation\n# =============================================================================\n\n.PHONY: generate\ngenerate:\n\tmise run generate\n\n.PHONY: gen-mocks\ngen-mocks:\n\tmockery\n\n# =============================================================================\n# Coglet (Rust) targets\n# =============================================================================\n\n.PHONY: fmt-coglet\nfmt-coglet:\n\tmise run fmt:rust:fix\n\n.PHONY: check-fmt-coglet\ncheck-fmt-coglet:\n\tmise run fmt:rust\n\n.PHONY: lint-coglet\nlint-coglet:\n\tmise run lint:rust\n\n# =============================================================================\n# Documentation\n# =============================================================================\n\n.PHONY: run-docs-server\nrun-docs-server:\n\tmise run docs:serve\n\n# =============================================================================\n# Clean\n# =============================================================================\n\n.PHONY: clean\nclean:\n\tmise run clean\n\n.PHONY: clean-coglet\nclean-coglet:\n\tmise run clean:rust\n" }, { "path": "README.md", "content": "# Cog: Containers for machine learning\n\nCog is an open-source tool that lets you package machine learning models in a standard, production-ready container.\n\nYou can deploy your packaged model to your own infrastructure, or to [Replicate](https://replicate.com/).\n\n## Highlights\n\n- 📦 **Docker containers without the pain.** Writing your own `Dockerfile` can be a bewildering process. With Cog, you define your environment with a [simple configuration file](#how-it-works) and it generates a Docker image with all the best practices: Nvidia base images, efficient caching of dependencies, installing specific Python versions, sensible environment variable defaults, and so on.\n\n- 🤬️ **No more CUDA hell.** Cog knows which CUDA/cuDNN/PyTorch/Tensorflow/Python combos are compatible and will set it all up correctly for you.\n\n- ✅ **Define the inputs and outputs for your model with standard Python.** Then, Cog generates an OpenAPI schema and validates the inputs and outputs.\n\n- 🎁 **Automatic HTTP prediction server**: Your model's types are used to dynamically generate a RESTful HTTP API using a high-performance Rust/Axum server.\n\n- 🚀 **Ready for production.** Deploy your model anywhere that Docker images run. Your own infrastructure, or [Replicate](https://replicate.com).\n\n## How it works\n\nDefine the Docker environment your model runs in with `cog.yaml`:\n\n```yaml\nbuild:\n gpu: true\n system_packages:\n - \"libgl1-mesa-glx\"\n - \"libglib2.0-0\"\n python_version: \"3.13\"\n python_requirements: requirements.txt\npredict: \"predict.py:Predictor\"\n```\n\nDefine how predictions are run on your model with `predict.py`:\n\n```python\nfrom cog import BasePredictor, Input, Path\nimport torch\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n self.model = torch.load(\"./weights.pth\")\n\n # The arguments and types the model takes as input\n def predict(self,\n image: Path = Input(description=\"Grayscale input image\")\n ) -> Path:\n \"\"\"Run a single prediction on the model\"\"\"\n processed_image = preprocess(image)\n output = self.model(processed_image)\n return postprocess(output)\n```\n\nIn the above we accept a path to the image as an input, and return a path to our transformed image after running it through our model.\n\nNow, you can run predictions on this model:\n\n```console\n$ cog predict -i image=@input.jpg\n--> Building Docker image...\n--> Running Prediction...\n--> Output written to output.jpg\n```\n\nOr, build a Docker image for deployment:\n\n```console\n$ cog build -t my-classification-model\n--> Building Docker image...\n--> Built my-classification-model:latest\n\n$ docker run -d -p 5000:5000 --gpus all my-classification-model\n\n$ curl http://localhost:5000/predictions -X POST \\\n -H 'Content-Type: application/json' \\\n -d '{\"input\": {\"image\": \"https://.../input.jpg\"}}'\n```\n\nOr, combine build and run via the `serve` command:\n\n```console\n$ cog serve -p 8080\n\n$ curl http://localhost:8080/predictions -X POST \\\n -H 'Content-Type: application/json' \\\n -d '{\"input\": {\"image\": \"https://.../input.jpg\"}}'\n```\n\n\n\n## Why are we building this?\n\nIt's really hard for researchers to ship machine learning models to production.\n\nPart of the solution is Docker, but it is so complex to get it to work: Dockerfiles, pre-/post-processing, Flask servers, CUDA versions. More often than not the researcher has to sit down with an engineer to get the damn thing deployed.\n\n[Andreas](https://github.com/andreasjansson) and [Ben](https://github.com/bfirsh) created Cog. Andreas used to work at Spotify, where he built tools for building and deploying ML models with Docker. Ben worked at Docker, where he created [Docker Compose](https://github.com/docker/compose).\n\nWe realized that, in addition to Spotify, other companies were also using Docker to build and deploy machine learning models. [Uber](https://eng.uber.com/michelangelo-pyml/) and others have built similar systems. So, we're making an open source version so other people can do this too.\n\nHit us up if you're interested in using it or want to collaborate with us. [We're on Discord](https://discord.gg/replicate) or email us at [team@replicate.com](mailto:team@replicate.com).\n\n## Prerequisites\n\n- **macOS, Linux or Windows 11**. Cog works on macOS, Linux and Windows 11 with [WSL 2](docs/wsl2/wsl2.md)\n- **Docker**. Cog uses Docker to create a container for your model. You'll need to [install Docker](https://docs.docker.com/get-docker/) before you can run Cog. If you install Docker Engine instead of Docker Desktop, you will need to [install Buildx](https://docs.docker.com/build/architecture/#buildx) as well.\n\n## Install\n\nIf you're using macOS, you can install Cog using Homebrew:\n\n```console\nbrew install replicate/tap/cog\n```\n\nYou can also download and install the latest release using our\n[install script](https://cog.run/install):\n\n```sh\n# bash, zsh, and other shells\nsh <(curl -fsSL https://cog.run/install.sh)\n\n# fish shell\nsh (curl -fsSL https://cog.run/install.sh | psub)\n\n# download with wget and run in a separate command\nwget -qO- https://cog.run/install.sh\nsh ./install.sh\n```\n\nYou can manually install the latest release of Cog directly from GitHub\nby running the following commands in a terminal:\n\n```console\nsudo curl -o /usr/local/bin/cog -L \"https://github.com/replicate/cog/releases/latest/download/cog_$(uname -s)_$(uname -m)\"\nsudo chmod +x /usr/local/bin/cog\n```\n\nOr if you are on docker:\n\n```\nRUN sh -c \"INSTALL_DIR=\\\"/usr/local/bin\\\" SUDO=\\\"\\\" $(curl -fsSL https://cog.run/install.sh)\"\n```\n\n## Upgrade\n\nIf you're using macOS and you previously installed Cog with Homebrew, run the following:\n\n```console\nbrew upgrade replicate/tap/cog\n```\n\nOtherwise, you can upgrade to the latest version by running the same commands you used to install it.\n\n## Development\n\nSee [CONTRIBUTING.md](CONTRIBUTING.md) for how to set up a development environment and build from source.\n\n## Next steps\n\n- [Get started with an example model](docs/getting-started.md)\n- [Get started with your own model](docs/getting-started-own-model.md)\n- [Using Cog with notebooks](docs/notebooks.md)\n- [Using Cog with Windows 11](docs/wsl2/wsl2.md)\n- [Take a look at some examples of using Cog](https://github.com/replicate/cog-examples)\n- [Deploy models with Cog](docs/deploy.md)\n- [`cog.yaml` reference](docs/yaml.md) to learn how to define your model's environment\n- [Prediction interface reference](docs/python.md) to learn how the `Predictor` interface works\n- [Training interface reference](docs/training.md) to learn how to add a fine-tuning API to your model\n- [HTTP API reference](docs/http.md) to learn how to use the HTTP API that models serve\n\n## Need help?\n\n[Join us in #cog on Discord.](https://discord.gg/replicate)\n\n[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/replicate/cog)\n\n## Contributors ✨\n\nThanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):\n\n\n\n\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
\"Ben
Ben Firshman
💻 📖		\"Andreas
Andreas Jansson
💻 📖 🚧		\"Zeke
Zeke Sikelianos
💻 📖 🔧		\"Rory
Rory Byrne
💻 📖 ⚠️		\"Michael
Michael Floering
💻 📖 🤔		\"Ben
Ben Evans
📖		\"shashank
shashank agarwal
💻 📖
\"VictorXLR\"/
VictorXLR
💻 📖 ⚠️		\"hung
hung anna
🐛		\"Brian
Brian Whitman
🐛		\"JimothyJohn\"/
JimothyJohn
🐛		\"ericguizzo\"/
ericguizzo
🐛		\"Dominic
Dominic Baggott
💻 ⚠️		\"Dashiell
Dashiell Stander
🐛 💻 ⚠️
\"Shuwei
Shuwei Liang
🐛 💬		\"Eric
Eric Allam
🤔		\"Iván
Iván Perdomo
🐛		\"Charles
Charles Frye
📖		\"Luan
Luan Pham
🐛 📖		\"TommyDew\"/
TommyDew
💻		\"Jesse
Jesse Andrews
💻 📖 ⚠️
\"Nick
Nick Stenning
💻 📖 🎨 🚇 ⚠️		\"Justin
Justin Merrell
📖		\"Rurik
Rurik Ylä-Onnenvuori
🐛		\"Youka\"/
Youka
🐛		\"Clay
Clay Mullis
📖		\"Mattt\"/
Mattt
💻 📖 🚇		\"Eng
Eng Zer Jun
⚠️
\"BB\"/
BB
💻		\"williamluer\"/
williamluer
📖		\"Simon
Simon Eskildsen
💻		\"F\"/
F
🐛 💻		\"Philip
Philip Potter
🐛 💻		\"Joanne
Joanne Chen
📖		\"technillogue\"/
technillogue
💻
\"Aron
Aron Carroll
📖 💻 🤔		\"Bohdan
Bohdan Mykhailenko
📖 🐛		\"Daniel
Daniel Radu
📖 🐛		\"Itay
Itay Etelis
💻		\"Gennaro
Gennaro Schiano
📖		\"André
André Knörig
📖		\"Dan
Dan Fairs
💻
\n\n\n\n\n\n\nThis project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!\n" }, { "path": "architecture/00-overview.md", "content": "# Cog Architecture Overview\n\nCog packages machine learning models into production-ready OCI images.\n\n## The Big Picture\n\n```mermaid\nflowchart LR\n subgraph input[\"What you write\"]\n model[\"Model Code
+ cog.yaml\"]\n end\n \n subgraph cog[\"Cog\"]\n cli[\"CLI\"]\n sdk[\"Python SDK\"]\n end\n \n subgraph output[\"What you get\"]\n image[\"Container Image\"]\n api[\"HTTP API\"]\n end\n \n model --> cli\n cli -->|\"builds\"| image\n image -->|\"runs\"| sdk\n sdk -->|\"serves\"| api\n```\n\n## Components\n\n### Model Source\n\nWhat the model author provides: `cog.yaml` for environment config, a Predictor class with `setup()` and `predict()` methods, and optionally model weights.\n\n**Deep dive**: [Model Source](./01-model-source.md)\n\n---\n\n### Schema\n\nAn OpenAPI specification generated from the predictor's type hints. Describes what inputs the model accepts and what outputs it produces.\n\n**Deep dive**: [Schema](./02-schema.md)\n\n---\n\n### Prediction API\n\nThe HTTP interface for running predictions. A fixed envelope format (`PredictionRequest`/`PredictionResponse`) wraps model-specific inputs and outputs.\n\n**Deep dives**:\n- [Legacy API](./legacy/03-prediction-api.md) - FastAPI implementation details\n- [FFI API](./ffi/03-prediction-api.md) - Rust/Axum implementation details\n\n---\n\n### Container Runtime\n\nThe runtime that runs inside the container: an HTTP server, worker process isolation, and prediction execution. Cog has two runtime implementations:\n\n- **Legacy (Python/FastAPI)**: Current default implementation - [Documentation](./legacy/)\n- **FFI (Rust/PyO3)**: Next-generation experimental implementation - [Documentation](./ffi/)\n\n**Deep dives**:\n- [Legacy Runtime](./legacy/04-container-runtime.md) - FastAPI/Uvicorn two-process architecture\n- [FFI Runtime](./ffi/04-container-runtime.md) - Rust/Axum with PyO3 FFI bridge\n\n---\n\n### Build System\n\nTransforms `cog.yaml` and user code into a Docker image with the right Python version, CUDA libraries, and dependencies.\n\n**Deep dive**: [Build System](./05-build-system.md)\n\n---\n\n### CLI\n\nThe command-line tool for building, testing, and deploying models.\n\n**Deep dive**: [CLI](./06-cli.md)\n\n---\n\n## How It Fits Together\n\n```mermaid\nflowchart TB\n subgraph source[\"Model Source\"]\n yaml[\"cog.yaml\"]\n code[\"predict.py\"]\n weights[\"weights\"]\n end\n \n subgraph build[\"Build Time\"]\n config[\"Config Parser\"]\n generator[\"Dockerfile Generator\"]\n schema_gen[\"Schema Generator\"]\n end\n \n subgraph image[\"Container Image\"]\n layers[\"Base + Deps + Code\"]\n schema[\"OpenAPI Schema
(label)\"]\n end\n \n subgraph runtime[\"Runtime\"]\n server[\"HTTP Server\"]\n worker[\"Worker Process\"]\n predictor[\"Predictor\"]\n end\n \n yaml --> config\n config --> generator\n generator --> layers\n code --> layers\n weights --> layers\n \n layers --> schema_gen\n schema_gen --> schema\n \n image --> server\n server --> worker\n worker --> predictor\n```\n\n## Terminology\n\n| Term | Meaning |\n|------|---------|\n| **Predictor** | User's model class with `setup()` and `predict()` methods |\n| **Schema** | OpenAPI spec describing the model's input/output interface |\n| **Envelope** | Fixed request/response structure wrapping model-specific data |\n| **Worker** | Isolated subprocess running user code |\n| **Setup** | One-time model initialization at container start |\n\n## Runtime Implementations\n\nCog supports two runtime implementations:\n\n### Legacy Runtime (Python/FastAPI)\n- **Status**: Current default\n- **Use when**: Running standard Cog containers\n- **Implementation**: `python/cog/server/`\n- **Documentation**: [legacy/](./legacy/)\n\n### FFI Runtime (Rust/PyO3)\n- **Status**: Experimental (in development)\n- **Use when**: Set `USE_COGLET` environment variable\n- **Implementation**: `crates/coglet/`\n- **Documentation**: [ffi/](./ffi/)\n- **Benefits**: Better performance, stability, and resource management\n\nBoth runtimes expose the same HTTP API and support the same model code. The FFI runtime is a drop-in replacement with improved internals.\n\n## Reading Order\n\nFor understanding Cog's architecture, we recommend reading in this order:\n\n1. [Model Source](./01-model-source.md) - What users write\n2. [Schema](./02-schema.md) - How the interface is described\n3. **Choose a runtime path**:\n - **Legacy**: [Prediction API](./legacy/03-prediction-api.md) → [Container Runtime](./legacy/04-container-runtime.md)\n - **FFI**: [Prediction API](./ffi/03-prediction-api.md) → [Container Runtime](./ffi/04-container-runtime.md)\n4. [Build System](./05-build-system.md) - How images are built\n5. [CLI](./06-cli.md) - How users interact with it all\n" }, { "path": "architecture/01-model-source.md", "content": "# Model Source\n\nThis document covers what a model author provides to Cog and the primitives they work with.\n\n## What Users Write\n\nA Cog model consists of:\n\n```\nmy-model/\n├── cog.yaml # Environment configuration\n├── predict.py # Predictor class\n└── weights/ # Model weights (optional, can be downloaded)\n```\n\n## cog.yaml\n\nDeclares the runtime environment:\n\n```yaml\nbuild:\n python_version: \"3.11\"\n gpu: true\n python_packages:\n - torch==2.1.0\n - transformers==4.35.0\n system_packages:\n - ffmpeg\n run:\n - curl -o /src/model.bin https://example.com/model.bin\n\npredict: \"predict.py:Predictor\"\n\nconcurrency:\n max: 1\n```\n\n| Field | Purpose |\n|-------|---------|\n| `build.python_version` | Python interpreter version (3.10-3.13) |\n| `build.gpu` | Enable CUDA support |\n| `build.python_packages` | pip packages to install |\n| `build.system_packages` | apt packages to install |\n| `build.run` | Arbitrary shell commands during build |\n| `predict` | Path to predictor class (`module:ClassName`) |\n| `train` | Path to training class (optional) |\n| `concurrency.max` | Max concurrent predictions (requires async) |\n\nThe [Build System](./05-build-system.md) uses this configuration to produce an image containing all necessary dependencies, libraries, and the correct Python/CUDA versions.\n\n## The Predictor Class\n\nA predictor is a Python class with two methods:\n\n```python\nfrom cog import BasePredictor, Input, Path\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load model into memory. Called once at container start.\"\"\"\n self.model = load_model(\"./weights\")\n \n def predict(self, prompt: str, steps: int = 50) -> Path:\n \"\"\"Run inference. Called for each prediction request.\"\"\"\n output = self.model.generate(prompt, steps=steps)\n output.save(\"/tmp/output.png\")\n return Path(\"/tmp/output.png\")\n```\n\n### setup()\n\n- Called **once** when the container starts\n- Used to load model weights, initialize GPU contexts, warm up caches\n- Runs before the HTTP server accepts requests\n- Optional: if omitted, Cog proceeds directly to serving\n\n### predict()\n\n- Called **for each prediction request**\n- Signature defines the model's input schema (via type hints)\n- Return type defines the output schema\n- Can be sync (`def`) or async (`async def`)\n\n### train() (optional)\n\n- Same contract as `predict()` but for fine-tuning workflows\n- Configured separately in `cog.yaml` with `train:` key\n\n## Input Types\n\nThe types used in `predict()` parameters become the model's input schema.\n\n### Basic Types\n\n```python\ndef predict(\n self,\n text: str, # String input\n count: int, # Integer\n temperature: float, # Float\n verbose: bool, # Boolean\n) -> str:\n```\n\n### File Inputs (cog.Path)\n\nURLs are automatically downloaded to local files:\n\n```python\nfrom cog import Path\n\ndef predict(self, image: Path) -> Path:\n # Client sends: {\"input\": {\"image\": \"https://example.com/photo.jpg\"}}\n # Cog downloads the URL, `image` is a local path like /tmp/inputabc123.jpg\n img = PIL.Image.open(image)\n ...\n```\n\n`cog.Path` extends `pathlib.Path`. At runtime:\n- HTTP/HTTPS URLs are downloaded to temp files\n- Data URLs are decoded\n- The predictor receives a local filesystem path\n\n### Secrets (cog.Secret)\n\nFor sensitive values that shouldn't appear in logs:\n\n```python\nfrom cog import Secret\n\ndef predict(self, api_key: Secret) -> str:\n # Value is masked in logs and webhooks\n client = SomeAPI(api_key.get_secret_value())\n ...\n```\n\n### Input Constraints\n\nUse `Input()` to add metadata and validation:\n\n```python\nfrom cog import Input\n\ndef predict(\n self,\n prompt: str = Input(description=\"The text prompt\"),\n steps: int = Input(default=50, ge=1, le=100, description=\"Inference steps\"),\n style: str = Input(choices=[\"photo\", \"art\", \"sketch\"]),\n) -> str:\n```\n\n| Parameter | Effect |\n|-----------|--------|\n| `description` | Shown in UI and schema |\n| `default` | Default value if not provided |\n| `ge`, `le` | Numeric bounds (greater/less than or equal) |\n| `min_length`, `max_length` | String length bounds |\n| `choices` | Enum values (deprecated: prefer `Literal`) |\n\n### Enums with Literal\n\n```python\nfrom typing import Literal\n\ndef predict(\n self,\n size: Literal[\"small\", \"medium\", \"large\"] = \"medium\",\n) -> str:\n```\n\n### Lists\n\n```python\nfrom typing import List\nfrom cog import Path\n\ndef predict(\n self,\n images: List[Path], # Multiple file inputs\n tags: List[str], # Multiple strings\n) -> str:\n```\n\n### Optional Inputs\n\n```python\nfrom typing import Optional\n\ndef predict(\n self,\n seed: Optional[int] = None, # Can be omitted or null\n) -> str:\n```\n\n## Output Types\n\nThe return type annotation defines what the model produces.\n\n### Basic Types\n\n```python\ndef predict(self, prompt: str) -> str:\n return \"Generated text...\"\n```\n\n### File Outputs\n\nReturn `cog.Path` pointing to a generated file:\n\n```python\nfrom cog import Path\n\ndef predict(self, prompt: str) -> Path:\n # Generate file\n output_path = \"/tmp/output.png\"\n self.model.generate(prompt).save(output_path)\n return Path(output_path)\n```\n\nAt runtime, Cog uploads the file and returns a URL to the client.\n\n### Multiple Outputs\n\nReturn a list:\n\n```python\nfrom typing import List\nfrom cog import Path\n\ndef predict(self, prompt: str) -> List[Path]:\n paths = []\n for i in range(4):\n path = f\"/tmp/output_{i}.png\"\n self.model.generate(prompt, seed=i).save(path)\n paths.append(Path(path))\n return paths\n```\n\n### Streaming with Iterator\n\nYield values progressively:\n\n```python\nfrom typing import Iterator\n\ndef predict(self, prompt: str) -> Iterator[str]:\n for token in self.model.generate_stream(prompt):\n yield token\n```\n\nThe schema marks this as `x-cog-array-type: iterator`. Clients receive outputs as they're produced via webhooks or streaming responses.\n\n### Streaming Text with ConcatenateIterator\n\nFor LLM-style token streaming where outputs should be concatenated:\n\n```python\nfrom cog import ConcatenateIterator\n\ndef predict(self, prompt: str) -> ConcatenateIterator[str]:\n for token in self.model.generate(prompt):\n yield token # \"Hello\", \" \", \"world\", \"!\"\n # Client sees progressive: \"Hello\" -> \"Hello \" -> \"Hello world\" -> \"Hello world!\"\n```\n\nThe schema includes `x-cog-array-display: concatenate` to signal that outputs should be joined rather than displayed as a list.\n\n## Weights\n\nModel weights can be loaded in several ways:\n\n### Bundled in the Image\n\nInclude weights in your source directory - they're copied into the image during build:\n\n```\nmy-model/\n├── cog.yaml\n├── predict.py\n└── weights/\n └── model.safetensors\n```\n\n```python\ndef setup(self):\n self.model = load(\"./weights/model.safetensors\")\n```\n\n### Downloaded at Runtime\n\nWeights can be fetched during `setup()` rather than bundled. Common approaches:\n\n**Using the `weights` parameter** (Cog's built-in mechanism):\n\n```python\nclass Predictor(BasePredictor):\n def setup(self, weights: Path):\n self.model = load(weights)\n```\n\nThe `weights` value comes from `COG_WEIGHTS` env var or falls back to `./weights`:\n\n```bash\nCOG_WEIGHTS=https://example.com/model.tar cog predict ...\n```\n\n**Using pget** (parallel download tool, included in Cog images):\n\n```python\nimport subprocess\n\ndef setup(self):\n subprocess.run([\"pget\", \"https://example.com/model.tar\", \"./weights\"])\n self.model = load(\"./weights/model.safetensors\")\n```\n\n**Direct download in setup**:\n\n```python\ndef setup(self):\n # Using requests, huggingface_hub, or any other method\n snapshot_download(repo_id=\"meta-llama/Llama-2-7b\", local_dir=\"./weights\")\n self.model = load(\"./weights\")\n```\n\nThe choice depends on your deployment needs - bundled weights make images larger but start faster; downloaded weights keep images small but require network access at startup.\n\n## Async Predictors\n\nFor concurrent predictions, use async:\n\n```python\nclass Predictor(BasePredictor):\n async def setup(self):\n self.model = await load_model_async()\n \n async def predict(self, prompt: str) -> str:\n return await self.model.generate(prompt)\n```\n\nRequires:\n- Python 3.11+\n- `concurrency.max > 1` in cog.yaml\n\nSee [Container Runtime](./04-container-runtime.md) for concurrency details.\n\n## Code References\n\n| File | Purpose |\n|------|---------|\n| `python/cog/__init__.py` | Public API exports |\n| `python/cog/base_predictor.py` | BasePredictor class |\n| `python/cog/types.py` | Input, Path, Secret, ConcatenateIterator |\n| `python/cog/predictor.py` | Type introspection, weights handling |\n| `pkg/config/config.go` | cog.yaml parsing |\n" }, { "path": "architecture/02-schema.md", "content": "# Schema\n\nThe schema is an **OpenAPI 3.0.2 specification** that describes a model's interface. It's the contract between the model and everything that interacts with it.\n\n## Why the Schema Exists\n\nEvery Cog model uses the same [Prediction API](./ffi/03-prediction-api.md) envelope format, but the `input` and `output` fields are model-specific. The schema captures what each model expects and produces.\n\n```\n┌─────────────────────────────────────────────────┐\n│ PredictionRequest (fixed envelope) │\n│ ┌─────────────────────────────────────────┐ │\n│ │ \"input\": { ... } <- model-specific │ │\n│ └─────────────────────────────────────────┘ │\n└─────────────────────────────────────────────────┘\n ↑\n Schema defines this part\n```\n\nWithout the schema, consumers would have no way to know:\n- What inputs the model accepts\n- What types those inputs should be\n- What constraints apply (required fields, min/max values, allowed choices)\n- What the output looks like\n\n### How It's Used Today\n\n| Consumer | What They Use the Schema For |\n|----------|------------------------------|\n| **Replicate platform** | Generate input forms in the web UI, validate requests before routing to models |\n| **HTTP server (coglet)** | Validate incoming JSON, reject malformed requests before they reach user code |\n| **CLI (`cog predict`)** | Parse `-i key=value` flags into correctly-typed Python objects |\n| **Docker label** | Extract model interface without running the container |\n| **API clients** | Know what to send and what to expect back without reading source code |\n\n## How It's Generated\n\nCog supports two schema generation paths:\n\n### Legacy Runtime Path (default)\n\nThe **legacy path** boots the built Docker container and runs `python -m cog.command.openapi_schema` to introspect the model at runtime using pydantic. This is the default for all builds. It works with any Python type that pydantic can serialize, including third-party types, complex inheritance, and dynamically constructed classes.\n\n### Static Path (opt-in)\n\nThe **static path** parses Python source code at `cog build` time using [tree-sitter](https://tree-sitter.github.io/tree-sitter/) in Go. No Python runtime is invoked. This makes schema generation deterministic, fast, and independent of the model's dependencies.\n\nEnable it by setting the `COG_STATIC_SCHEMA` environment variable:\n\n```bash\nCOG_STATIC_SCHEMA=1 cog build -t my-model\n```\n\nThe static path requires SDK >= 0.17.0. When opted in, if the static parser encounters a type it cannot resolve, it **falls back to the legacy runtime path** automatically with a warning — so builds never fail due to static parser limitations.\n\nFor local commands (`cog train`, `cog serve`, `cog predict`), the static path is always used regardless of the `COG_STATIC_SCHEMA` flag, because these commands return before the post-build legacy generation step — the CLI needs the schema to parse `-i` input flags.\n\n```mermaid\nflowchart LR\n subgraph source[\"Model Source\"]\n predict[\"predict.py\"]\n types[\"output_types.py\"]\n end\n\n subgraph parser[\"Go Static Parser\"]\n ts[\"tree-sitter Python\"]\n resolve[\"Type Resolver\"]\n cross[\"Cross-File Resolver\"]\n end\n\n subgraph output[\"Schema\"]\n spec[\"OpenAPI 3.0.2 JSON\"]\n end\n\n predict --> ts\n types --> cross\n ts --> resolve\n cross --> resolve\n resolve --> spec\n```\n\n### Static Path Pipeline Steps\n\n1. **Parse** the predictor file with tree-sitter (concrete syntax tree, not AST)\n2. **Collect imports** — track where each name came from (`from cog import Path`, `from pydantic import BaseModel`)\n3. **Collect module scope** — resolve module-level variable assignments (for default values, choices lists)\n4. **Collect BaseModel subclasses** — find all classes that inherit from `BaseModel` (cog or pydantic) in the current file\n5. **Resolve cross-file models** — for imported names not found locally, find the `.py` file on disk, parse it, and extract its BaseModel definitions\n6. **Extract inputs** — walk the `predict()` / `train()` method parameters, resolve types, defaults, and `Input()` metadata\n7. **Resolve output type** — recursively resolve the return type annotation into a `SchemaType`\n8. **Generate OpenAPI** — convert the extracted `PredictorInfo` into a full OpenAPI 3.0.2 JSON document\n\nIf any step fails with an unresolvable type, the build falls back to the legacy runtime path.\n\n### Cross-File Resolution\n\nWhen a predictor imports types from other project files, the schema generator resolves them automatically:\n\n```python\n# output_types.py\nfrom pydantic import BaseModel\n\nclass Prediction(BaseModel):\n text: str\n score: float\n tags: list[str]\n```\n\n```python\n# predict.py\nfrom cog import BasePredictor\nfrom output_types import Prediction\n\nclass Predictor(BasePredictor):\n def predict(self, prompt: str) -> Prediction:\n ...\n```\n\nThe resolver handles every permutation of local imports:\n\n| Import Style | File Resolved |\n|-------------|---------------|\n| `from output_types import X` | `/output_types.py` |\n| `from .output_types import X` | `/output_types.py` |\n| `from models.output import X` | `/models/output.py` |\n| `from .models.output import X` | `/models/output.py` |\n| `from output_types import X as Y` | `/output_types.py` (alias tracked) |\n\n**How it distinguishes local from external**: the resolver converts the module path to a filesystem path and checks if the file exists. If `output_types.py` exists in the project directory, it's local. If not (e.g., `from transformers import ...`), it's external. Known external packages (stdlib, torch, numpy, etc.) are skipped without a filesystem check.\n\n**Error messages**: when a type can't be resolved, the error includes the import source:\n```\ncannot resolve output type 'WeirdType' (imported from 'some_package') —\nexternal types cannot be statically analyzed. Define it as a BaseModel\nsubclass in your predict file, or provide a .pyi stub\n```\n\n## SchemaType: The Type System\n\nOutput types are represented as a recursive algebraic data type (`SchemaType`) that composes arbitrarily:\n\n```\nSchemaType\n├── SchemaPrimitive — str, int, float, bool, Path\n├── SchemaAny — untyped (bare dict, Any)\n├── SchemaArray — list[T], with Items → SchemaType\n├── SchemaDict -- dict[str, V], with ValueType -> SchemaType\n├── SchemaObject — BaseModel subclass, with Fields → OrderedMap[name, SchemaField]\n├── SchemaIterator — Iterator[T], with Elem → SchemaType\n└── SchemaConcatIterator — ConcatenateIterator[str]\n```\n\nThis recursive structure means nested types like `dict[str, list[dict[str, int]]]` are fully representable and produce correct JSON Schema:\n\n```json\n{\n \"type\": \"object\",\n \"additionalProperties\": {\n \"type\": \"array\",\n \"items\": {\n \"type\": \"object\",\n \"additionalProperties\": {\n \"type\": \"integer\"\n }\n }\n }\n}\n```\n\n### JSON Schema Generation\n\nEach `SchemaType` produces its JSON Schema fragment via `JSONSchema()`:\n\n| SchemaType Kind | JSON Schema |\n|-----------------|-------------|\n| `SchemaPrimitive(str)` | `{\"type\": \"string\"}` |\n| `SchemaPrimitive(Path)` | `{\"type\": \"string\", \"format\": \"uri\"}` |\n| `SchemaAny` | `{\"type\": \"object\"}` |\n| `SchemaArray(items)` | `{\"type\": \"array\", \"items\": items.JSONSchema()}` |\n| `SchemaDict(valueType)` | `{\"type\": \"object\", \"additionalProperties\": valueType.JSONSchema()}` |\n| `SchemaObject(fields)` | `{\"type\": \"object\", \"properties\": {...}, \"required\": [...]}` |\n| `SchemaIterator(elem)` | `{\"type\": \"array\", \"items\": elem.JSONSchema(), \"x-cog-array-type\": \"iterator\"}` |\n| `SchemaConcatIterator` | `{\"type\": \"array\", \"items\": {\"type\": \"string\"}, \"x-cog-array-type\": \"iterator\", \"x-cog-array-display\": \"concatenate\"}` |\n\n## Type Mappings\n\n### Input Types\n\n| Python | JSON Schema | Notes |\n|--------|-------------|-------|\n| `str` | `{\"type\": \"string\"}` | |\n| `int` | `{\"type\": \"integer\"}` | |\n| `float` | `{\"type\": \"number\"}` | |\n| `bool` | `{\"type\": \"boolean\"}` | |\n| `cog.Path` | `{\"type\": \"string\", \"format\": \"uri\"}` | URLs downloaded at runtime |\n| `cog.File` | `{\"type\": \"string\", \"format\": \"uri\"}` | File uploads |\n| `cog.Secret` | `{\"type\": \"string\", \"format\": \"password\", \"x-cog-secret\": true}` | Masked in logs |\n| `list[T]` | `{\"type\": \"array\", \"items\": {...}}` | |\n| `Optional[T]` | Type T + not in `required` | Input fields only |\n| `Literal[\"a\", \"b\"]` / `choices=[...]` | `{\"enum\": [\"a\", \"b\"]}` | |\n\n### Output Types\n\n| Python | SchemaType | JSON Schema |\n|--------|------------|-------------|\n| `str` | `SchemaPrimitive` | `{\"type\": \"string\"}` |\n| `int` | `SchemaPrimitive` | `{\"type\": \"integer\"}` |\n| `float` | `SchemaPrimitive` | `{\"type\": \"number\"}` |\n| `bool` | `SchemaPrimitive` | `{\"type\": \"boolean\"}` |\n| `Path` | `SchemaPrimitive` | `{\"type\": \"string\", \"format\": \"uri\"}` |\n| `dict` (bare) | `SchemaAny` | `{\"type\": \"object\"}` |\n| `dict[str, V]` | `SchemaDict` | `{\"type\": \"object\", \"additionalProperties\": V}` |\n| `list` (bare) | `SchemaArray(SchemaAny)` | `{\"type\": \"array\", \"items\": {\"type\": \"object\"}}` |\n| `list[T]` | `SchemaArray` | `{\"type\": \"array\", \"items\": T}` |\n| `BaseModel` subclass | `SchemaObject` | `{\"type\": \"object\", \"properties\": {...}}` |\n| `Iterator[T]` | `SchemaIterator` | `{\"type\": \"array\", \"items\": T, \"x-cog-array-type\": \"iterator\"}` |\n| `ConcatenateIterator[str]` | `SchemaConcatIterator` | Streaming token output |\n| Nested types | Recursive | `dict[str, list[dict[str, int]]]` fully supported |\n\n### Unsupported Output Types\n\n| Python | Error |\n|--------|-------|\n| `Optional[T]` / `T \\| None` | Predictions must succeed with a value or fail with an error |\n| `Union[A, B]` | Ambiguous for downstream consumers |\n| External package types | Cannot be statically analyzed — define as BaseModel or use .pyi stub |\n\n## Cog-Specific Extensions\n\n| Extension | Purpose |\n|-----------|---------|\n| `x-order` | Preserves parameter order from function signature |\n| `x-cog-array-type` | Marks iterators vs regular arrays |\n| `x-cog-array-display` | Hints for how to display streaming output |\n| `x-cog-secret` | Marks sensitive inputs |\n\n## Where the Schema Lives\n\n### In the Image\n\nEmbedded as a Docker label during build:\n\n```bash\ndocker inspect my-model | jq -r '.[0].Config.Labels[\"run.cog.openapi_schema\"]'\n```\n\nAlso written to `.cog/openapi_schema.json` inside the image for the runtime to serve.\n\n### At Runtime\n\n| Endpoint | Format |\n|----------|--------|\n| `GET /openapi.json` | Raw OpenAPI spec |\n\n### Override and Configuration\n\n| Environment Variable | Purpose |\n|---------------------|---------|\n| `COG_STATIC_SCHEMA=1` | Opt in to the static Go tree-sitter schema generator (falls back to legacy on failure) |\n| `COG_OPENAPI_SCHEMA=path` | Skip generation entirely and use a pre-built schema file |\n\n```bash\n# Use static schema generation\nCOG_STATIC_SCHEMA=1 cog build -t my-model\n\n# Use a pre-built schema file\nCOG_OPENAPI_SCHEMA=my_schema.json cog build\n```\n\n## Schema Structure\n\nA simplified example showing a multi-file predictor with structured output:\n\n```json\n{\n \"openapi\": \"3.0.2\",\n \"info\": { \"title\": \"Cog\", \"version\": \"0.1.0\" },\n \"paths\": {\n \"/predictions\": {\n \"post\": {\n \"requestBody\": {\n \"content\": {\n \"application/json\": {\n \"schema\": { \"$ref\": \"#/components/schemas/PredictionRequest\" }\n }\n }\n }\n }\n }\n },\n \"components\": {\n \"schemas\": {\n \"Input\": {\n \"type\": \"object\",\n \"properties\": {\n \"prompt\": {\n \"type\": \"string\",\n \"description\": \"Text prompt\",\n \"x-order\": 0\n },\n \"steps\": {\n \"type\": \"integer\",\n \"default\": 50,\n \"minimum\": 1,\n \"maximum\": 100,\n \"x-order\": 1\n }\n },\n \"required\": [\"prompt\"]\n },\n \"Output\": {\n \"type\": \"object\",\n \"properties\": {\n \"text\": { \"type\": \"string\", \"title\": \"Text\" },\n \"score\": { \"type\": \"number\", \"title\": \"Score\" }\n },\n \"required\": [\"text\", \"score\"]\n },\n \"PredictionRequest\": { \"...\" : \"...\" },\n \"PredictionResponse\": { \"...\" : \"...\" }\n }\n }\n}\n```\n\n## Code References\n\n| File | Purpose |\n|------|---------|\n| `pkg/schema/schema_type.go` | `SchemaType` ADT, `ResolveSchemaType()`, `JSONSchema()` generation |\n| `pkg/schema/types.go` | `PredictorInfo`, `PrimitiveType`, `FieldType`, `InputField`, `ImportContext` |\n| `pkg/schema/python/parser.go` | Tree-sitter Python parser, `ParsePredictor()`, cross-file resolution |\n| `pkg/schema/openapi.go` | OpenAPI document assembly from `PredictorInfo` |\n| `pkg/schema/generator.go` | Top-level `Generate()`, `GenerateCombined()`, `Parser` type |\n| `pkg/schema/errors.go` | Typed error kinds (`ErrUnresolvableType`, `ErrOptionalOutput`, etc.) |\n| `pkg/image/build.go` | `canUseStaticSchemaGen()` — opt-in gate, `generateStaticSchema()` — entry point, fallback to legacy on `ErrUnresolvableType` |\n| `pkg/image/openapi_schema.go` | `GenerateOpenAPISchema()` — legacy runtime path (boots container, runs `python -m cog.command.openapi_schema`) |\n| `python/cog/_adt.py` | Internal ADT types for Python-side predictor introspection |\n| `python/cog/_inspector.py` | Python-side predictor inspector (runtime introspection for legacy path) |\n| `python/cog/_schemas.py` | Python-side OpenAPI schema generation from inspected predictor info |\n| `python/cog/command/openapi_schema.py` | CLI entry point for `python -m cog.command.openapi_schema` (invoked by legacy runtime path) |\n" }, { "path": "architecture/05-build-system.md", "content": "# Build System\n\nThe build system transforms [Model Source](./01-model-source.md) (cog.yaml + predict.py + weights) into a production-ready OCI image containing the [Container Runtime](./04-container-runtime.md).\n\n## Build Flow\n\n```mermaid\nflowchart TB\n subgraph input[\"Inputs\"]\n yaml[\"cog.yaml\"]\n code[\"predict.py\"]\n weights[\"weights\"]\n end\n \n subgraph cli[\"CLI (pkg/cli/build.go)\"]\n parse[\"Parse Config\"]\n validate[\"Validate\"]\n end\n \n subgraph generate[\"Dockerfile Generation (pkg/dockerfile/)\"]\n generator[\"Generator\"]\n baseimage[\"Base Image Selection\"]\n compat[\"Compatibility Matrix\"]\n wheel[\"Embedded Python Wheel\"]\n end\n \n subgraph docker[\"Docker Build\"]\n buildkit[\"Buildkit\"]\n image[\"Container Image\"]\n end\n \n subgraph post[\"Post-Build\"]\n schema[\"Generate OpenAPI Schema\"]\n freeze[\"pip freeze\"]\n labels[\"Apply Labels\"]\n end\n \n yaml --> parse --> validate\n validate --> generator\n compat --> generator\n baseimage --> generator\n wheel --> generator\n generator -->|\"Dockerfile\"| buildkit\n code --> buildkit\n weights --> buildkit\n buildkit --> image\n image --> schema\n image --> freeze\n schema --> labels\n freeze --> labels\n labels -->|\"Final Image\"| output[\"Tagged Image\"]\n```\n\n## Key Components\n\n### 1. Config Parsing & Validation\n\nReads `cog.yaml` and validates/completes the configuration:\n- Validates Python version (3.10-3.13)\n- Auto-detects CUDA version from PyTorch/TensorFlow requirements\n- Resolves package versions against compatibility matrix\n\n```\ncog.yaml (user provides) → Config (completed)\n───────────────────────── ─────────────────\ngpu: true gpu: true\npython_packages: cuda: \"12.1\" ← auto-detected\n - torch==2.1.0 cudnn: \"8\" ← auto-detected\n```\n\n---\n\n### 2. Dockerfile Generator\n\nThe generator produces a Dockerfile from the validated config.\n\n#### Generated Dockerfile Sections\n\n```dockerfile\n# 1. Base image (cog-base, CUDA, or python-slim)\nFROM r8.im/cog-base:cuda12.1-python3.11-torch2.1.0\n\n# 2. System packages\nRUN apt-get update && apt-get install -y ffmpeg\n\n# 3. Python packages\nRUN pip install -r requirements.txt\n\n# 4. Cog wheel (embedded in CLI binary)\nCOPY cog-0.12.0-py3-none-any.whl /tmp/\nRUN pip install /tmp/cog-0.12.0-py3-none-any.whl\n\n# 5. User run commands\nRUN echo \"custom setup\"\n\n# 6. Copy source\nWORKDIR /src\nCOPY . /src\n\n# 7. Entrypoint\nENTRYPOINT [\"/sbin/tini\", \"--\"]\nCMD [\"python\", \"-m\", \"cog.server.http\"]\n```\n\n---\n\n### 3. Compatibility Matrix\n\nPyTorch and TensorFlow releases are built against specific CUDA/cuDNN versions. The compatibility matrix captures these relationships from upstream release notes.\n\n```mermaid\nflowchart LR\n subgraph input[\"User specifies\"]\n torch[\"torch==2.1.0\"]\n end\n \n subgraph matrix[\"Compatibility Matrix\"]\n lookup[\"torch_compatibility_matrix.json\"]\n end\n \n subgraph output[\"Cog determines\"]\n cuda[\"CUDA 12.1\"]\n cudnn[\"cuDNN 8\"]\n python[\"Python 3.10-3.13\"]\n end\n \n torch --> lookup\n lookup --> cuda\n lookup --> cudnn\n lookup --> python\n```\n\n**Data files** (embedded JSON, generated by `tools/compatgen/`):\n- `pkg/config/torch_compatibility_matrix.json` - PyTorch ↔ CUDA mappings\n- `pkg/config/tf_compatibility_matrix.json` - TensorFlow ↔ CUDA mappings \n- `pkg/config/cuda_base_images.json` - Available NVIDIA base image tags\n\nThese are regenerated when new framework versions are released and embedded into the CLI binary at build time.\n\n**What it stores** (for each framework release):\n- Framework version (e.g., `torch==2.1.0`)\n- Compatible CUDA versions\n- Compatible cuDNN versions\n- Compatible Python versions\n- Package index URLs (for CUDA-specific wheels)\n\n---\n\n### 4. Base Image Selection\n\nBase image selection uses the compatibility matrix to find a pre-built image that matches the required Python/CUDA/PyTorch combination.\n\n```mermaid\nflowchart TD\n start[\"Config has Python + CUDA + Torch versions\"] --> lookup{\"Matching cog-base
image exists?\"}\n lookup -->|\"Yes\"| cogbase[\"Use cog-base image
r8.im/cog-base:cuda12.1-python3.11-torch2.1.0\"]\n lookup -->|\"No\"| gpu{\"GPU enabled?\"}\n gpu -->|\"Yes\"| cuda[\"Use NVIDIA CUDA image
nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
(install Python + Torch in Dockerfile)\"]\n gpu -->|\"No\"| slim[\"Use Python slim image
python:3.11-slim\"]\n```\n\n#### Cog Base Images\n\nPre-built images hosted at `r8.im/cog-base` with Python, CUDA, cuDNN, and PyTorch already installed.\n\n- Format: `r8.im/cog-base:cuda-python-torch`\n- Generated from the compatibility matrix (`BaseImageConfigurations()`)\n- Includes common system packages (ffmpeg, git, curl, etc.)\n- Faster builds since heavy dependencies are pre-installed\n\n#### Fallback: NVIDIA CUDA Images\n\nWhen no matching cog-base exists (e.g., unusual version combination):\n- Uses official `nvidia/cuda` images\n- Dockerfile installs Python via pyenv\n- Dockerfile installs PyTorch and other packages via pip\n- Slower builds but supports any valid combination\n\n---\n\n### 5. Embedded Python Wheel\n\nThe Cog Python SDK is embedded in the Go binary at compile time and injected into images during build.\n\nDuring build, the wheel is:\n1. Selected based on configuration (see below)\n2. Copied into the Docker build context\n3. Installed via pip\n\n#### Wheel Selection\n\nThe `COG_SDK_WHEEL` environment variable controls which cog SDK wheel is installed:\n\n| Value | Source |\n|-------|--------|\n| (unset) | Latest `cog` from PyPI (auto-detects local wheel in dev builds) |\n| `pypi` | Latest `cog` from PyPI |\n| `pypi:0.18.0` | Specific version from PyPI |\n| `https://...` | Download wheel from URL |\n| `/path/to/file.whl` | Use local wheel file |\n\nThis allows testing development versions of the SDK or pinning to specific releases. The `build.sdk_version` field in `cog.yaml` provides the same version-pinning capability without requiring environment variables.\n\n---\n\n### 6. Post-Build: Labels & Schema\n\nAfter the main build, Cog:\n\n1. **Runs the container** to generate OpenAPI schema\n2. **Runs pip freeze** to capture installed packages\n3. **Applies labels** with metadata\n\n#### Image Labels\n\n| Label | Content |\n|-------|---------|\n| `run.cog.version` | Cog CLI version |\n| `run.cog.config` | Serialized cog.yaml |\n| `run.cog.openapi_schema` | OpenAPI spec from type hints |\n| `run.cog.pip_freeze` | Installed package versions |\n\nThese labels can be fetched from a remote registry or local image store (like containerd) without pulling the full image. This allows tooling - both the Cog CLI during development and production infrastructure - to inspect model metadata and make decisions about how to run a model before booting it.\n\n---\n\n## Image Layer Structure\n\nA built Cog image has layers in this order (bottom to top):\n\n```\n┌─────────────────────────────────────────────────┐\n│ COPY . /src │ ← User code + weights\n├─────────────────────────────────────────────────┤\n│ RUN commands (from cog.yaml) │ ← Custom build steps\n├─────────────────────────────────────────────────┤\n│ pip install (python_packages) │ ← Python dependencies\n├─────────────────────────────────────────────────┤\n│ Cog wheel install │ ← Cog runtime\n├─────────────────────────────────────────────────┤\n│ apt-get install (system_packages) │ ← System dependencies\n├─────────────────────────────────────────────────┤\n│ tini init │ ← Process manager\n├─────────────────────────────────────────────────┤\n│ │\n│ Base image │ ← Largest layer\n│ (OS, Python, CUDA, cuDNN, PyTorch) │ ~5-15 GB for GPU images\n│ │\n└─────────────────────────────────────────────────┘\n```\n\nThe base image is by far the largest layer. Using a matching `cog-base` image means this layer is shared across builds and doesn't need to be re-downloaded or rebuilt.\n\n---\n\n## Code Reference\n\n| Component | Location |\n|-----------|----------|\n| CLI command | `pkg/cli/build.go` |\n| Build orchestration | `pkg/image/build.go` |\n| Dockerfile generator | `pkg/dockerfile/standard_generator.go` |\n| Base image selection | `pkg/dockerfile/base.go` |\n| Compatibility matrix | `pkg/config/compatibility.go` |\n| Embedded wheels | `pkg/wheels/wheels.go` |\n| Label definitions | `pkg/docker/command/manifest.go` |\n" }, { "path": "architecture/06-cli.md", "content": "# CLI\n\nThe Cog CLI is a Go binary that provides commands for the full model lifecycle: development, building, testing, and deployment. This document covers what each command does and how it connects to the systems described in previous docs.\n\n**Important**: Model code always runs inside a container, never on the host machine. Commands like `cog predict`, `cog train`, and `cog serve` build an image, start a container, and interact with it via the [Prediction API](./03-prediction-api.md). The CLI orchestrates this, but the model execution happens in the containerized [Container Runtime](./04-container-runtime.md).\n\n## Commands Overview\n\n| Command | Job To Be Done |\n|---------|----------------|\n| `cog init` | Bootstrap a new model project |\n| `cog build` | Create a container image |\n| `cog predict` | Run a prediction in a container |\n| `cog train` | Run training in a container |\n| `cog run` | Run arbitrary commands in a container |\n| `cog serve` | Start HTTP server in a container |\n| `cog push` | Deploy to Replicate |\n| `cog login` | Authenticate with Replicate |\n\n## Development Commands\n\n### cog init\n\n**Job**: Create a starter `cog.yaml` and `predict.py` for a new model.\n\n```bash\ncog init\n```\n\nCreates:\n- `cog.yaml` with sensible defaults\n- `predict.py` with a skeleton Predictor class\n\n**Code**: `pkg/cli/init.go`\n\n---\n\n### cog predict\n\n**Job**: Run a prediction in a container.\n\n```bash\ncog predict -i prompt=\"A photo of a cat\" -i steps=50\n```\n\nWhat happens:\n1. Builds the image (if needed)\n2. Starts a container running the [Container Runtime](./04-container-runtime.md)\n3. Parses `-i` flags against the [Schema](./02-schema.md)\n4. Sends a [PredictionRequest](./03-prediction-api.md) to the container's HTTP API\n5. Streams output back to terminal\n\nInput types are inferred from the schema:\n- Strings: `-i prompt=\"hello\"`\n- Numbers: `-i steps=50`\n- Files: `-i image=@photo.jpg` (uploaded to container)\n- URLs: `-i image=https://example.com/photo.jpg`\n\n**Code**: `pkg/cli/predict.go`\n\n---\n\n### cog train\n\n**Job**: Run training in a container.\n\n```bash\ncog train -i data=@dataset.zip -i epochs=10\n```\n\nSame as `cog predict` but calls the `train()` method instead of `predict()`.\n\n**Code**: `pkg/cli/train.go`\n\n---\n\n### cog run\n\n**Job**: Run arbitrary commands in a container.\n\n```bash\ncog run python -c \"import torch; print(torch.cuda.is_available())\"\ncog run bash\n```\n\nBuilds the image (if needed), starts a container, and runs the specified command inside it. Useful for:\n- Debugging the container environment\n- Running one-off scripts\n- Interactive exploration\n\n**Code**: `pkg/cli/run.go`\n\n---\n\n### cog serve\n\n**Job**: Start the HTTP server in a container for testing.\n\n```bash\ncog serve\n# Server running at http://localhost:5000\n```\n\nBuilds the image (if needed) and starts a container running the [Container Runtime](./04-container-runtime.md). The container's port 5000 is exposed to the host. You can then:\n- Send requests to `POST /predictions`\n- View Swagger UI at `/docs`\n- Test webhooks\n\n**Code**: `pkg/cli/serve.go`\n\n## Build Commands\n\n### cog build\n\n**Job**: Build a container image from [Model Source](./01-model-source.md).\n\n```bash\ncog build -t my-model\n```\n\nWhat happens (see [Build System](./05-build-system.md) for details):\n\n1. **Parse** `cog.yaml`\n2. **Resolve** CUDA/cuDNN versions from compatibility matrix\n3. **Generate** Dockerfile\n4. **Build** image via Docker/Buildkit\n5. **Run** container to extract [Schema](./02-schema.md)\n6. **Apply** labels (schema, config, pip freeze)\n\nKey flags:\n- `-t, --tag`: Image tag\n- `--no-cache`: Disable Docker cache\n- `--separate-weights`: Exclude weights from image (for separate upload)\n\n**Code**: `pkg/cli/build.go`, `pkg/image/build.go`\n\n## Deployment Commands\n\n### cog push\n\n**Job**: Build and push to Replicate.\n\n```bash\ncog push r8.im/username/model-name\n```\n\nWhat happens:\n1. Builds image (like `cog build`)\n2. Pushes to Replicate's registry\n3. Registers model with Replicate API\n\nThe image tag must be a Replicate model reference (`r8.im/owner/name`).\n\n**Code**: `pkg/cli/push.go`, `pkg/api/client.go`\n\n---\n\n### cog login\n\n**Job**: Authenticate with Replicate.\n\n```bash\ncog login\n# or\ncog login --token-stdin < token.txt\n```\n\nStores credentials for `cog push`.\n\n**Code**: `pkg/cli/login.go`\n\n## How CLI Commands Interact with Containers\n\nCommands like `predict`, `train`, and `serve` follow the same pattern: build an image, start a container, communicate via HTTP. The CLI never runs model code directly.\n\n```mermaid\nsequenceDiagram\n participant CLI as cog CLI (host)\n participant Docker\n participant Container as Container (runtime)\n\n CLI->>CLI: Parse -i flags, load cog.yaml\n CLI->>Docker: Build image (if needed)\n Docker-->>CLI: Image ready\n \n CLI->>Docker: Start container\n Docker->>Container: python -m cog.server.http\n Container->>Container: Run setup()\n \n loop Until READY\n CLI->>Container: GET /health-check\n Container-->>CLI: Status (STARTING/READY)\n end\n \n CLI->>Container: POST /predictions\n Container->>Container: Run predict()\n Container-->>CLI: Response JSON\n \n CLI->>Docker: Stop container\n```\n\nFor what happens inside the container (setup, predict, IPC), see [Container Runtime](./04-container-runtime.md).\n\n## CLI Architecture\n\nThe CLI is built with [Cobra](https://github.com/spf13/cobra) (Go CLI framework).\n\n```\ncmd/cog/\n└── cog.go # Entry point\n\npkg/cli/\n├── root.go # Root command, subcommand registration\n├── build.go # cog build\n├── predict.go # cog predict\n├── train.go # cog train\n├── run.go # cog run\n├── serve.go # cog serve\n├── push.go # cog push\n├── login.go # cog login\n└── init.go # cog init\n```\n\nCommands delegate to packages:\n- `pkg/image/` - Image building\n- `pkg/dockerfile/` - Dockerfile generation\n- `pkg/docker/` - Docker client operations\n- `pkg/config/` - cog.yaml parsing\n- `pkg/api/` - Replicate API client\n- `pkg/predict/` - Local prediction execution\n\n## Code References\n\n| File | Purpose |\n|------|---------|\n| `pkg/cli/root.go` | Command registration |\n| `pkg/cli/build.go` | Build command |\n| `pkg/cli/predict.go` | Predict command, input parsing |\n| `pkg/cli/push.go` | Push command |\n| `pkg/image/build.go` | Build orchestration |\n| `pkg/predict/predictor.go` | Local prediction client |\n" }, { "path": "architecture/ffi/03-prediction-api.md", "content": "# Prediction API (FFI/Rust)\n\nThe FFI runtime implements the same Prediction API as the legacy runtime, using the same envelope format and endpoints. This document highlights FFI-specific behavior and implementation details.\n\n> **Note**: The API surface is identical to the [legacy implementation](../legacy/03-prediction-api.md). Clients don't need to change code when switching runtimes.\n\n## Endpoints\n\n| Endpoint | Method | Purpose | FFI Notes |\n|----------|--------|---------|-----------|\n| `POST /predictions` | Create | Start a new prediction | Uses `SyncPredictionGuard` for automatic cancellation |\n| `PUT /predictions/{id}` | Create (idempotent) | Start or retrieve existing prediction | Concurrent-safe with DashMap |\n| `POST /predictions/{id}/cancel` | Cancel | Cancel a running prediction | Uses cancel tokens propagated to worker |\n| `GET /health-check` | Health | Check server status | Returns health state machine status |\n| `GET /` | Index | List available endpoints | Static route |\n| `GET /openapi.json` | Schema | OpenAPI specification | Cached from worker `Ready` event |\n\n## FFI-Specific Behaviors\n\n### Connection Drop Handling\n\n**Key difference from legacy**: Synchronous predictions automatically cancel when the client connection drops.\n\n```rust\n// SyncPredictionGuard is RAII - drops when connection closes\nlet guard = handle.sync_guard();\nlet result = service.predict(slot, input).await;\n// If connection drops here, guard.drop() cancels the prediction\n```\n\nThis prevents wasted computation on predictions where the client is no longer listening.\n\n### Health States\n\nThe FFI runtime uses a more detailed health state machine. The `/health-check` endpoint always returns HTTP 200 with the status in the JSON body:\n\n| State | JSON `status` | Condition |\n|-------|---------------|-----------|\n| `STARTING` | `\"STARTING\"` | Worker subprocess initializing |\n| `READY` | `\"READY\"` | Worker ready, slots available |\n| `BUSY` | `\"BUSY\"` | All slots occupied (backpressure) |\n| `SETUP_FAILED` | `\"SETUP_FAILED\"` | `setup()` threw exception |\n| `DEFUNCT` | `\"DEFUNCT\"` | Fatal error, worker crashed |\n\n**New behavior**: When all concurrency slots are occupied, new predictions receive `409 Conflict` instead of queuing. Clients should implement retry with backoff.\n\n> **Note**: Prediction endpoints return 503 when health is not `READY`.\n\n### Idempotent PUT Behavior\n\nThe FFI runtime uses a concurrent-safe DashMap for prediction state:\n\n```rust\n// Atomic check-or-insert\nmatch service.get_prediction_response(id) {\n Some(response) => return 202 + response, // Already exists\n None => {\n service.submit_prediction(id, input, webhook); // Create new\n return 202 + starting_state;\n }\n}\n```\n\nThis is fully thread-safe without locks, unlike the legacy runtime which uses Python's asyncio locks.\n\n## Request Flow Differences\n\n### Sync Prediction (POST /predictions)\n\n**Legacy**:\n```python\n# Connection drop has no effect\nresult = await runner.predict(input)\nreturn result\n```\n\n**FFI**:\n```rust\n// Connection drop triggers guard.drop() → cancellation\nlet guard = handle.sync_guard(); // RAII guard\nlet result = service.predict(slot, input).await;\ndrop(guard); // Or automatic on scope exit\nreturn result;\n```\n\n### Async Prediction (Prefer: respond-async)\n\nBehavior is identical to legacy, but implemented differently:\n\n**Legacy**: Uses asyncio tasks\n**FFI**: Uses tokio tasks with cancel tokens\n\n```rust\ntokio::spawn(async move {\n let _result = service.predict(slot, input).await;\n // Prediction state is already updated by predict() internally\n // Webhooks fire automatically from Prediction mutation methods\n service.remove_prediction(id);\n});\n```\n\n### Cancellation Propagation\n\n**Legacy**: Sends `SIGUSR1` signal to child process\n\n**FFI**: Uses IPC message + different strategies for sync vs async predictors:\n\n```\nParent: ControlRequest::Cancel { slot }\n │\n └─▶ Worker: handler.cancel(slot)\n```\n\n**Sync Predictors:**\n```\nhandler.cancel(slot)\n │\n ├─▶ Set CANCEL_REQUESTED flag for slot\n │\n ├─▶ Send SIGUSR1 to self\n │\n └─▶ Signal handler: raise KeyboardInterrupt (if in cancelable region)\n\nPrediction code:\n with CancelableGuard(): # Sets CANCELABLE=true\n predictor.predict() # Can be interrupted\n # CANCELABLE=false on exit\n```\n\n**Async Predictors:**\n```\nhandler.cancel(slot)\n │\n ├─▶ Get future from slot state\n └─▶ future.cancel()\n │\n └─▶ Python raises asyncio.CancelledError\n```\n\nThis provides more reliable cancellation with proper handling for both sync and async execution models.\n\n## Concurrency Model\n\n### Slot-Based Permits\n\nThe FFI runtime uses explicit permit tokens instead of async task limits:\n\n```rust\n// Acquire permit (blocks if all slots busy)\nlet permit = permit_pool.acquire().await?;\n\n// Permit is held during prediction\nlet slot_id = permit.slot_id();\nlet result = orchestrator.predict(slot_id, input).await;\n\n// Permit automatically returned on drop\ndrop(permit); // Or automatic on scope exit\n```\n\n**Advantages**:\n- Fixed, predictable concurrency\n- Fair queuing (FIFO permit acquisition)\n- Observable slot usage in metrics\n- No task explosion\n\n### Configuration\n\n```yaml\n# cog.yaml\nconcurrency:\n max: 5\n```\n\nThis creates 5 slots in the PermitPool. Each slot corresponds to one Unix socket connection to the worker subprocess.\n\n## File Handling\n\nFile handling is identical to legacy (URLs downloaded to temp files, outputs uploaded), but the implementation differs:\n\n**Legacy**: Uses Python `aiohttp` + `requests`\n\n**FFI**: Uses Rust `reqwest` with connection pooling:\n\n```rust\n// Download input files\nlet client = reqwest::Client::builder()\n .connection_pool_idle_timeout(Duration::from_secs(30))\n .build()?;\n\nlet bytes = client.get(url).send().await?.bytes().await?;\ntokio::fs::write(temp_path, bytes).await?;\n```\n\nThis provides better performance for large file downloads.\n\n## Webhooks\n\nWebhook delivery is similar but with improvements:\n\n### Retry Logic\n\n**Legacy**: Simple exponential backoff\n\n**FFI**: Structured retry with observability:\n\n```rust\nlet retry_policy = ExponentialBackoff::builder()\n .max_elapsed_time(Some(Duration::from_secs(60)))\n .build();\n\nlet webhook_sender = WebhookSender::new(client, retry_policy);\nwebhook_sender.send_with_retry(url, payload).await?;\n```\n\n### Trace Context Propagation\n\nThe FFI runtime automatically propagates OpenTelemetry trace context in webhook headers:\n\n```rust\n// Automatic trace propagation\nheaders.insert(\"traceparent\", trace_id);\nheaders.insert(\"tracestate\", trace_state);\n```\n\nThis enables distributed tracing across prediction → webhook → downstream services.\n\n## Status Lifecycle\n\nThe status lifecycle is identical to legacy:\n\n```mermaid\nstateDiagram-v2\n [*] --> starting: Request received\n starting --> processing: predict() called\n processing --> succeeded: predict() returns\n processing --> failed: predict() raises exception\n processing --> canceled: Cancel requested\n succeeded --> [*]\n failed --> [*]\n canceled --> [*]\n```\n\nState transitions happen on the `Prediction` struct directly, which fires webhooks as a side effect:\n\n```rust\n// State transitions fire webhooks automatically\npred.set_processing(); // fires Start webhook\n// ... prediction runs, logs/outputs append ...\npred.set_succeeded(output); // fires terminal Completed webhook\n```\n\n## Dynamic Payload Handling\n\nInput validation and output serialization work the same as legacy:\n\n1. **Parse JSON** → Extract `input` from request body\n2. **Validate against schema** → Pydantic checks types (in worker subprocess)\n3. **Download files** → Rust HTTP client fetches URLs\n4. **Send to worker** → JSON-framed message via Unix socket\n5. **Call predict()** → Python worker executes user code\n6. **Capture output** → Worker sends back via slot channel\n7. **Upload files** → Rust uploads to storage\n8. **Serialize** → Return JSON response\n\nThe key difference is that steps 1, 3, 7, 8 happen in Rust (faster), while steps 2, 5, 6 happen in Python (same as legacy).\n\n## Error Handling\n\n### Worker Crashes\n\n**Legacy**: Parent process becomes unstable, may need restart\n\n**FFI**: Server marks health as `DEFUNCT` but continues serving other endpoints:\n\n```rust\n// Worker process died\nmatch worker.wait().await {\n Ok(status) if !status.success() => {\n health.set(Health::Defunct);\n // HTTP server still runs, returns 503 for predictions\n }\n}\n```\n\n### Setup Failures\n\nBoth runtimes mark the container as unhealthy, but FFI provides more detail:\n\n```rust\n// Detailed setup failure\nmatch control_rx.recv().await? {\n ControlResponse::Failed { error } => {\n health.set(Health::SetupFailed { reason: error });\n // Include error in health-check response\n }\n}\n```\n\n## Performance Characteristics\n\n| Operation | Legacy | FFI | Improvement |\n|-----------|--------|-----|-------------|\n| Request parsing | Pydantic (Python) | serde (Rust) | ~3x faster |\n| File download | aiohttp | reqwest | ~2x faster |\n| Concurrency overhead | asyncio tasks | Tokio + permits | ~50% less memory |\n| Webhook delivery | Sequential retries | Concurrent + backoff | Better throughput |\n| State management | asyncio locks | DashMap (lock-free) | No contention |\n\n## Environment Variables\n\nFFI-specific variables:\n\n| Variable | Default | Purpose |\n|----------|---------|---------|\n| `USE_COGLET` | unset | Enable FFI runtime (set to any value) |\n| `COG_CONCURRENCY_SLOTS` | 1 | Number of prediction slots |\n| `COG_WORKER_TIMEOUT` | 300s | Worker subprocess timeout |\n| `RUST_LOG` | info | Rust logging (tracing crate) |\n\nLegacy variables like `COG_MAX_CONCURRENCY` are ignored when using FFI.\n\n## Code References\n\n| File | Purpose |\n|------|---------|\n| `crates/coglet/src/transport/http/routes.rs` | HTTP endpoint handlers |\n| `crates/coglet/src/prediction.rs` | Prediction state + webhook firing |\n| `crates/coglet/src/webhook.rs` | Webhook delivery with retries |\n| `crates/coglet/src/bridge/protocol.rs` | IPC message types |\n| `crates/coglet/src/permit/pool.rs` | Slot-based concurrency |\n\n## Migration Notes\n\nWhen switching from legacy to FFI runtime:\n\n✅ **No changes needed**:\n- HTTP API endpoints\n- Request/response format\n- Predictor code\n- Client code\n\n⚠️ **Behavioral differences**:\n- Sync predictions cancel on connection drop\n- 409 responses when at capacity (not queuing)\n- Different health state granularity\n- Different environment variables\n\n📈 **Improvements**:\n- ~2x faster HTTP layer\n- Better resource management\n- More reliable cancellation\n- Worker crash resilience\n" }, { "path": "architecture/ffi/04-container-runtime.md", "content": "# Container Runtime (FFI/Rust)\n\nThis document covers the FFI runtime implementation using Rust with PyO3 bindings. This is a complete rewrite of the HTTP server, moving from Python/FastAPI to Rust/Axum with a PyO3 ABI3 wheel.\n\n## Overview\n\nThe FFI runtime provides significant improvements over the legacy Python runtime:\n- **Rust HTTP server (Axum)**: Faster request handling, better backpressure management\n- **Worker isolation**: Python predictor crashes don't kill the server\n- **Slot-based concurrency**: Predictable resource control with permit pools\n- **Same API surface**: Drop-in replacement for the legacy runtime\n- **Subprocess reuse**: Predictor stays loaded between requests\n\n## High-Level Architecture\n\n```\n┌─────────────────────────────────────────────────────────────────────────────────┐\n│ HTTP Transport (axum) │\n│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ ┌─────────────────────────┐ │\n│ │ POST │ │ PUT │ │ POST │ │ GET │ │\n│ │ /predictions│ │ /predictions│ │ /cancel │ │ /health-check │ │\n│ │ │ │ /{id} │ │ │ │ /openapi.json │ │\n│ └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ └───────────┬─────────────┘ │\n└─────────┼────────────────┼────────────────┼─────────────────────┼───────────────┘\n │ │ │ │\n ▼ ▼ ▼ ▼\n┌─────────────────────────────────────────────────────────────────────────────────┐\n│ PredictionService │\n│ ┌────────────────────────────────────────────────────────────────────────────┐ │\n│ │ Active Predictions (DashMap) │ │\n│ │ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │ │\n│ │ �� PredictionEntry │ │ PredictionEntry │ │ PredictionEntry │ ... │ │\n│ │ │ ─────────────── │ │ ─────────────── │ │ ─────────────── │ │ │\n│ │ │ prediction (Arc)│ │ prediction (Arc)│ │ prediction (Arc)│ │ │\n│ │ │ cancel_token │ │ cancel_token │ │ cancel_token │ │ │\n│ │ │ input │ │ input │ │ input │ │ │\n│ │ └─────────────────┘ └─────────────────┘ └─────────────────┘ │ │\n│ └────────────────────────────────────────────────────────────────────────────┘ │\n│ │\n│ ┌────────────────────────────────────────────────────────────────────────────┐ │\n│ │ PermitPool │ │\n│ │ ┌────────┐ ┌────────┐ ┌────────┐ │ │\n│ │ │ Permit │ │ Permit │ │ Permit │ (concurrency control) │ │\n│ │ │ slot_0 │ │ slot_1 │ │ slot_2 │ │ │\n│ │ └────────┘ └────────┘ └────────┘ │ │\n│ └────────────────────────────────────────────────────────────────────────────┘ │\n│ │\n│ ┌────────────────────────────────────────────────────────────────────────────┐ │\n│ │ OrchestratorHandle │ │\n│ │ (slot_ids, control_tx for worker comms) │ │\n│ └────────────────────────────────────────────────────────────────────────────┘ │\n└──────────────────────────────────┬──────────────────────────────────────────────┘\n │\n Unix Socket (slot) + stdin/stdout (control)\n │\n ▼\n┌─────────────────────────────────────────────────────────────────────────────────┐\n│ Worker Subprocess (Python) │\n│ ┌────────────────────────────────────────────────────────────────────────────┐ │\n│ │ Predictor │ │\n│ │ ┌─────────────────────────────────────────────────────────────────────┐ │ │\n│ │ │ setup() → runs once at startup │ │ │\n│ │ │ predict() → handles SlotRequest::Predict │ │ │\n│ │ └─────────────────────────────────────────────────────────────────────┘ │ │\n│ └────────────────────────────────────────────────────────────────────────────┘ │\n└─────────────────────────────────────────────────────────────────────────────────┘\n```\n\n## Component Ownership\n\nThe FFI runtime uses clear ownership patterns to manage prediction lifecycle:\n\n```\n═══════════════════════════════════════════════════════════════════════════════════\n COMPONENT OWNERSHIP\n═══════════════════════════════════════════════════════════════════════════════════\n PredictionService (single owner of prediction state)\n ├── owns: DashMap (active predictions)\n ├── owns: OrchestratorState (pool + orchestrator handle)\n ├── owns: health, setup_result, schema\n └── method: cancel() fires token + delegates to orchestrator\n\n PredictionEntry (in DashMap)\n ├── has: Arc> (the real state — single source of truth)\n ├── has: CancellationToken\n └── has: input (for API responses)\n\n Prediction (state machine — webhooks fire from mutation methods)\n ├── owns: status, logs, outputs, output, error, metrics\n ├── owns: WebhookSender (fires on set_processing, append_log, etc.)\n └── owns: completion notifier (for waiting on result)\n\n PredictionSlot (RAII container)\n ├── owns: Arc> (shared with DashMap entry)\n ├── owns: Permit (concurrency token, returns to pool on drop)\n └── Drop: marks permit idle, releases back to pool\n\n PredictionHandle (returned to route handler)\n ├── has: CancellationToken clone\n └── method: sync_guard(service) → SyncPredictionGuard (cancels on drop)\n\n Cancellation (via OrchestratorHandle)\n ├── Sync predictors: ControlRequest::Cancel → SIGUSR1 → KeyboardInterrupt\n └── Async predictors: ControlRequest::Cancel → future.cancel() → CancelledError\n═══════════════════════════════════════════════════════════════════════════════════\n```\n\n## Worker Subprocess Protocol\n\nCommunication between the Rust server and Python worker uses two channels:\n\n### Control Channel (stdin/stdout - JSON framed)\n\n| Parent → Child | Child → Parent |\n|----------------|----------------|\n| `Init { predictor_ref, num_slots, ... }` | `Ready { slots, schema }` |\n| `Cancel { slot }` | `Log { source, data }` |\n| `Shutdown` | `Idle { slot }` |\n| | `Failed { slot, error }` |\n| | `ShuttingDown` |\n\n### Slot Channel (Unix socket per slot - JSON framed)\n\n| Parent → Child | Child → Parent |\n|----------------|----------------|\n| `Predict { id, input }` | `Log { data }` |\n| | `Output { value }` (streaming) |\n| | `Done { output }` |\n| | `Failed { error }` |\n| | `Cancelled` |\n\n## Health State Machine\n\n```mermaid\nstateDiagram-v2\n [*] --> STARTING: Container start\n note right of STARTING: Predictions return 503\n \n STARTING --> READY: setup() succeeds\n STARTING --> SETUP_FAILED: setup() raises exception\n \n READY --> BUSY: All slots occupied\n note right of BUSY: New predictions get 409\n \n BUSY --> READY: Slot freed\n \n READY --> DEFUNCT: Fatal error / worker crash\n BUSY --> DEFUNCT: Fatal error / worker crash\n note right of DEFUNCT: Predictions return 503\n \n SETUP_FAILED --> [*]\n DEFUNCT --> [*]\n```\n\n### Health States\n\nThe health-check endpoint always returns HTTP 200 with the status in the JSON body. This allows load balancers and orchestrators to distinguish between \"server is running but not ready\" vs \"server is down\".\n\n| State | JSON `status` | Meaning |\n|-------|---------------|---------|\n| `STARTING` | `\"STARTING\"` | Worker subprocess initializing, `setup()` running |\n| `READY` | `\"READY\"` | Worker ready, at least one slot available |\n| `BUSY` | `\"BUSY\"` | All slots occupied, no capacity for new predictions |\n| `SETUP_FAILED` | `\"SETUP_FAILED\"` | `setup()` threw exception, cannot serve predictions |\n| `DEFUNCT` | `\"DEFUNCT\"` | Fatal error or worker crash, server unusable |\n\n> **Note**: Prediction endpoints (`/predictions`) return 503 when health is not `READY`.\n\n## Prediction Flow\n\n### Sync Request (POST /predictions)\n\n```mermaid\nsequenceDiagram\n participant Client\n participant Routes\n participant Service\n participant Worker\n \n Client->>Routes: POST /predictions\n Routes->>Service: submit_prediction(id, input, webhook)\n Service-->>Routes: PredictionHandle + slot\n \n Note over Routes: SyncPredictionGuard held
(cancels on connection drop)\n \n Routes->>Service: predict(slot, input)\n Service->>Worker: predict(slot, input)\n Worker-->>Service: result\n Note over Service: Prediction.set_succeeded() fires webhook\n \n Routes-->>Client: 200 {output}\n```\n\n**Key behavior**: The `SyncPredictionGuard` is held for the duration of the request. If the client connection drops, the guard is dropped and the prediction is automatically cancelled.\n\n### Async Request (Prefer: respond-async)\n\n```mermaid\nsequenceDiagram\n participant Client\n participant Routes\n participant Service\n participant Worker\n \n Client->>Routes: POST + respond-async\n Routes->>Service: submit_prediction(id, input, webhook)\n Service-->>Routes: PredictionHandle + slot\n \n Routes-->>Client: 202 {status: \"starting\"}\n \n Note over Routes,Worker: spawned task continues independently\n \n par Background Task\n Service->>Worker: predict(slot, input)\n Worker-->>Service: result\n Note over Service: Prediction mutations fire webhooks automatically\n end\n \n Service-->>Client: webhook (completed)\n```\n\n**Key behavior**: No guard is held. The prediction continues even if the client disconnects.\n\n### Idempotent PUT (PUT /predictions/{id})\n\n```mermaid\nsequenceDiagram\n participant Client\n participant Routes\n participant Service\n \n Client->>Routes: PUT /predictions/X\n Routes->>Service: get_prediction_response(\"X\")\n \n alt Prediction exists\n Service-->>Routes: existing state\n Routes-->>Client: 202 + full state\n else Prediction doesn't exist\n Routes->>Service: submit_prediction + predict\n Routes-->>Client: 202 + starting state\n end\n```\n\n### Connection Drop (Sync Mode)\n\n```mermaid\nsequenceDiagram\n participant Client\n participant Routes\n participant Service\n participant Worker\n \n Client->>Routes: POST /predictions\n Note over Routes: SyncPredictionGuard armed\n Routes->>Worker: predict(slot)\n \n Client-xRoutes: ✕ connection drops\n \n Note over Routes: guard.drop()\n Routes->>Service: cancel(id)\n Service->>Worker: Cancel\n Worker-->>Service: Cancelled\n```\n\n## File Structure\n\n```\ncrates/coglet/src/\n├── lib.rs # Public API exports\n├── service.rs # PredictionService (single owner of prediction state)\n├── prediction.rs # Prediction state (logs, outputs, status)\n├── health.rs # Health enum + SetupResult\n├── orchestrator.rs # Worker subprocess management\n├── permit/\n│ ├── mod.rs\n│ ├── pool.rs # PermitPool (concurrency control)\n│ └── slot.rs # PredictionSlot (Prediction + Permit RAII)\n├── bridge/\n│ ├── mod.rs\n│ ├── protocol.rs # Control/Slot request/response types\n│ ├── codec.rs # JSON length-delimited framing\n│ └── transport.rs # Unix socket transport\n├── transport/\n│ └── http/\n│ ├── mod.rs\n│ ├── server.rs # Axum server setup\n│ └── routes.rs # HTTP handlers (uses service)\n├── webhook.rs # WebhookSender (retry logic, trace context)\n├── worker.rs # run_worker, PredictHandler trait, SetupError\n└── version.rs # VersionInfo\n\ncrates/coglet-python/src/\n└── lib.rs # PyO3 bindings (coglet.server.serve())\n```\n\n## Invocation Path\n\nHow coglet gets invoked when running a Cog container:\n\n```\n┌─────────────────────────────────────────────────────────────────────────────┐\n│ cog predict / cog run │\n│ (CLI) │\n└─────────────────────────────────┬───────────────────────────────────────────┘\n │\n ▼\n┌─────────────────────────────────────────────────────────────────────────────┐\n│ python -m cog.server.http │\n│ │\n│ if USE_COGLET env var: │\n│ import coglet │\n│ coglet.server.serve(predictor_ref, port=5000) ──────────────────┐ │\n│ else: │ │\n│ # original Python FastAPI server │ │\n│ uvicorn.run(app, port=5000) │ │\n└─────────────────────────────────────────────────────────────────────────┼───┘\n │\n ▼\n┌─────────────────────────────────────────────────────────────────────────────┐\n│ coglet (Rust) │\n│ │\n│ ┌───────────────────────────────────────────────────────────────────┐ │\n│ │ HTTP Server (axum) :5000 │ │\n│ │ /predictions, /health-check, etc. │ │\n│ └───────────────────────────────────────────────────────────────────┘ │\n│ │ │\n│ ▼ │\n│ ┌───────────────────────────────────────────────────────────────────┐ │\n│ │ PredictionService (state, webhooks, permits) │ │\n│ └───────────────────────────────────────────────────────────────────┘ │\n│ │ │\n│ Unix socket + pipes │\n│ │ │\n│ ▼ │\n│ ┌───────────────────────────────────────────────────────────────────┐ │\n│ │ Worker subprocess (Python) │ │\n│ │ - loads predictor_ref │ │\n│ │ - runs setup() │ │\n│ │ - handles predict() requests │ │\n│ └───────────────────────────────────────────────────────────────────┘ │\n└─────────────────────────────────────────────────────────────────────────────┘\n```\n\n## Key Design Decisions\n\n### Why Rust?\n- **Performance**: Axum is faster than Uvicorn/FastAPI for HTTP handling\n- **Stability**: Server doesn't crash when user code fails\n- **Resource management**: Better backpressure and concurrency control\n- **Memory safety**: No Python GIL contention in HTTP layer\n\n### Why PyO3 FFI?\n- **ABI3 wheel**: Single wheel works across Python 3.10-3.13\n- **Native performance**: Direct C API calls, no serialization overhead\n- **Same predictor code**: Users don't change anything\n- **Drop-in replacement**: Same HTTP API, same behavior\n\n### Why Subprocess (not in-process)?\n- **Isolation**: Python crashes/segfaults don't kill server\n- **CUDA context**: Clean GPU initialization per worker\n- **Memory**: Fresh address space for model loading\n- **Restart potential**: Architecture enables future worker restart on fatal errors (not yet implemented)\n\n### Why Slots (not async tasks)?\n- **Predictable**: Fixed number of concurrent predictions\n- **Fair**: Permits prevent starvation\n- **Observable**: Easy to monitor slot usage\n- **Simple**: No async complexity in worker subprocess\n\n## Environment Variables\n\n| Variable | Default | Purpose |\n|----------|---------|---------|\n| `USE_COGLET` | unset | Enable FFI runtime (set to any value) |\n| `PORT` | 5000 | HTTP server port |\n| `COG_LOG_LEVEL` | INFO | Logging verbosity |\n| `COG_CONCURRENCY_SLOTS` | 1 | Number of concurrent prediction slots |\n\n## Comparison to Legacy Runtime\n\n| Aspect | Legacy (Python) | FFI (Rust) |\n|--------|----------------|------------|\n| HTTP Server | FastAPI/Uvicorn | Axum |\n| Language | Pure Python | Rust + PyO3 |\n| IPC | multiprocessing.Pipe (pickled) | Unix socket + pipes (JSON) |\n| Concurrency | async tasks | Slot-based permits |\n| Cancellation | SIGUSR1 signal | IPC message + SIGUSR1 (sync) / future.cancel() (async) |\n| Connection drop | No effect on prediction | Cancels sync predictions |\n| Worker crash | Server unstable | Server stays up, marks DEFUNCT |\n| Performance | Baseline | ~2x faster HTTP layer |\n\n## Code References\n\n| File | Purpose |\n|------|---------|\n| `crates/coglet/src/service.rs` | Main orchestrator: PredictionService |\n| `crates/coglet/src/prediction.rs` | Prediction state machine + webhook firing |\n| `crates/coglet/src/transport/http/routes.rs` | HTTP endpoint handlers |\n| `crates/coglet/src/permit/pool.rs` | Slot-based concurrency control |\n| `crates/coglet/src/orchestrator.rs` | Worker subprocess spawn/management |\n| `crates/coglet/src/bridge/protocol.rs` | IPC message definitions |\n| `crates/coglet-python/src/lib.rs` | PyO3 Python bindings |\n| `python/cog/server/http.py` | Entry point (checks USE_COGLET) |\n" }, { "path": "architecture/ffi/README.md", "content": "# FFI Runtime (Rust + PyO3)\n\nThis directory documents the next-generation Cog runtime implementation using Rust with PyO3 FFI bindings.\n\n## Status\n\nThis is an **experimental** runtime implementation currently in development. It provides significant improvements in:\n- Performance and resource management\n- Worker process isolation and stability\n- Concurrency control with slot-based permits\n- Graceful cancellation and connection drop handling\n\n## When to Use\n\nEnable this implementation by setting the `USE_COGLET` environment variable when running Cog containers.\n\n## Key Improvements\n\n- **Rust HTTP server (Axum)**: Faster, better backpressure handling\n- **Worker isolation**: Python crashes don't kill the server\n- **Slot-based concurrency**: Predictable resource management with permit pool\n- **Subprocess reuse**: Predictor stays loaded between requests\n- **Better cancellation**: Sync predictions cancel on connection drop via RAII guards\n\n## Architecture Overview\n\n```\nHTTP Server (Rust/Axum)\n ↓\nPredictionService (state, webhooks, DashMap)\n ↓\nPermitPool (slot-based concurrency)\n ↓\nOrchestrator → Worker Subprocess (Python)\n ↓ (Unix socket + pipes)\nPredictor (setup/predict)\n```\n\n## Documentation\n\n- [Prediction API](./03-prediction-api.md) - HTTP endpoints with coglet-specific behavior\n- [Container Runtime](./04-container-runtime.md) - Complete FFI architecture and flow\n\n## Implementation\n\nPrimary code location: `crates/coglet/`\n- `src/transport/http/` - Axum HTTP server\n- `src/service.rs` - PredictionService (single owner of prediction state)\n- `src/permit/` - Slot-based concurrency control\n- `src/orchestrator.rs` - Worker subprocess management\n- `src/bridge/` - IPC protocol and transport\n- `src/worker/` - Worker implementation\n\nPython bindings: `crates/coglet-python/src/lib.rs`\n" }, { "path": "architecture/legacy/03-prediction-api.md", "content": "# Prediction API\n\nThe Prediction API is the HTTP interface for running model inference. It uses a fixed **envelope format** that wraps model-specific inputs and outputs, allowing a uniform API across all Cog models.\n\n## Endpoints\n\n| Endpoint | Method | Purpose |\n|----------|--------|---------|\n| `POST /predictions` | Create | Start a new prediction |\n| `PUT /predictions/{id}` | Create (idempotent) | Start or retrieve existing prediction |\n| `POST /predictions/{id}/cancel` | Cancel | Cancel a running prediction |\n| `GET /health-check` | Health | Check server status |\n| `GET /` | Index | List available endpoints |\n| `GET /openapi.json` | Schema | OpenAPI specification |\n\nBy default, `POST /predictions` blocks until completion. For long-running predictions, use async mode with `Prefer: respond-async` header - the response returns immediately with status `processing`, and progress updates are delivered via webhook.\n\n## The Envelope Pattern\n\nEvery Cog model exposes the same endpoints with the same request/response structure. The model-specific parts (input fields, output type) are defined by the [Schema](./02-schema.md) and validated at runtime.\n\n```\n┌────────────────────────────────────────────────────────┐\n│ Fixed Envelope (same for all models) │\n│ ┌──────────────────────────────────────────────────┐ │\n│ │ id, status, created_at, logs, metrics, ... │ │\n│ └──────────────────────────────────────────────────┘ │\n│ ┌──────────────────────────────────────────────────┐ │\n│ │ input: { ... } ← model-specific (from schema)│ │\n│ └──────────────────────────────────────────────────┘ │\n│ ┌──────────────────────────────────────────────────┐ │\n│ │ output: ... ← model-specific (from schema)│ │\n│ └──────────────────────────────────────────────────┘ │\n└────────────────────────────────────────────────────────┘\n```\n\nThis pattern means:\n- Clients use the same code to call any Cog model\n- Platforms can route requests without understanding model internals\n- Input validation is schema-driven, not hardcoded\n\n## PredictionRequest\n\nWhat clients send to start a prediction:\n\n```json\n{\n \"id\": \"abc-123\",\n \"input\": {\n \"prompt\": \"A photo of a cat\",\n \"steps\": 50\n },\n \"webhook\": \"https://example.com/webhook\",\n \"webhook_events_filter\": [\"start\", \"output\", \"logs\", \"completed\"]\n}\n```\n\n| Field | Type | Purpose |\n|-------|------|---------|\n| `id` | string (optional) | Client-provided ID for idempotency |\n| `input` | object | **Model-specific** - validated against schema |\n| `webhook` | URL (optional) | Where to send progress updates |\n| `webhook_events_filter` | array (optional) | Which events to send |\n| `created_at` | datetime (optional) | Client-provided timestamp |\n\nThe `input` object is validated against the `Input` schema generated from the predictor's `predict()` signature. Unknown fields are rejected; missing required fields raise validation errors.\n\n## PredictionResponse\n\nWhat comes back from the API:\n\n```json\n{\n \"id\": \"abc-123\",\n \"status\": \"succeeded\",\n \"input\": {\n \"prompt\": \"A photo of a cat\",\n \"steps\": 50\n },\n \"output\": \"https://storage.example.com/output.png\",\n \"logs\": \"Loading model...\\nGenerating image...\\nDone.\",\n \"error\": null,\n \"metrics\": {\n \"predict_time\": 4.52\n },\n \"created_at\": \"2024-01-15T10:30:00Z\",\n \"started_at\": \"2024-01-15T10:30:01Z\",\n \"completed_at\": \"2024-01-15T10:30:05Z\"\n}\n```\n\n| Field | Type | Purpose |\n|-------|------|---------|\n| `id` | string | Prediction identifier |\n| `status` | enum | `starting`, `processing`, `succeeded`, `canceled`, `failed` |\n| `input` | object | Echo of the input (for reference) |\n| `output` | any | **Model-specific** - type defined by schema |\n| `logs` | string | Captured stdout/stderr from predict() |\n| `error` | string | Error message if status is `failed` |\n| `metrics` | object | Timing and other metrics |\n| `created_at` | datetime | When request was received |\n| `started_at` | datetime | When prediction began |\n| `completed_at` | datetime | When prediction finished |\n\n## Status Lifecycle\n\n```mermaid\nstateDiagram-v2\n [*] --> starting: Request received\n starting --> processing: predict() called\n processing --> succeeded: predict() returns\n processing --> failed: predict() raises exception\n processing --> canceled: Cancel requested\n succeeded --> [*]\n failed --> [*]\n canceled --> [*]\n```\n\n## Dynamic Payload Handling\n\nThe magic of the envelope pattern is that the `input` and `output` fields are dynamically typed based on the schema.\n\n### Input Validation Flow\n\n```mermaid\nflowchart LR\n subgraph request[\"Incoming Request\"]\n json[\"JSON body\"]\n end\n \n subgraph validation[\"Validation\"]\n schema[\"Schema (Input type)\"]\n pydantic[\"Pydantic Model\"]\n end\n \n subgraph transform[\"Transformation\"]\n download[\"Download URLs → Files\"]\n coerce[\"Type Coercion\"]\n end\n \n subgraph predict[\"predict()\"]\n kwargs[\"**kwargs\"]\n end\n \n json --> pydantic\n schema --> pydantic\n pydantic --> download\n download --> coerce\n coerce --> kwargs\n```\n\n1. **Parse JSON** - Extract `input` from request body\n2. **Validate against schema** - Pydantic checks types, required fields, constraints\n3. **Download files** - URLs in `cog.Path` fields are fetched to local temp files\n4. **Coerce types** - Strings become Paths, etc.\n5. **Call predict()** - Validated input passed as `**kwargs`\n\n### Output Handling Flow\n\n```mermaid\nflowchart LR\n subgraph predict[\"predict()\"]\n result[\"Return value / yields\"]\n end\n \n subgraph transform[\"Transformation\"]\n upload[\"Upload files → URLs\"]\n serialize[\"JSON serialization\"]\n end\n \n subgraph response[\"Response\"]\n output[\"output field\"]\n end\n \n result --> upload\n upload --> serialize\n serialize --> output\n```\n\n1. **Capture output** - Return value or yielded values from predict()\n2. **Upload files** - `cog.Path` outputs are uploaded, replaced with URLs\n3. **Serialize** - Convert to JSON-compatible format\n4. **Return** - Place in `output` field of response\n\n### File Handling\n\nInput files (cog.Path):\n```\nClient sends: {\"input\": {\"image\": \"https://example.com/photo.jpg\"}}\nServer downloads: /tmp/inputabc123.jpg\npredict() sees: image = Path(\"/tmp/inputabc123.jpg\")\n```\n\nOutput files (cog.Path):\n```\npredict() returns: Path(\"/tmp/output.png\")\nServer uploads: https://storage.example.com/output-xyz.png\nClient receives: {\"output\": \"https://storage.example.com/output-xyz.png\"}\n```\n\n## Webhooks\n\nFor async predictions, progress is delivered via webhooks:\n\n```mermaid\nsequenceDiagram\n participant Client\n participant Cog\n participant Webhook\n \n Client->>Cog: POST /predictions (Prefer: respond-async)\n Cog-->>Client: 202 {status: \"starting\"}\n \n Cog->>Webhook: {status: \"starting\"}\n Note over Cog: predict() starts\n Cog->>Webhook: {status: \"processing\"}\n \n loop Output yields\n Cog->>Webhook: {output: \"partial...\", logs: \"...\"}\n end\n \n Cog->>Webhook: {status: \"succeeded\", output: \"final\"}\n```\n\n### Webhook Events\n\n| Event | When | Payload Contains |\n|-------|------|------------------|\n| `start` | Prediction begins | `status: starting` |\n| `output` | Each yield from iterator | Partial `output` |\n| `logs` | Log lines captured | Updated `logs` |\n| `completed` | Prediction finishes | Final `status`, `output`, `metrics` |\n\nFilter events with `webhook_events_filter`:\n```json\n{\n \"input\": {...},\n \"webhook\": \"https://...\",\n \"webhook_events_filter\": [\"completed\"]\n}\n```\n\n## Streaming Output\n\nFor models that yield output progressively:\n\n```python\ndef predict(self, prompt: str) -> Iterator[str]:\n for token in generate(prompt):\n yield token\n```\n\nThe API can deliver these as:\n\n1. **Webhooks** - Each yield triggers an `output` webhook\n2. **Server-Sent Events** - Stream via `Accept: text/event-stream`\n3. **Final array** - Sync response collects all yields into `output: [\"a\", \"b\", \"c\"]`\n\n## Training API\n\nThe training API (`/trainings`) uses the same envelope pattern:\n\n- `TrainingRequest` extends `PredictionRequest`\n- `TrainingResponse` extends `PredictionResponse`\n- Calls `train()` method instead of `predict()`\n\n## Code References\n\n| File | Purpose |\n|------|---------|\n| `python/cog/schema.py` | `PredictionRequest`, `PredictionResponse`, `Status` |\n| `python/cog/server/http.py` | HTTP endpoints, request handling |\n| `python/cog/server/runner.py` | Prediction orchestration |\n| `python/cog/server/webhook.py` | Webhook delivery |\n" }, { "path": "architecture/legacy/04-container-runtime.md", "content": "# Container Runtime\n\nThis document covers what happens when a Cog container runs. It's where the [Model Source](./01-model-source.md), [Schema](./02-schema.md), and [Prediction API](./03-prediction-api.md) come together.\n\n## Overview\n\nWhen a Cog container runs, it executes a **two-process architecture** with a minimal init system. The design isolates user model code from the HTTP server for stability, resource management, and clean shutdown handling.\n\n## High-Level Architecture\n\n```mermaid\nflowchart TB\n subgraph container[\"Cog Container\"]\n subgraph init[\"tini (PID 1)\"]\n tini[\"Signal forwarding & zombie reaping\"]\n end\n \n subgraph parent[\"Parent Process (HTTP Server)\"]\n direction TB\n subgraph components[\"Components\"]\n direction LR\n fastapi[\"FastAPI/Uvicorn
port 5000\"]\n worker[\"Worker
(parent-side)\"]\n runner[\"PredictionRunner
(orchestrator)\"]\n end\n subgraph threads[\"Thread Pools\"]\n direction LR\n t1[\"Event consumer\"]\n t2[\"Prediction start\"]\n t3[\"Input download
(8 threads)\"]\n end\n end\n \n pipe[[\"multiprocessing.Pipe
(bidirectional IPC)\"]]\n \n subgraph child[\"Child Process (_ChildWorker)\"]\n direction TB\n subgraph child_components[\"Components\"]\n direction LR\n predictor[\"User Predictor
(predict.py)
---
setup()
predict()
train()\"]\n redirector[\"StreamRedirector
stdout/stderr
capture\"]\n eventloop[\"Event Loop
(sync/async)\"]\n end\n end\n end\n \n init --> parent\n parent <--> pipe\n pipe <--> child\n```\n\n## Process Roles\n\n### tini (PID 1)\n- **What**: Minimal init system (~30KB binary)\n- **Why**: Proper signal forwarding to children, zombie process reaping\n- **Entry**: `ENTRYPOINT [\"/sbin/tini\", \"--\"]`\n\n### Parent Process (HTTP Server)\n- **What**: Python process running FastAPI/Uvicorn\n- **Entry**: `CMD [\"python\", \"-m\", \"cog.server.http\"]`\n- **Responsibilities**:\n - HTTP API on port 5000\n - Request validation (Pydantic)\n - Input file downloading (from URLs)\n - Webhook delivery\n - Output file uploads\n - Health state management\n - Child process lifecycle\n\n### Child Process (_ChildWorker)\n- **What**: Isolated Python process for user code\n- **Spawned via**: `multiprocessing.get_context(\"spawn\").Process`\n- **Responsibilities**:\n - Load user's predictor module\n - Run `setup()` once at startup\n - Execute `predict()` / `train()` methods\n - Capture stdout/stderr\n - Send events back to parent\n\n## Why Two Processes?\n\n1. **Isolation**: User code crashes don't bring down the HTTP server\n2. **Memory**: Fresh address space for each model load (spawn vs fork)\n3. **CUDA**: Clean GPU context initialization in child\n4. **Cleanup**: Parent can restart child if it dies\n5. **Monitoring**: Parent tracks child health independently\n\n## Inter-Process Communication\n\n```mermaid\nflowchart LR\n subgraph parent[\"Parent Process\"]\n Worker\n end\n \n subgraph child[\"Child Process\"]\n ChildWorker[\"_ChildWorker\"]\n end\n \n Worker -->|\"PredictionInput
Cancel
Shutdown\"| ChildWorker\n ChildWorker -->|\"Log
PredictionOutput
PredictionOutputType
PredictionMetric
Done\"| Worker\n```\n\nCommunication uses Python's `multiprocessing.Pipe()` with pickled `Envelope` objects:\n\n```python\n@define\nclass Envelope:\n event: Union[Cancel, PredictionInput, Shutdown, Log, ...]\n tag: Optional[str] = None # Routes concurrent predictions\n```\n\n### Event Types\n\n| Event | Direction | Purpose |\n|-------|-----------|---------|\n| `PredictionInput` | Parent → Child | Start prediction with input payload |\n| `Cancel` | Parent → Child | Abort the current prediction |\n| `Shutdown` | Parent → Child | Graceful termination signal |\n| `PredictionOutputType` | Child → Parent | Declares the output type (once per prediction) |\n| `PredictionOutput` | Child → Parent | Output value (multiple for generators) |\n| `Log` | Child → Parent | Captured stdout/stderr line |\n| `PredictionMetric` | Child → Parent | Timing/performance metrics |\n| `Done` | Child → Parent | Prediction complete (success or failure) |\n\n## Request Flow: Prediction Lifecycle\n\n```mermaid\nsequenceDiagram\n participant Client\n participant FastAPI\n participant Runner as PredictionRunner\n participant Worker as Worker (parent)\n participant Pool as ThreadPool\n participant Child as _ChildWorker\n participant Predictor as User predict()\n\n Client->>FastAPI: POST /predictions
{\"input\": {\"prompt\": \"...\"}}\n FastAPI->>Runner: predict(request)\n Runner->>Worker: predict(payload, tag)\n \n Worker->>Pool: Download input URLs\n Pool-->>Worker: Local file paths\n \n Worker->>Child: PredictionInput event\n Child->>Predictor: predict(**payload)\n \n loop Generator yields / prints\n Predictor-->>Child: yield output / print()\n Child-->>Worker: PredictionOutput / Log events\n Worker-->>Runner: handle_event()\n Runner-->>Client: Webhook (if configured)\n end\n \n Predictor-->>Child: return\n Child-->>Worker: Done event\n Worker-->>Runner: handle_event()\n \n Runner->>Runner: Upload output files\n Runner->>Client: Send final webhook\n \n Runner-->>FastAPI: PredictTask complete\n FastAPI-->>Client: Response JSON
{\"output\": \"...\", \"status\": \"succeeded\"}\n```\n\n## Key Components Deep Dive\n\n### HTTP Server (`http.py`)\n\n| Endpoint | Method | Purpose |\n|----------|--------|---------|\n| `/` | GET | API index |\n| `/health-check` | GET | Health status |\n| `/predictions` | POST | New prediction |\n| `/predictions/{id}` | PUT | Idempotent create |\n| `/predictions/{id}/cancel` | POST | Cancel running |\n| `/shutdown` | POST | Graceful shutdown |\n\n### Health States\n\n```mermaid\nstateDiagram-v2\n [*] --> STARTING: Container start\n \n STARTING --> READY: setup() succeeds\n STARTING --> SETUP_FAILED: setup() raises exception\n \n READY --> BUSY: prediction starts\n BUSY --> READY: prediction completes\n \n READY --> DEFUNCT: child dies unexpectedly\n BUSY --> DEFUNCT: child dies unexpectedly\n \n SETUP_FAILED --> [*]\n DEFUNCT --> [*]\n```\n\n### StreamRedirector (Output Capture)\n\nThe child process captures stdout/stderr including native library output (CUDA, etc.):\n\n```mermaid\nflowchart LR\n subgraph child[\"Child Process\"]\n subgraph usercode[\"User Code\"]\n predict[\"predict()\"]\n end\n \n subgraph redirector[\"StreamRedirector\"]\n original[\"Original fd 1/2
(saved)\"]\n pipewrite[\"Pipe write end
(replaces fd 1/2)\"]\n reader[\"Reader Thread\"]\n end\n \n predict -->|\"print()
CUDA logs\"| pipewrite\n pipewrite --> reader\n end\n \n reader -->|\"Log events\"| parent[\"To Parent Process\"]\n```\n\n## Concurrency Model\n\n### Default: Sequential (`max_concurrency=1`)\n- One prediction at a time\n- Sync `def predict()` supported\n- Cancellation via `SIGUSR1` signal\n\n### Concurrent (`max_concurrency > 1`)\n- Requires `async def predict()`\n- Python 3.11+ for `asyncio.TaskGroup`\n- Configure in `cog.yaml`:\n ```yaml\n concurrency:\n max: 5\n ```\n\n```mermaid\ngantt\n title max_concurrency=1 (Sequential)\n dateFormat X\n axisFormat %s\n section Predictions\n Prediction 1 :0, 3\n Prediction 2 :3, 6\n Prediction 3 :6, 9\n```\n\n```mermaid\ngantt\n title max_concurrency=5 (Concurrent)\n dateFormat X\n axisFormat %s\n section Predictions\n Prediction 1 :0, 4\n Prediction 2 :1, 4\n Prediction 3 :2, 6\n Prediction 4 :0, 5\n Prediction 5 :3, 5\n```\n\n## Environment Variables\n\n| Variable | Default | Purpose |\n|----------|---------|---------|\n| `PORT` | 5000 | HTTP server port |\n| `COG_LOG_LEVEL` | INFO | Logging verbosity |\n| `COG_MAX_CONCURRENCY` | 1 | Max concurrent predictions |\n| `COG_THROTTLE_RESPONSE_INTERVAL` | 0.5s | Webhook rate limit |\n\n## File Locations\n\n| Path | Purpose |\n|------|---------|\n| `/var/run/cog/ready` | K8s readiness probe touch file |\n| `/src` | User code (WORKDIR) |\n| `/src/weights` | Common weights location |\n\n## Code References\n\n| File | Purpose |\n|------|---------|\n| `python/cog/server/http.py` | FastAPI app, endpoints |\n| `python/cog/server/worker.py` | Worker, _ChildWorker |\n| `python/cog/server/runner.py` | PredictionRunner |\n| `python/cog/server/webhook.py` | Webhook delivery |\n| `python/cog/server/stream_redirector.py` | Output capture |\n" }, { "path": "architecture/legacy/README.md", "content": "# Legacy Python Runtime (FastAPI)\n\nThis directory documents the original Cog runtime implementation using Python's FastAPI/Uvicorn HTTP server.\n\n## Status\n\nThis is the **current default** runtime implementation. It uses a two-process architecture with:\n- Parent process: FastAPI/Uvicorn HTTP server\n- Child process: User predictor code in isolated subprocess\n- IPC: Python `multiprocessing.Pipe` with pickled events\n\n## When to Use\n\nThis implementation is used by default when running Cog containers unless the `USE_COGLET` environment variable is set.\n\n## Documentation\n\n- [Prediction API](./03-prediction-api.md) - HTTP endpoints and request/response format\n- [Container Runtime](./04-container-runtime.md) - Two-process architecture and execution flow\n\n## Implementation\n\nPrimary code location: `python/cog/server/`\n- `http.py` - FastAPI application and endpoints\n- `worker.py` - Worker process management\n- `runner.py` - Prediction orchestration\n- `webhook.py` - Webhook delivery\n- `stream_redirector.py` - Output capture\n" }, { "path": "cmd/cog/cog.go", "content": "package main\n\nimport (\n\t\"github.com/replicate/cog/pkg/cli\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc main() {\n\tcmd, err := cli.NewRootCommand()\n\tif err != nil {\n\t\tconsole.Fatalf(\"%f\", err)\n\t}\n\n\tif err = cmd.Execute(); err != nil {\n\t\tconsole.Fatalf(\"%s\", err)\n\t}\n}\n" }, { "path": "crates/.gitignore", "content": "/target/\n" }, { "path": "crates/Cargo.toml", "content": "[workspace]\nresolver = \"2\"\nmembers = [\"coglet\", \"coglet-python\"]\n\n[workspace.package]\nversion = \"0.17.0-rc.2\"\nedition = \"2024\"\nlicense = \"Apache-2.0\"\nrepository = \"https://github.com/replicate/cog\"\nhomepage = \"https://cog.run\"\ndocumentation = \"https://cog.run/docs\"\nkeywords = [\"machine-learning\", \"inference\", \"containers\", \"prediction\"]\ncategories = [\"development-tools\", \"web-programming\"]\n\n[workspace.dependencies]\n# Async runtime\ntokio = { version = \"1\", features = [\"full\"] }\ntokio-util = \"0.7\"\nfutures = \"0.3\"\n\n# HTTP server\naxum = \"0.8\"\n\n# HTTP client\nreqwest = { version = \"0.12\", default-features = false, features = [\"json\", \"rustls-tls-native-roots\"] }\n\n# Serialization\nserde = { version = \"1\", features = [\"derive\"] }\nserde_json = \"1\"\n\n# Identifiers\nuuid = { version = \"1\", features = [\"v4\", \"serde\"] }\n\n# Observability\ntracing = \"0.1\"\ntracing-subscriber = { version = \"0.3\", features = [\"env-filter\", \"json\"] }\n\n# Error handling\nthiserror = \"2\"\nanyhow = \"1\"\n\n# Python bindings\npyo3 = { version = \"0.27\", features = [\"abi3-py310\"] }\npyo3-async-runtimes = { version = \"0.27\", features = [\"tokio-runtime\"] }\npyo3-stub-gen = \"0.18\"\n\n# Testing\ninsta = { version = \"1\", features = [\"json\"] }\n" }, { "path": "crates/README.md", "content": "# Coglet: Rust Runtime for Cog\n\nCoglet is the Rust-based prediction server that powers Cog's subprocess isolation model.\nIt provides process isolation, concurrent slot management, and high-performance IPC for\nrunning ML predictions.\n\n## Architecture Overview\n\n```\n┌─────────────────────────────────────────────────────────────────────────────┐\n│ Parent Process │\n│ ┌─────────────┐ ┌──────────────┐ ┌─────────────────────────────────┐ │\n│ │ HTTP Server │───▶│ Prediction │───▶│ Orchestrator │ │\n│ │ (axum) │ │ Service │ │ - Spawns worker subprocess │ │\n│ └─────────────┘ └──────────────┘ │ - Routes predictions to slots │ │\n│ │ - Handles worker lifecycle │ │\n│ └───────────────┬─────────────────┘ │\n│ │ │\n│ ┌───────────────────────────────┼───────────────┐ │\n│ │ Control Channel (stdin/stdout - JSON lines) │ │\n│ │ - Init, Ready, Cancel, Shutdown │ │\n│ └───────────────────────────────┼───────────────┘ │\n│ │ │\n│ ┌───────────────────────────────┼───────────────┐ │\n│ │ Slot Sockets (Unix domain - per slot) │ │\n│ │ - Predict requests │ │\n│ │ - Streaming logs, outputs │ │\n│ │ - Done/Failed/Cancelled responses │ │\n│ └───────────────────────────────┼───────────────┘ │\n└──────────────────────────────────────────────────────────┼───────────────────┘\n │\n┌──────────────────────────────────────────────────────────┼───────────────────┐\n│ Worker Subprocess │ │\n│ ┌─────────────────────────────────────────────────────────────────────────┐ │\n│ │ Python Runtime (GIL) │ │\n│ │ ┌─────────────────┐ ┌─────────────────┐ ┌───────────────────────┐ │ │\n│ │ │ PythonPredictor │ │ SlotLogWriter │ │ Audit Hook │ │ │\n│ │ │ - load() │ │ (sys.stdout/err)│ │ - Protects streams │ │ │\n│ │ │ - setup() │ │ Routes via │ │ - Tee pattern for │ │ │\n│ │ │ - predict() │ │ ContextVar │ │ user overrides │ │ │\n│ │ └─────────────────┘ └─────────────────┘ └───────────────────────┘ │ │\n│ └─────────────────────────────────────────────────────────────────────────┘ │\n│ │\n│ ┌──────────────────────────────────────────────────────────────────────┐ │\n│ │ Tokio Runtime │ │\n│ │ - Async event loop for slot socket I/O │ │\n│ │ - Releases GIL during I/O (py.detach) │ │\n│ │ - Single async executor for async predictors │ │\n│ └──────────────────────────────────────────────────────────────────────┘ │\n└──────────────────────────────────────────────────────────────────────────────┘\n```\n\n## Prediction Flow\n\n```\nHTTP Request Parent Process Worker Subprocess\n │ │ │\n │ POST /predictions │ │\n ├───────────────────────────────▶│ │\n │ │ │\n │ ┌───────────┴───────────┐ │\n │ │ 1. Acquire slot permit│ │\n │ │ 2. Register prediction│ │\n │ └───────────┬───────────┘ │\n │ │ │\n │ │ SlotRequest::Predict │\n │ │ {id, input} │\n │ ├─────────────────────────────────▶│\n │ │ (slot socket) │\n │ │ │\n │ │ ┌───────────┴───────────┐\n │ │ │ 3. Set ContextVar │\n │ │ │ 4. Call predict() │\n │ │ └───────────┬───────────┘\n │ │ │\n │ │ SlotResponse::Log │\n │ │◀─────────────────────────────────┤ (streaming)\n │ │ │\n │ │ SlotResponse::Output │\n │ │◀─────────────────────────────────┤ (generators)\n │ │ │\n │ │ SlotResponse::Done │\n │ │◀─────────────────────────────────┤\n │ │ {id, output, predict_time} │\n │ │ │\n │ ┌───────────┴───────────┐ │\n │ │ 5. Update prediction │ │\n │ │ 6. Release permit │ │\n │ │ 7. Send webhook │ │\n │ └───────────┬───────────┘ │\n │ │ │\n │ 200 OK {output} │ │\n │◀───────────────────────────────┤ │\n │ │ │\n```\n\n## Startup Sequence\n\n```\n1. coglet.server.serve() called from Python\n │\n ├─▶ Start HTTP server immediately (health returns STARTING until ready)\n │\n └─▶ Spawn orchestrator task\n │\n ├─▶ Create slot transport (Unix sockets)\n │\n ├─▶ Spawn worker: python -c \"import coglet; coglet.server._run_worker()\"\n │\n ├─▶ Send Init message (predictor_ref, num_slots, transport_info)\n │ │\n │ │ ┌────────────────────────────────────────────────┐\n │ └──▶│ Worker: connect sockets, install log writers, │\n │ │ install audit hook, load predictor, run setup │\n │ └────────────────────────────────────────────────┘\n │\n ├─▶ Wait for Ready {slots, schema} or Failed {error}\n │\n ├─▶ Populate PermitPool with slot sockets\n │\n ├─▶ Start event loop (routes responses to predictions)\n │\n └─▶ Set health = READY, start accepting predictions\n```\n\n## Components\n\n### coglet (core library)\nPure Rust library with no Python dependencies. Provides:\n- **orchestrator.rs** - Spawns worker, manages lifecycle, routes messages\n- **worker.rs** - Child-side event loop, prediction execution\n- **service.rs** - Transport-agnostic prediction service\n- **permit/** - Slot-based concurrency control (PermitPool)\n- **bridge/** - IPC protocol and transport (Unix sockets + JSON codec)\n- **transport/http/** - Axum-based HTTP server and routes\n\n### coglet-python (PyO3 bindings)\nBridges coglet to Python via PyO3. Provides:\n- **lib.rs** - Python module with `serve()`, `active()`, `_run_worker()`\n- **predictor.rs** - Wraps Python predictor class (sync/async detection)\n- **worker_bridge.rs** - Implements `PredictHandler` trait for Python\n- **log_writer.rs** - ContextVar-based stdout/stderr routing\n- **audit.rs** - Protects runtime streams from user code\n- **cancel.rs** - SIGUSR1-based cancellation for sync predictors\n\n## Directory Structure\n\n```\ncrates/\n├── Cargo.toml # Workspace manifest\n├── Cargo.lock\n├── deny.toml # cargo-deny configuration\n│\n├── coglet/ # Core Rust library\n│ ├── Cargo.toml\n│ └── src/\n│ ├── lib.rs # Public API exports\n│ ├── health.rs # Health/SetupStatus types\n│ ├── prediction.rs # Prediction state machine\n│ ├── predictor.rs # PredictionResult, PredictionError\n│ ├── service.rs # PredictionService\n│ ├── webhook.rs # WebhookSender (retry, trace context)\n│ ├── version.rs # Version info\n│ ├── webhook.rs # Webhook sender\n│ ├── orchestrator.rs # Worker lifecycle, event loop (parent)\n│ ├── worker.rs # Worker event loop (child)\n│ ├── bridge/\n│ │ ├── mod.rs\n│ │ ├── codec.rs # JSON line codec\n│ │ ├── protocol.rs # Message types (ControlRequest, SlotResponse, etc.)\n│ │ └── transport.rs # Unix socket transport\n│ ├── permit/\n│ │ ├── mod.rs\n│ │ ├── pool.rs # PermitPool (concurrency control)\n│ │ └── slot.rs # PredictionSlot (permit + prediction)\n│ └── transport/\n│ ├── mod.rs\n│ └── http/\n│ ├── mod.rs\n│ ├── server.rs # Axum server setup\n│ └── routes.rs # HTTP handlers\n│\n└── coglet-python/ # PyO3 bindings\n ├── Cargo.toml\n ├── coglet.pyi # Type stubs for Python\n └── src/\n ├── lib.rs # Python module definition\n ├── predictor.rs # PythonPredictor wrapper\n ├── worker_bridge.rs # PredictHandler impl\n ├── input.rs # Input processing (Pydantic/ADT)\n ├── output.rs # Output serialization\n ├── log_writer.rs # SlotLogWriter, ContextVar routing\n ├── audit.rs # Audit hook, TeeWriter\n └── cancel.rs # Cancellation support\n```\n\n## Bridge Protocol\n\nTwo communication channels between parent and worker:\n\n### Control Channel (stdin/stdout)\n\nUsed for lifecycle messages. JSON lines, one message per line.\n\n**Parent → Worker:**\n```json\n{\"type\": \"init\", \"predictor_ref\": \"predict.py:Predictor\", \"num_slots\": 2, ...}\n{\"type\": \"cancel\", \"slot\": \"uuid\"}\n{\"type\": \"shutdown\"}\n```\n\n**Worker → Parent:**\n```json\n{\"type\": \"ready\", \"slots\": [\"uuid1\", \"uuid2\"], \"schema\": {...}}\n{\"type\": \"log\", \"source\": \"stdout\", \"data\": \"Loading model...\"}\n{\"type\": \"idle\", \"slot\": \"uuid\"}\n{\"type\": \"failed\", \"slot\": \"uuid\", \"error\": \"Setup failed: ...\"}\n{\"type\": \"shutting_down\"}\n```\n\n### Slot Sockets (Unix domain)\n\nPer-slot bidirectional sockets for prediction data. Avoids head-of-line blocking.\n\n**Parent → Worker:**\n```json\n{\"type\": \"predict\", \"id\": \"pred_123\", \"input\": {\"prompt\": \"Hello\"}}\n```\n\n**Worker → Parent:**\n```json\n{\"type\": \"log\", \"source\": \"stdout\", \"data\": \"Processing...\"}\n{\"type\": \"output\", \"output\": \"chunk\"}\n{\"type\": \"done\", \"id\": \"pred_123\", \"output\": \"Hello, world!\", \"predict_time\": 0.5}\n{\"type\": \"failed\", \"id\": \"pred_123\", \"error\": \"ValueError: ...\"}\n{\"type\": \"cancelled\", \"id\": \"pred_123\"}\n```\n\n## Key Design Decisions\n\n### Subprocess Isolation\nWorker runs in a separate process. Benefits:\n- Crash isolation (worker crash → restart, parent survives)\n- Memory isolation (GPU memory leaks don't accumulate)\n- Clean shutdown (SIGKILL if needed)\n\n### Single Worker Mode\nAlways exactly one worker subprocess. No dynamic scaling - the parent is\nlightweight, all the heavy lifting happens in the worker.\n\n### Slot-Based Concurrency\nEach slot is a Unix socket pair. `max_concurrency` determines slot count.\nPermits control access - at most one prediction per slot at a time.\n\n### ContextVar-Based Log Routing\nAsync predictions may spawn tasks. ContextVar propagates prediction ID\nthrough the call stack, allowing log routing even from spawned tasks.\n\n### Audit Hook Protection\nUser code might replace `sys.stdout`. The audit hook intercepts this and\nwraps their stream in a TeeWriter, preserving our log routing while\nallowing their code to work as expected.\n" }, { "path": "crates/coglet/Cargo.toml", "content": "[package]\nname = \"coglet\"\ndescription = \"High-performance prediction server for Cog ML models\"\nversion.workspace = true\nedition.workspace = true\nlicense.workspace = true\nrepository.workspace = true\nhomepage.workspace = true\ndocumentation.workspace = true\nkeywords.workspace = true\ncategories.workspace = true\n\n[dependencies]\n# Async runtime\ntokio.workspace = true\ntokio-util = { workspace = true, features = [\"codec\"] }\nfutures.workspace = true\nasync-trait = \"0.1\"\n\n# Serialization\nserde.workspace = true\nserde_json.workspace = true\n\n# Encoding\nbase64 = \"0.22.1\"\nmime_guess = \"2.0.5\"\n\n# Identifiers\nuuid.workspace = true\n\n# HTTP server\naxum.workspace = true\n\n# HTTP client (webhooks)\nreqwest.workspace = true\nureq = { version = \"3\", default-features = false, features = [\"json\", \"rustls\", \"platform-verifier\"] }\n\n# Time\nchrono = { version = \"0.4\", features = [\"serde\"] }\n\n# Error handling\nthiserror.workspace = true\nanyhow.workspace = true\n\n# Input validation\njsonschema = \"0.29\"\n\n# Concurrent collections\ndashmap = \"6\"\n\n# Observability\ntracing.workspace = true\ntracing-subscriber.workspace = true\n\n[target.'cfg(unix)'.dependencies]\nnix = { version = \"0.30\", features = [\"signal\", \"fs\"] }\n\n[dev-dependencies]\ninsta.workspace = true\ntempfile = \"3\"\nwiremock = \"0.6\"\ntower = { version = \"0.5\", features = [\"util\"] }\nhttp-body-util = \"0.1\"\n" }, { "path": "crates/coglet/README.md", "content": "# coglet\n\nCore Rust library for the coglet prediction server. Pure Rust with no Python\ndependencies - the Python bindings live in `coglet-python`.\n\n## Architecture\n\n```\n coglet\n ┌─────────────────────────────────────────────────────────────────┐\n │ │\n │ ┌─────────────────────────────────────────────────────────┐ │\n │ │ transport/http │ │\n │ │ ┌──────────────┐ ┌─────────────────────────────────┐ │ │\n │ │ │ server.rs │ │ routes.rs │ │ │\n │ │ │ Axum setup │ │ /health, /predictions, /cancel │ │ │\n │ │ └──────────────┘ └─────────────────────────────────┘ │ │\n │ └───────────────────────────────┬─────────────────────────┘ │\n │ │ │\n │ ┌───────────────────────────────▼─────────────────────────┐ │\n │ │ service.rs │ │\n │ │ PredictionService: health, permits, state, webhooks │ │\n │ └───────────────────────────────┬─────────────────────────┘ │\n │ │ │\n │ ┌────────────────────────┼────────────────┐ │\n │ │ │ │ │\n │ ▼ ▼ ▼ │\n │ ┌─────────────┐ ┌────────────────────┐ ┌──────────┐ │\n │ │ permit/ │ │ orchestrator.rs │ │webhook.rs│ │\n │ │ PermitPool │ │ Parent-side: │ │ Sender │ │\n │ │ Slot alloc │ │ spawn, route │ │ Retry │ │\n │ └─────────────┘ └─────────┬──────────┘ └──────────┘ │\n │ │ │\n │ ┌────────────────────────────▼────────────────────────────┐ │\n │ │ bridge/ │ │\n │ │ ┌──────────────┐ ┌─────────────┐ ┌────────────────┐ │ │\n │ │ │ protocol.rs │ │ codec.rs │ │ transport.rs │ │ │\n │ │ │ Message types│ │ JSON lines │ │ Unix sockets │ │ │\n │ │ └──────────────┘ └─────────────┘ └────────────────┘ │ │\n │ └─────────────────────────────────────────────────────────┘ │\n │ │\n │ ┌─────────────────────────────────────────────────────────┐ │\n │ │ worker.rs │ │\n │ │ Child-side: PredictHandler trait, run_worker loop │ │\n │ └─────────────────────────────────────────────────────────┘ │\n │ │\n └─────────────────────────────────────────────────────────────────┘\n```\n\n## Directory Structure\n\n```\ncoglet/\n└── src/\n ├── lib.rs # Public API exports\n │\n │ # Core Types\n ├── health.rs # Health, SetupStatus, SetupResult\n ├── prediction.rs # Prediction state machine\n ├── predictor.rs # PredictionResult, PredictionError, PredictionOutput\n ├── version.rs # VersionInfo\n │\n │ # Service Layer\n ├── service.rs # PredictionService - lifecycle, state, webhooks\n ├── webhook.rs # WebhookSender, webhook types\n │\n │ # Orchestrator (Parent Process)\n ├── orchestrator.rs # spawn_worker, OrchestratorHandle, event loop\n │\n │ # Worker (Child Process) \n ├── worker.rs # run_worker, PredictHandler trait, SetupError\n │\n │ # Concurrency Control\n ├── permit/\n │ ├── mod.rs\n │ ├── pool.rs # PermitPool - slot permit management\n │ └── slot.rs # PredictionSlot - permit + prediction binding\n │\n │ # IPC Bridge\n ├── bridge/\n │ ├── mod.rs\n │ ├── protocol.rs # ControlRequest, ControlResponse, SlotRequest, SlotResponse\n │ ├── codec.rs # JsonCodec - newline-delimited JSON\n │ └── transport.rs # Unix socket transport, ChildTransportInfo\n │\n │ # HTTP Transport\n └── transport/\n ├── mod.rs\n └── http/\n ├── mod.rs\n ├── server.rs # ServerConfig, serve()\n └── routes.rs # Route handlers, request/response types\n```\n\n## Key Components\n\n### PredictionService (`service.rs`)\n\nSingle owner of prediction state. Manages:\n- Health state (Unknown → Starting → Ready/SetupFailed)\n- PermitPool + Orchestrator reference\n- Active predictions (DashMap — single source of truth)\n- Cancellation (cancel tokens + orchestrator delegation)\n- Webhooks fire from Prediction mutation methods (no dual state)\n\n```rust\nlet service = PredictionService::new_no_pool()\n .with_health(Health::Starting)\n .with_version(version);\n\n// Later, after worker is ready:\nservice.set_orchestrator(pool, handle).await;\nservice.set_health(Health::Ready).await;\n```\n\n### Orchestrator (`orchestrator.rs`)\n\nParent-side worker lifecycle management.\n\n```\nspawn_worker(config)\n │\n ├─▶ Create Unix socket transport (N slots)\n ├─▶ Spawn: python -c \"import coglet; coglet.server._run_worker()\"\n ├─▶ Send Init message via stdin\n ├─▶ Wait for worker to connect sockets\n ├─▶ Wait for Ready message (with timeout)\n ├─▶ Populate PermitPool with slot writers\n ├─▶ Spawn event loop task\n └─▶ Return OrchestratorReady {pool, schema, handle}\n```\n\nEvent loop handles:\n- `ControlResponse::Idle` - Slot ready for next prediction\n- `ControlResponse::Failed` - Slot poisoned, mark unavailable \n- `SlotResponse::Log/Output/Done/Failed` - Route to prediction\n- Worker crash - Fail all in-flight predictions\n\n### Worker (`worker.rs`)\n\nChild-side event loop. Implements `PredictHandler` trait.\n\n```\nrun_worker(handler, config)\n │\n ├─▶ Connect to slot sockets (from env)\n ├─▶ Setup control channel (stdin/stdout)\n ├─▶ Run handler.setup() with log routing\n ├─▶ Send Ready {slots, schema}\n ├─▶ Enter event loop:\n │ - ControlRequest::Cancel → handler.cancel(slot)\n │ - ControlRequest::Shutdown → exit\n │ - SlotRequest::Predict → spawn prediction task\n └─▶ Exit on shutdown or all slots poisoned\n```\n\n### PermitPool (`permit/pool.rs`)\n\nSlot-based concurrency control.\n\n```rust\nlet pool = PermitPool::new(max_concurrency);\n\n// Add slot with its socket writer\npool.add_permit(slot_id, writer);\n\n// Acquire permit (returns None if at capacity)\nlet permit = pool.try_acquire()?;\n\n// Send prediction request\npermit.send(SlotRequest::Predict { id, input }).await?;\n\n// Return permit when done\ndrop(permit);\n```\n\n### Bridge Protocol (`bridge/protocol.rs`)\n\nMessage types for parent-worker communication.\n\n**Control Channel:**\n- `ControlRequest`: Init, Cancel, Shutdown\n- `ControlResponse`: Ready, Log, Idle, Failed, Cancelled, ShuttingDown\n\n**Slot Channel:**\n- `SlotRequest`: Predict\n- `SlotResponse`: Log, Output, Done, Failed, Cancelled\n\nAll messages are JSON with `{\"type\": \"...\"}` discriminator.\n\n## Behaviors\n\n### Health States\n\n```\nUnknown ──▶ Starting ──┬──▶ Ready ◀──▶ Busy\n │\n └──▶ SetupFailed ──▶ Defunct\n```\n\n- **Unknown**: Initial state, health-check returns status in body\n- **Starting**: Setup in progress\n- **Ready**: Accepting predictions\n- **Busy**: Ready but all slots in use (HTTP 409 on new predictions)\n- **SetupFailed**: setup() raised exception\n- **Defunct**: Unrecoverable error\n\n### Prediction States\n\n```\nStarting ──▶ Processing ──┬──▶ Succeeded\n ├──▶ Failed\n └──▶ Canceled\n```\n\n### Cancellation\n\n1. HTTP DELETE /predictions/{id} or PUT /predictions/{id}/cancel\n2. Parent sends `ControlRequest::Cancel { slot }`\n3. Worker calls `handler.cancel(slot)`\n4. For sync: SIGUSR1 raises KeyboardInterrupt\n5. For async: `future.cancel()` on the asyncio task\n6. Prediction returns with `SlotResponse::Cancelled`\n\n### Shutdown\n\n**Graceful (SIGTERM with await_explicit_shutdown):**\n1. Stop accepting new predictions\n2. Wait for in-flight to complete\n3. Send `ControlRequest::Shutdown`\n4. Worker responds `ShuttingDown`, exits\n5. Parent exits\n\n**Immediate (SIGTERM without flag):**\n1. Send `ControlRequest::Shutdown`\n2. Cancel in-flight predictions\n3. Exit\n\n**Worker crash:**\n1. Control channel closes\n2. Event loop detects, fails all in-flight predictions\n3. Health → Defunct\n\n### Slot Poisoning\n\nIf a slot socket has an error (write fails, etc.), the slot is marked poisoned.\nIt won't receive new predictions. If all slots are poisoned, worker exits.\n\n```rust\nenum SlotOutcome {\n Idle(SlotId), // Ready for next prediction\n Poisoned { slot, error }, // Slot is dead\n}\n```\n" }, { "path": "crates/coglet/src/bridge/codec.rs", "content": "//! Framed codec for worker communication.\n//!\n//! Uses LengthDelimitedCodec for framing + serde_json for serialization.\n//! Works over any AsyncRead/AsyncWrite (pipes, sockets, etc).\n\nuse std::io;\nuse std::marker::PhantomData;\n\nuse serde::{Serialize, de::DeserializeOwned};\nuse tokio_util::bytes::{Bytes, BytesMut};\nuse tokio_util::codec::{Decoder, Encoder, LengthDelimitedCodec};\n\n/// Codec that frames messages with length prefix and serializes with JSON.\n///\n/// Wraps LengthDelimitedCodec and adds serde_json serialization.\npub struct JsonCodec {\n inner: LengthDelimitedCodec,\n _phantom: PhantomData,\n}\n\nimpl Default for JsonCodec {\n fn default() -> Self {\n Self::new()\n }\n}\n\nimpl JsonCodec {\n pub fn new() -> Self {\n Self {\n inner: LengthDelimitedCodec::builder()\n .length_field_length(4)\n .new_codec(),\n _phantom: PhantomData,\n }\n }\n}\n\nimpl Decoder for JsonCodec {\n type Item = T;\n type Error = io::Error;\n\n fn decode(&mut self, src: &mut BytesMut) -> Result, Self::Error> {\n match self.inner.decode(src)? {\n Some(bytes) => {\n let item = serde_json::from_slice(&bytes)\n .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;\n Ok(Some(item))\n }\n None => Ok(None),\n }\n }\n}\n\nimpl Encoder for JsonCodec {\n type Error = io::Error;\n\n fn encode(&mut self, item: T, dst: &mut BytesMut) -> Result<(), Self::Error> {\n let json =\n serde_json::to_vec(&item).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;\n let json_len = json.len();\n // SAFETY: These logs must NOT be shipped over IPC (would create feedback loop).\n // WorkerTracingLayer filters out coglet::bridge::codec target to prevent encoding\n // a WorkerLog message from triggering another log that creates another WorkerLog, etc.\n tracing::trace!(json_size_bytes = json_len, \"Encoding frame\");\n if json_len > 100_000 {\n tracing::info!(\n // This log line should be shipped across the IPC to be emitted, unlike the\n // above trace line. This is a real indicator that we've encoded a large\n // frame and is generally useful.\n target: \"coglet::bridge::codec::large_frame\",\n json_size_bytes = json_len,\n json_size_kb = json_len / 1024,\n \"Large frame being encoded\"\n );\n }\n self.inner.encode(Bytes::from(json), dst)\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use crate::bridge::protocol::{\n ControlRequest, ControlResponse, SlotId, SlotRequest, SlotResponse,\n };\n\n #[test]\n fn codec_roundtrip_control_request() {\n let mut codec = JsonCodec::::new();\n let mut buf = BytesMut::new();\n\n let slot = SlotId::new();\n let req = ControlRequest::Cancel { slot };\n codec.encode(req, &mut buf).unwrap();\n let decoded = codec.decode(&mut buf).unwrap().unwrap();\n\n assert!(matches!(decoded, ControlRequest::Cancel { .. }));\n }\n\n #[test]\n fn codec_roundtrip_control_response() {\n let mut codec = JsonCodec::::new();\n let mut buf = BytesMut::new();\n\n let slots = vec![SlotId::new()];\n let resp = ControlResponse::Ready {\n slots,\n schema: None,\n };\n codec.encode(resp, &mut buf).unwrap();\n let decoded = codec.decode(&mut buf).unwrap().unwrap();\n\n assert!(matches!(decoded, ControlResponse::Ready { .. }));\n }\n\n #[test]\n fn codec_roundtrip_slot_request() {\n let mut codec = JsonCodec::::new();\n let mut buf = BytesMut::new();\n\n let req = SlotRequest::Predict {\n id: \"test\".to_string(),\n input: Some(serde_json::json!({\"x\": 1})),\n input_file: None,\n output_dir: \"/tmp/coglet/predictions/test/outputs\".to_string(),\n context: Default::default(),\n };\n\n codec.encode(req.clone(), &mut buf).unwrap();\n let decoded = codec.decode(&mut buf).unwrap().unwrap();\n\n match (req, decoded) {\n (\n SlotRequest::Predict {\n id: id1,\n input: input1,\n input_file: file1,\n output_dir: dir1,\n ..\n },\n SlotRequest::Predict {\n id: id2,\n input: input2,\n input_file: file2,\n output_dir: dir2,\n ..\n },\n ) => {\n assert_eq!(id1, id2);\n assert_eq!(input1, input2);\n assert_eq!(file1, file2);\n assert_eq!(dir1, dir2);\n }\n }\n }\n\n #[test]\n fn codec_roundtrip_slot_response() {\n let mut codec = JsonCodec::::new();\n let mut buf = BytesMut::new();\n\n let resp = SlotResponse::Done {\n id: \"test\".to_string(),\n output: Some(serde_json::json!(\"result\")),\n predict_time: 1.5,\n is_stream: false,\n };\n codec.encode(resp, &mut buf).unwrap();\n let decoded = codec.decode(&mut buf).unwrap().unwrap();\n\n match decoded {\n SlotResponse::Done {\n id,\n output,\n predict_time,\n is_stream,\n } => {\n assert_eq!(id, \"test\");\n assert_eq!(output, Some(serde_json::json!(\"result\")));\n assert!((predict_time - 1.5).abs() < 0.001);\n assert!(!is_stream);\n }\n _ => panic!(\"wrong variant\"),\n }\n }\n}\n" }, { "path": "crates/coglet/src/bridge/mod.rs", "content": "//! IPC bridge for coglet parent-worker communication.\n//!\n//! This module provides the wire protocol and codec for communication between\n//! the coglet orchestrator (parent) and worker subprocess.\n//!\n//! # Architecture\n//!\n//! - **protocol**: Message types (ControlRequest/Response, SlotRequest/Response)\n//! - **codec**: JSON framing codec for AsyncRead/AsyncWrite\n\npub mod codec;\npub mod protocol;\npub mod transport;\n" }, { "path": "crates/coglet/src/bridge/protocol.rs", "content": "//! Wire protocol types for parent-worker communication.\n//!\n//! Two channels:\n//! - **Control channel** (stdin/stdout): Init, Cancel, Shutdown, Ready, Idle\n//! - **Slot sockets**: Prediction data, streaming logs (per-slot to avoid HOL blocking)\n\nuse std::collections::HashMap;\n\nuse serde::{Deserialize, Serialize};\n\nuse super::transport::ChildTransportInfo;\n\n/// Unique identifier for a prediction slot.\n///\n/// UUID v4 avoids confusion with array indices and prevents accidental reuse.\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]\n#[serde(transparent)]\npub struct SlotId(uuid::Uuid);\n\nimpl SlotId {\n pub fn new() -> Self {\n Self(uuid::Uuid::new_v4())\n }\n\n pub fn as_uuid(&self) -> &uuid::Uuid {\n &self.0\n }\n\n pub fn parse(s: &str) -> Result {\n let uuid = uuid::Uuid::parse_str(s)?;\n Ok(Self(uuid))\n }\n}\n\nimpl Default for SlotId {\n fn default() -> Self {\n Self::new()\n }\n}\n\nimpl std::fmt::Display for SlotId {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n write!(f, \"{}\", self.0)\n }\n}\n\n/// Maximum payload size (input or output) that can be sent inline over the IPC\n/// slot socket. Payloads exceeding this threshold are spilled to disk. The\n/// `LengthDelimitedCodec` default frame limit is 8 MiB, so 6 MiB provides a\n/// 2 MiB safety margin for framing overhead and other message fields.\npub const MAX_INLINE_IPC_SIZE: usize = 1024 * 1024 * 6; // 6MiB\n\nconst MAX_WORKER_LOG_SIZE: usize = 1024 * 1024 * 4; // 4MIB\nconst WORKER_LOG_TRUNCATE_NOTICE: &str = \"[**** LOG LINE TRUNCATED AT 4 MiB ****]\";\n\n/// To ensure no panics happen due to oversized log lines, we truncate at 4 MiB. 1 MiB\n/// let alone 4 MiB log line boarder/exceed usefulness from a readability standpoint.\npub fn truncate_worker_log(mut log_message: String) -> String {\n if log_message.len() > MAX_WORKER_LOG_SIZE {\n let boundary =\n log_message.floor_char_boundary(MAX_WORKER_LOG_SIZE - WORKER_LOG_TRUNCATE_NOTICE.len());\n log_message.truncate(boundary);\n log_message.push_str(WORKER_LOG_TRUNCATE_NOTICE);\n }\n log_message\n}\n\n/// Control messages from parent to worker.\n#[derive(Debug, Clone, Serialize, Deserialize)]\n#[serde(tag = \"type\", rename_all = \"snake_case\")]\npub enum ControlRequest {\n /// Initial configuration sent immediately after spawn (must be first message).\n Init {\n predictor_ref: String,\n num_slots: usize,\n transport_info: ChildTransportInfo,\n is_train: bool,\n is_async: bool,\n },\n\n Cancel {\n slot: SlotId,\n },\n\n /// Request user-defined healthcheck execution.\n Healthcheck {\n id: String,\n },\n\n Shutdown,\n}\n\n/// Control messages from worker to parent.\n#[derive(Debug, Clone, Serialize, Deserialize)]\n#[serde(tag = \"type\", rename_all = \"snake_case\")]\npub enum ControlResponse {\n Ready {\n /// Slot IDs in socket order - parent uses these for all subsequent communication.\n slots: Vec,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n schema: Option,\n },\n\n /// Setup-phase logs (before slots are active).\n Log {\n source: LogSource,\n data: String,\n },\n\n /// Worker tracing log (Rust structured logging).\n WorkerLog {\n target: String,\n level: String,\n message: String,\n },\n\n /// Slot completed and is ready for next prediction.\n Idle {\n slot: SlotId,\n },\n\n Cancelled {\n slot: SlotId,\n },\n\n /// Slot is poisoned and will not accept more predictions.\n Failed {\n slot: SlotId,\n error: String,\n },\n\n /// Worker unrecoverable error - parent should poison all slots and fail all\n /// in-flight predictions. The worker will abort immediately after sending this.\n ///\n /// Reason explains *why* (e.g. \"slots mutex poisoned: cannot guarantee slot isolation\").\n Fatal {\n reason: String,\n },\n\n /// System diagnostic: logs dropped due to backpressure.\n DroppedLogs {\n count: usize,\n interval_millis: u64,\n },\n\n /// Result of user-defined healthcheck execution.\n HealthcheckResult {\n id: String,\n status: HealthcheckStatus,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n error: Option,\n },\n\n ShuttingDown,\n}\n\n/// Status of a user-defined healthcheck.\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]\n#[serde(rename_all = \"snake_case\")]\npub enum HealthcheckStatus {\n /// Healthcheck passed (returned True or no healthcheck defined).\n Healthy,\n /// Healthcheck failed (returned False, raised exception, or timed out).\n Unhealthy,\n}\n\n/// Type-safe slot completion - ensures poisoned slots produce Failed, not Idle.\n#[derive(Debug)]\npub enum SlotOutcome {\n Idle(SlotId),\n Poisoned { slot: SlotId, error: String },\n}\n\nimpl SlotOutcome {\n pub fn idle(slot: SlotId) -> Self {\n Self::Idle(slot)\n }\n\n pub fn poisoned(slot: SlotId, error: impl Into) -> Self {\n Self::Poisoned {\n slot,\n error: error.into(),\n }\n }\n\n pub fn slot_id(&self) -> SlotId {\n match self {\n Self::Idle(slot) => *slot,\n Self::Poisoned { slot, .. } => *slot,\n }\n }\n\n pub fn is_poisoned(&self) -> bool {\n matches!(self, Self::Poisoned { .. })\n }\n\n pub fn into_control_response(self) -> ControlResponse {\n match self {\n Self::Idle(slot) => ControlResponse::Idle { slot },\n Self::Poisoned { slot, error } => ControlResponse::Failed { slot, error },\n }\n }\n}\n\n/// Messages from parent to worker on slot socket.\n#[derive(Debug, Clone, Serialize, Deserialize)]\n#[serde(tag = \"type\", rename_all = \"snake_case\")]\npub enum SlotRequest {\n Predict {\n id: String,\n /// Inline input payload (present when input fits within the IPC frame limit).\n #[serde(skip_serializing_if = \"Option::is_none\")]\n input: Option,\n /// Path to a spill file containing the JSON input (present when input exceeds\n /// `MAX_INLINE_IPC_SIZE`). The worker reads, deserializes, and deletes the file.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n input_file: Option,\n /// Directory for writing file outputs (created by coglet before dispatch).\n /// Not included in API responses — internal transport detail.\n output_dir: String,\n /// Per-prediction context from the request body (`dict[str, str]`).\n /// Made available to predictors via `current_scope().context`.\n #[serde(default)]\n context: HashMap,\n },\n}\n\nimpl SlotRequest {\n /// Returns the prediction ID without consuming the request.\n pub fn prediction_id(&self) -> &str {\n match self {\n SlotRequest::Predict { id, .. } => id,\n }\n }\n\n /// Rehydrate the input from either inline value or spill file.\n ///\n /// Returns `(id, input, output_dir, context)`. If the input was spilled to disk,\n /// reads the file, deserializes, and deletes it.\n pub fn rehydrate_input(\n self,\n ) -> std::io::Result<(String, serde_json::Value, String, HashMap)> {\n match self {\n SlotRequest::Predict {\n id,\n input: Some(value),\n output_dir,\n context,\n ..\n } => Ok((id, value, output_dir, context)),\n SlotRequest::Predict {\n id,\n input: None,\n input_file: Some(path),\n output_dir,\n context,\n } => {\n let bytes = std::fs::read(&path)?;\n // Clean up spill file immediately — bytes are already in memory.\n // Do this before parsing so the file is removed even if JSON is corrupt.\n if let Err(e) = std::fs::remove_file(&path) {\n tracing::warn!(path = %path, error = %e, \"Failed to remove input spill file\");\n }\n let value: serde_json::Value = serde_json::from_slice(&bytes)\n .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;\n Ok((id, value, output_dir, context))\n }\n SlotRequest::Predict { .. } => Err(std::io::Error::new(\n std::io::ErrorKind::InvalidData,\n \"SlotRequest::Predict has neither input nor input_file\",\n )),\n }\n }\n}\n\n#[derive(Debug, Clone, Serialize, Deserialize)]\n#[serde(tag = \"type\", rename_all = \"snake_case\")]\npub enum FileOutputKind {\n /// Output is a file-like return type (e.g. File, Path)\n FileType,\n /// Output exceeds size threshold for bridge codec serialization but is not a file-like return type\n Oversized,\n}\n\n/// Accumulation mode for user metrics.\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]\n#[serde(rename_all = \"snake_case\")]\npub enum MetricMode {\n /// Replace existing value (default).\n Replace,\n /// Add to existing numeric value.\n Increment,\n /// Append to existing array.\n Append,\n}\n\n/// Messages from worker to parent on slot socket.\n#[derive(Debug, Clone, Serialize, Deserialize)]\n#[serde(tag = \"type\", rename_all = \"snake_case\")]\npub enum SlotResponse {\n Log {\n source: LogSource,\n data: String,\n },\n\n /// Output for a file/path-like output return type or an output that exceeds the size threshold\n /// for bridge codec serialization.\n FileOutput {\n filename: String,\n kind: FileOutputKind,\n /// Explicit MIME type from the predictor. Falls back to mime_guess when None.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n mime_type: Option,\n },\n\n /// Streaming output chunk (for generators).\n Output {\n output: serde_json::Value,\n },\n\n /// User-emitted metric from the prediction.\n ///\n /// Metrics are key-value pairs attached to the prediction response.\n /// Supports dot-path keys (e.g., \"timing.preprocess\") that the server\n /// resolves into nested objects. The mode controls how values are merged:\n /// - Replace: overwrite existing value\n /// - Increment: add to existing numeric value\n /// - Append: push onto existing array\n Metric {\n name: String,\n value: serde_json::Value,\n mode: MetricMode,\n },\n\n Done {\n id: String,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n output: Option,\n predict_time: f64,\n /// Predictor signal: true when the output is a list, generator, or\n /// iterator — used as fallback when the schema Output type is `Any`\n /// or unavailable.\n #[serde(default, skip_serializing_if = \"std::ops::Not::not\")]\n is_stream: bool,\n },\n\n Failed {\n id: String,\n error: String,\n },\n\n Cancelled {\n id: String,\n },\n}\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]\n#[serde(rename_all = \"snake_case\")]\npub enum LogSource {\n Stdout,\n Stderr,\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use serde_json::json;\n use std::path::PathBuf;\n\n fn test_slot_id() -> SlotId {\n SlotId(uuid::Uuid::parse_str(\"550e8400-e29b-41d4-a716-446655440000\").unwrap())\n }\n\n #[test]\n fn control_init_serializes() {\n let req = ControlRequest::Init {\n predictor_ref: \"predict.py:Predictor\".to_string(),\n num_slots: 2,\n transport_info: ChildTransportInfo::NamedSockets {\n dir: PathBuf::from(\"/tmp/coglet-123\"),\n num_slots: 2,\n },\n is_train: false,\n is_async: true,\n };\n insta::assert_json_snapshot!(req);\n }\n\n #[test]\n fn control_cancel_serializes() {\n let req = ControlRequest::Cancel {\n slot: test_slot_id(),\n };\n insta::assert_json_snapshot!(req);\n }\n\n #[test]\n fn control_shutdown_serializes() {\n let req = ControlRequest::Shutdown;\n insta::assert_json_snapshot!(req);\n }\n\n #[test]\n fn control_healthcheck_serializes() {\n let req = ControlRequest::Healthcheck {\n id: \"hc_123\".to_string(),\n };\n insta::assert_json_snapshot!(req);\n }\n\n #[test]\n fn control_healthcheck_result_healthy_serializes() {\n let resp = ControlResponse::HealthcheckResult {\n id: \"hc_123\".to_string(),\n status: HealthcheckStatus::Healthy,\n error: None,\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn control_healthcheck_result_unhealthy_serializes() {\n let resp = ControlResponse::HealthcheckResult {\n id: \"hc_123\".to_string(),\n status: HealthcheckStatus::Unhealthy,\n error: Some(\"user healthcheck returned False\".to_string()),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn control_ready_serializes() {\n let resp = ControlResponse::Ready {\n slots: vec![test_slot_id()],\n schema: None,\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn control_ready_with_schema_serializes() {\n let resp = ControlResponse::Ready {\n slots: vec![test_slot_id()],\n schema: Some(json!({\n \"openapi\": \"3.0.2\",\n \"info\": {\"title\": \"Cog\", \"version\": \"0.1.0\"}\n })),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn control_idle_serializes() {\n let resp = ControlResponse::Idle {\n slot: test_slot_id(),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn control_cancelled_serializes() {\n let resp = ControlResponse::Cancelled {\n slot: test_slot_id(),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn control_failed_serializes() {\n let resp = ControlResponse::Failed {\n slot: test_slot_id(),\n error: \"segfault\".to_string(),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_predict_serializes() {\n let req = SlotRequest::Predict {\n id: \"pred_123\".to_string(),\n input: Some(json!({\"text\": \"hello\"})),\n input_file: None,\n output_dir: \"/tmp/coglet/predictions/pred_123/outputs\".to_string(),\n context: Default::default(),\n };\n insta::assert_json_snapshot!(req);\n }\n\n #[test]\n fn slot_predict_file_input_serializes() {\n let req = SlotRequest::Predict {\n id: \"pred_456\".to_string(),\n input: None,\n input_file: Some(\"/tmp/coglet/predictions/pred_456/inputs/spill_abc.json\".to_string()),\n output_dir: \"/tmp/coglet/predictions/pred_456/outputs\".to_string(),\n context: Default::default(),\n };\n insta::assert_json_snapshot!(req);\n }\n\n #[test]\n fn slot_log_serializes() {\n let resp = SlotResponse::Log {\n source: LogSource::Stdout,\n data: \"Processing...\".to_string(),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_output_serializes() {\n let resp = SlotResponse::Output {\n output: json!(\"chunk 1\"),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_done_serializes() {\n let resp = SlotResponse::Done {\n id: \"pred_123\".to_string(),\n output: Some(json!(\"final result\")),\n predict_time: 1.234,\n is_stream: false,\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_failed_serializes() {\n let resp = SlotResponse::Failed {\n id: \"pred_123\".to_string(),\n error: \"ValueError: invalid input\".to_string(),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_cancelled_serializes() {\n let resp = SlotResponse::Cancelled {\n id: \"pred_123\".to_string(),\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_metric_replace_serializes() {\n let resp = SlotResponse::Metric {\n name: \"temperature\".to_string(),\n value: json!(0.7),\n mode: MetricMode::Replace,\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_metric_increment_serializes() {\n let resp = SlotResponse::Metric {\n name: \"token_count\".to_string(),\n value: json!(1),\n mode: MetricMode::Increment,\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_metric_append_serializes() {\n let resp = SlotResponse::Metric {\n name: \"logprobs\".to_string(),\n value: json!(-1.2),\n mode: MetricMode::Append,\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_metric_delete_serializes() {\n let resp = SlotResponse::Metric {\n name: \"unwanted\".to_string(),\n value: json!(null),\n mode: MetricMode::Replace,\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn slot_metric_complex_value_serializes() {\n let resp = SlotResponse::Metric {\n name: \"timing\".to_string(),\n value: json!({\"preprocess\": 0.1, \"inference\": 0.8}),\n mode: MetricMode::Replace,\n };\n insta::assert_json_snapshot!(resp);\n }\n\n #[test]\n fn rehydrate_input_inline() {\n let req = SlotRequest::Predict {\n id: \"p1\".to_string(),\n input: Some(json!({\"text\": \"hello\"})),\n input_file: None,\n output_dir: \"/tmp/out\".to_string(),\n context: Default::default(),\n };\n let (id, input, output_dir, _context) = req.rehydrate_input().unwrap();\n assert_eq!(id, \"p1\");\n assert_eq!(input, json!({\"text\": \"hello\"}));\n assert_eq!(output_dir, \"/tmp/out\");\n }\n\n #[test]\n fn rehydrate_input_from_file() {\n let dir = tempfile::tempdir().unwrap();\n let spill_path = dir.path().join(\"spill_test.json\");\n std::fs::write(&spill_path, r#\"{\"key\":\"value\"}\"#).unwrap();\n\n let req = SlotRequest::Predict {\n id: \"p2\".to_string(),\n input: None,\n input_file: Some(spill_path.to_str().unwrap().to_string()),\n output_dir: \"/tmp/out\".to_string(),\n context: Default::default(),\n };\n let (id, input, output_dir, _context) = req.rehydrate_input().unwrap();\n assert_eq!(id, \"p2\");\n assert_eq!(input, json!({\"key\": \"value\"}));\n assert_eq!(output_dir, \"/tmp/out\");\n // Spill file should be deleted\n assert!(!spill_path.exists());\n }\n\n #[test]\n fn rehydrate_input_neither_errors() {\n let req = SlotRequest::Predict {\n id: \"p3\".to_string(),\n input: None,\n input_file: None,\n output_dir: \"/tmp/out\".to_string(),\n context: Default::default(),\n };\n let err = req.rehydrate_input().unwrap_err();\n assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);\n }\n\n #[test]\n fn rehydrate_input_corrupt_file_errors() {\n let dir = tempfile::tempdir().unwrap();\n let spill_path = dir.path().join(\"corrupt.json\");\n std::fs::write(&spill_path, \"not valid json!!!\").unwrap();\n\n let req = SlotRequest::Predict {\n id: \"p4\".to_string(),\n input: None,\n input_file: Some(spill_path.to_str().unwrap().to_string()),\n output_dir: \"/tmp/out\".to_string(),\n context: Default::default(),\n };\n let err = req.rehydrate_input().unwrap_err();\n assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);\n }\n\n #[test]\n fn truncate_worker_log_truncates_long_messages() {\n let emoji = \"🦀\"; // 4-byte UTF-8 character\n // known size of truncate target, add one more character\n let count = 1024 * 1024 * 1024 * 4 / emoji.len() + 1;\n let message: String = truncate_worker_log(emoji.repeat(count));\n assert!(\n message.ends_with(WORKER_LOG_TRUNCATE_NOTICE),\n \"log message didn't end with {}\",\n WORKER_LOG_TRUNCATE_NOTICE\n );\n }\n\n #[test]\n fn truncate_worker_log_does_not_truncate_short_messages() {\n let emoji = \"🦀\"; // 4-byte UTF-8 character\n // known size of truncate target, add one more character\n let count = 10;\n let message: String = truncate_worker_log(emoji.repeat(count));\n assert!(\n !message.ends_with(WORKER_LOG_TRUNCATE_NOTICE),\n \"short log message was truncated\"\n );\n }\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_cancel_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: req\n---\n{\n \"type\": \"cancel\",\n \"slot\": \"550e8400-e29b-41d4-a716-446655440000\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_cancelled_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"cancelled\",\n \"slot\": \"550e8400-e29b-41d4-a716-446655440000\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_failed_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"failed\",\n \"slot\": \"550e8400-e29b-41d4-a716-446655440000\",\n \"error\": \"segfault\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_healthcheck_result_healthy_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"healthcheck_result\",\n \"id\": \"hc_123\",\n \"status\": \"healthy\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_healthcheck_result_unhealthy_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"healthcheck_result\",\n \"id\": \"hc_123\",\n \"status\": \"unhealthy\",\n \"error\": \"user healthcheck returned False\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_healthcheck_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: req\n---\n{\n \"type\": \"healthcheck\",\n \"id\": \"hc_123\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_idle_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"idle\",\n \"slot\": \"550e8400-e29b-41d4-a716-446655440000\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_init_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: req\n---\n{\n \"type\": \"init\",\n \"predictor_ref\": \"predict.py:Predictor\",\n \"num_slots\": 2,\n \"transport_info\": {\n \"NamedSockets\": {\n \"dir\": \"/tmp/coglet-123\",\n \"num_slots\": 2\n }\n },\n \"is_train\": false,\n \"is_async\": true\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_ready_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"ready\",\n \"slots\": [\n \"550e8400-e29b-41d4-a716-446655440000\"\n ]\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_ready_with_schema_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"ready\",\n \"slots\": [\n \"550e8400-e29b-41d4-a716-446655440000\"\n ],\n \"schema\": {\n \"info\": {\n \"title\": \"Cog\",\n \"version\": \"0.1.0\"\n },\n \"openapi\": \"3.0.2\"\n }\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__control_shutdown_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: req\n---\n{\n \"type\": \"shutdown\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_cancelled_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"cancelled\",\n \"id\": \"pred_123\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_done_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"done\",\n \"id\": \"pred_123\",\n \"output\": \"final result\",\n \"predict_time\": 1.234\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_failed_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"failed\",\n \"id\": \"pred_123\",\n \"error\": \"ValueError: invalid input\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_log_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"log\",\n \"source\": \"stdout\",\n \"data\": \"Processing...\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_metric_append_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"metric\",\n \"name\": \"logprobs\",\n \"value\": -1.2,\n \"mode\": \"append\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_metric_complex_value_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"metric\",\n \"name\": \"timing\",\n \"value\": {\n \"inference\": 0.8,\n \"preprocess\": 0.1\n },\n \"mode\": \"replace\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_metric_delete_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"metric\",\n \"name\": \"unwanted\",\n \"value\": null,\n \"mode\": \"replace\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_metric_increment_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"metric\",\n \"name\": \"token_count\",\n \"value\": 1,\n \"mode\": \"increment\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_metric_replace_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"metric\",\n \"name\": \"temperature\",\n \"value\": 0.7,\n \"mode\": \"replace\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_output_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: resp\n---\n{\n \"type\": \"output\",\n \"output\": \"chunk 1\"\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_predict_file_input_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: req\n---\n{\n \"type\": \"predict\",\n \"id\": \"pred_456\",\n \"input_file\": \"/tmp/coglet/predictions/pred_456/inputs/spill_abc.json\",\n \"output_dir\": \"/tmp/coglet/predictions/pred_456/outputs\",\n \"context\": {}\n}\n" }, { "path": "crates/coglet/src/bridge/snapshots/coglet__bridge__protocol__tests__slot_predict_serializes.snap", "content": "---\nsource: coglet/src/bridge/protocol.rs\nexpression: req\n---\n{\n \"type\": \"predict\",\n \"id\": \"pred_123\",\n \"input\": {\n \"text\": \"hello\"\n },\n \"output_dir\": \"/tmp/coglet/predictions/pred_123/outputs\",\n \"context\": {}\n}\n" }, { "path": "crates/coglet/src/bridge/transport.rs", "content": "//! Slot socket transport for parent-worker IPC.\n//!\n//! Platform-specific implementations:\n//! - **NamedSocketTransport**: Filesystem sockets (macOS, Linux, BSD)\n//! - **AbstractSocketTransport**: Linux abstract namespace (no filesystem, auto-cleanup)\n\nuse std::io;\nuse std::path::PathBuf;\n\nuse serde::{Deserialize, Serialize};\nuse tokio::net::UnixStream;\n\n/// Information passed to child process for connecting to slot sockets.\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub enum ChildTransportInfo {\n NamedSockets {\n dir: PathBuf,\n num_slots: usize,\n },\n #[cfg(target_os = \"linux\")]\n AbstractSockets {\n prefix: String,\n num_slots: usize,\n },\n}\n\n/// Named socket transport using filesystem sockets.\n///\n/// Socket path format: `{temp_dir}/coglet-{pid}/slot-{n}.sock`\npub struct NamedSocketTransport {\n dir: PathBuf,\n sockets: Vec,\n listeners: Vec,\n is_parent: bool,\n}\n\nimpl NamedSocketTransport {\n /// Create transport on parent side, binding listeners for child to connect.\n pub async fn create(num_slots: usize) -> io::Result<(Self, ChildTransportInfo)> {\n use std::os::unix::net::UnixListener as StdUnixListener;\n use tokio::net::UnixListener;\n\n let dir = std::env::temp_dir().join(format!(\"coglet-{}\", std::process::id()));\n std::fs::create_dir_all(&dir)?;\n\n tracing::debug!(transport_type = \"named\", dir = %dir.display(), num_slots, \"Creating slot transport\");\n\n let mut listeners = Vec::with_capacity(num_slots);\n for i in 0..num_slots {\n let path = dir.join(format!(\"slot-{}.sock\", i));\n\n if path.exists() {\n std::fs::remove_file(&path)?;\n }\n\n let std_listener = StdUnixListener::bind(&path)?;\n std_listener.set_nonblocking(true)?;\n let listener = UnixListener::from_std(std_listener)?;\n\n tracing::trace!(slot = i, path = %path.display(), \"Bound socket\");\n listeners.push(listener);\n }\n\n let transport = Self {\n dir: dir.clone(),\n sockets: Vec::with_capacity(num_slots),\n listeners,\n is_parent: true,\n };\n\n let child_info = ChildTransportInfo::NamedSockets {\n dir: dir.clone(),\n num_slots,\n };\n\n Ok((transport, child_info))\n }\n\n /// Accept connections from child on all slots.\n pub async fn accept_connections(&mut self, num_slots: usize) -> io::Result<()> {\n for i in 0..num_slots {\n let listener = &self.listeners[i];\n tracing::trace!(slot = i, \"Waiting for child connection\");\n let (stream, _) = listener.accept().await?;\n self.sockets.push(stream);\n tracing::trace!(slot = i, \"Child connected\");\n }\n\n self.listeners.clear();\n Ok(())\n }\n\n /// Connect from child side.\n pub async fn connect(dir: PathBuf, num_slots: usize) -> io::Result {\n let mut sockets = Vec::with_capacity(num_slots);\n\n for i in 0..num_slots {\n let path = dir.join(format!(\"slot-{}.sock\", i));\n tracing::trace!(slot = i, path = %path.display(), \"Connecting to socket\");\n\n let stream = UnixStream::connect(&path).await?;\n sockets.push(stream);\n\n tracing::trace!(slot = i, \"Connected\");\n }\n\n Ok(Self {\n dir,\n sockets,\n listeners: Vec::new(),\n is_parent: false,\n })\n }\n\n pub fn slot_socket(&mut self, slot: usize) -> Option<&mut UnixStream> {\n self.sockets.get_mut(slot)\n }\n\n /// Returns owned sockets for splitting into read/write halves.\n pub fn drain_sockets(&mut self) -> Vec {\n std::mem::take(&mut self.sockets)\n }\n\n pub fn dir(&self) -> &PathBuf {\n &self.dir\n }\n\n pub fn num_slots(&self) -> usize {\n self.sockets.len()\n }\n\n pub fn cleanup(&mut self) -> io::Result<()> {\n if self.is_parent && self.dir.exists() {\n tracing::debug!(dir = %self.dir.display(), \"Cleaning up socket directory\");\n std::fs::remove_dir_all(&self.dir)?;\n }\n Ok(())\n }\n}\n\nimpl Drop for NamedSocketTransport {\n fn drop(&mut self) {\n if let Err(e) = self.cleanup() {\n tracing::warn!(error = %e, \"Failed to cleanup socket directory\");\n }\n }\n}\n\n/// Abstract namespace socket transport (Linux only).\n///\n/// No filesystem entries, auto-cleanup when all references close.\n#[cfg(target_os = \"linux\")]\npub struct AbstractSocketTransport {\n #[allow(dead_code)] // Kept for debugging/identification\n prefix: String,\n sockets: Vec,\n listeners: Vec,\n}\n\n#[cfg(target_os = \"linux\")]\nimpl AbstractSocketTransport {\n /// Create transport on parent side, binding listeners for child to connect.\n pub async fn create(num_slots: usize) -> io::Result<(Self, ChildTransportInfo)> {\n use std::os::linux::net::SocketAddrExt;\n use std::os::unix::net::{SocketAddr, UnixListener as StdUnixListener};\n use tokio::net::UnixListener;\n\n let prefix = format!(\"coglet-{}\", std::process::id());\n\n tracing::debug!(transport_type = \"abstract\", prefix = %prefix, num_slots, \"Creating slot transport\");\n\n let mut listeners = Vec::with_capacity(num_slots);\n for i in 0..num_slots {\n let name = format!(\"{}-{}\", prefix, i);\n let addr = SocketAddr::from_abstract_name(name.as_bytes())?;\n\n let std_listener = StdUnixListener::bind_addr(&addr)?;\n std_listener.set_nonblocking(true)?;\n let listener = UnixListener::from_std(std_listener)?;\n\n tracing::trace!(slot = i, name = %name, \"Bound abstract socket\");\n listeners.push(listener);\n }\n\n let transport = Self {\n prefix: prefix.clone(),\n sockets: Vec::with_capacity(num_slots),\n listeners,\n };\n\n let child_info = ChildTransportInfo::AbstractSockets { prefix, num_slots };\n\n Ok((transport, child_info))\n }\n\n /// Accept connections from child on all slots.\n pub async fn accept_connections(&mut self, num_slots: usize) -> io::Result<()> {\n for i in 0..num_slots {\n let listener = &self.listeners[i];\n tracing::trace!(slot = i, \"Waiting for child connection\");\n let (stream, _) = listener.accept().await?;\n self.sockets.push(stream);\n tracing::trace!(slot = i, \"Child connected\");\n }\n\n self.listeners.clear();\n Ok(())\n }\n\n /// Connect from child side.\n pub async fn connect(prefix: String, num_slots: usize) -> io::Result {\n use std::os::linux::net::SocketAddrExt;\n use std::os::unix::net::SocketAddr;\n\n let mut sockets = Vec::with_capacity(num_slots);\n\n for i in 0..num_slots {\n let name = format!(\"{}-{}\", prefix, i);\n let addr = SocketAddr::from_abstract_name(name.as_bytes())?;\n\n tracing::trace!(slot = i, name = %name, \"Connecting to abstract socket\");\n\n // tokio doesn't support abstract sockets directly\n let std_stream = std::os::unix::net::UnixStream::connect_addr(&addr)?;\n std_stream.set_nonblocking(true)?;\n let stream = UnixStream::from_std(std_stream)?;\n\n sockets.push(stream);\n tracing::trace!(slot = i, \"Connected\");\n }\n\n Ok(Self {\n prefix,\n sockets,\n listeners: Vec::new(),\n })\n }\n\n pub fn slot_socket(&mut self, slot: usize) -> Option<&mut UnixStream> {\n self.sockets.get_mut(slot)\n }\n\n pub fn drain_sockets(&mut self) -> Vec {\n std::mem::take(&mut self.sockets)\n }\n\n pub fn num_slots(&self) -> usize {\n self.sockets.len()\n }\n}\n\npub enum SlotTransport {\n Named(NamedSocketTransport),\n #[cfg(target_os = \"linux\")]\n Abstract(AbstractSocketTransport),\n}\n\nimpl SlotTransport {\n pub fn slot_socket(&mut self, slot: usize) -> Option<&mut UnixStream> {\n match self {\n Self::Named(t) => t.slot_socket(slot),\n #[cfg(target_os = \"linux\")]\n Self::Abstract(t) => t.slot_socket(slot),\n }\n }\n\n pub fn drain_sockets(&mut self) -> Vec {\n match self {\n Self::Named(t) => t.drain_sockets(),\n #[cfg(target_os = \"linux\")]\n Self::Abstract(t) => t.drain_sockets(),\n }\n }\n\n pub fn num_slots(&self) -> usize {\n match self {\n Self::Named(t) => t.num_slots(),\n #[cfg(target_os = \"linux\")]\n Self::Abstract(t) => t.num_slots(),\n }\n }\n\n pub async fn accept_connections(&mut self, num_slots: usize) -> io::Result<()> {\n match self {\n Self::Named(t) => t.accept_connections(num_slots).await,\n #[cfg(target_os = \"linux\")]\n Self::Abstract(t) => t.accept_connections(num_slots).await,\n }\n }\n}\n\n/// Create transport using platform default (abstract on Linux, named elsewhere).\npub async fn create_transport(num_slots: usize) -> io::Result<(SlotTransport, ChildTransportInfo)> {\n #[cfg(target_os = \"linux\")]\n {\n let (transport, info) = AbstractSocketTransport::create(num_slots).await?;\n Ok((SlotTransport::Abstract(transport), info))\n }\n\n #[cfg(not(target_os = \"linux\"))]\n {\n let (transport, info) = NamedSocketTransport::create(num_slots).await?;\n Ok((SlotTransport::Named(transport), info))\n }\n}\n\npub async fn connect_transport(info: ChildTransportInfo) -> io::Result {\n match info {\n ChildTransportInfo::NamedSockets { dir, num_slots } => {\n let transport = NamedSocketTransport::connect(dir, num_slots).await?;\n Ok(SlotTransport::Named(transport))\n }\n #[cfg(target_os = \"linux\")]\n ChildTransportInfo::AbstractSockets { prefix, num_slots } => {\n let transport = AbstractSocketTransport::connect(prefix, num_slots).await?;\n Ok(SlotTransport::Abstract(transport))\n }\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn child_transport_info_roundtrips() {\n let info = ChildTransportInfo::NamedSockets {\n dir: PathBuf::from(\"/tmp/coglet-123\"),\n num_slots: 3,\n };\n let json = serde_json::to_string(&info).unwrap();\n let parsed: ChildTransportInfo = serde_json::from_str(&json).unwrap();\n\n match parsed {\n ChildTransportInfo::NamedSockets { dir, num_slots } => {\n assert_eq!(dir, PathBuf::from(\"/tmp/coglet-123\"));\n assert_eq!(num_slots, 3);\n }\n #[cfg(target_os = \"linux\")]\n _ => panic!(\"Wrong variant\"),\n }\n }\n\n #[cfg(target_os = \"linux\")]\n #[test]\n fn abstract_socket_info_roundtrips() {\n let info = ChildTransportInfo::AbstractSockets {\n prefix: \"coglet-456\".to_string(),\n num_slots: 2,\n };\n let json = serde_json::to_string(&info).unwrap();\n let parsed: ChildTransportInfo = serde_json::from_str(&json).unwrap();\n\n match parsed {\n ChildTransportInfo::AbstractSockets { prefix, num_slots } => {\n assert_eq!(prefix, \"coglet-456\");\n assert_eq!(num_slots, 2);\n }\n _ => panic!(\"Wrong variant\"),\n }\n }\n}\n" }, { "path": "crates/coglet/src/fd_redirect.rs", "content": "//! File descriptor redirection for subprocess isolation.\n//!\n//! Worker uses fd 1 (stdout) as the control channel to the orchestrator. When Python's\n//! setup() spawns subprocesses (subprocess.Popen), they inherit fd 1 and corrupt the\n//! control channel by writing directly into it.\n//!\n//! We redirect fds early in startup: move control channel to high-numbered fds (99-101),\n//! replace fd 1/2 with capture pipes, spawn threads to route captured output through\n//! the log system.\n//!\n//! CRITICAL: Must be called before FFI initialization (Python, etc.).\n//!\n//! ## Safety contracts\n//!\n//! All `unsafe` blocks in this module rely on these guarantees:\n//! 1. Called early in worker before any predictor/FFI code runs (tokio runtime threads\n//! exist but aren't accessing fds 0/1/2)\n//! 2. Standard fds (0, 1, 2) are guaranteed open by the OS at process startup\n//! 3. High-numbered fds (99-101) won't conflict with application/library usage\n//! 4. Ownership transfer to threads via `from_raw_fd` + `forget` prevents double-close\n//!\n//! Cannot use Miri: This code makes actual syscalls (dup/dup2) which Miri can't execute.\n\n#[cfg(unix)]\nuse std::io;\n#[cfg(unix)]\nuse std::os::fd::{AsRawFd, BorrowedFd, FromRawFd, OwnedFd};\n\n#[cfg(unix)]\nuse nix::unistd::{dup, dup2, pipe};\n#[cfg(unix)]\nuse tokio::sync::mpsc;\n\n#[cfg(unix)]\nuse crate::bridge::protocol::{ControlResponse, LogSource};\n\n/// Chosen to be above the range typically used by libraries (avoiding conflicts with\n/// application fds or library-opened files).\n#[cfg(unix)]\nconst CONTROL_STDIN_FD: i32 = 99;\n#[cfg(unix)]\nconst CONTROL_STDOUT_FD: i32 = 100;\n#[cfg(unix)]\nconst WORKER_STDERR_FD: i32 = 101;\n\n#[cfg(unix)]\npub struct ControlChannelFds {\n pub stdin_fd: OwnedFd,\n pub stdout_fd: OwnedFd,\n}\n\n/// Redirect stdout/stderr for subprocess isolation.\n///\n/// CRITICAL: Must be called before FFI initialization. Child processes spawned after\n/// this will inherit the capture pipes (not the control channel).\n#[cfg(unix)]\npub fn redirect_fds_for_subprocess_isolation(\n setup_log_tx: mpsc::Sender,\n) -> io::Result {\n // Safety: Called early in worker startup before FFI initialization (tokio runtime threads\n // exist but aren't accessing fds 0/1/2). dup/dup2 are atomic. BorrowedFd::borrow_raw is\n // safe because we're borrowing standard fds (0, 1, 2) which are guaranteed to be open.\n\n tracing::debug!(\"Preserving control channel to high fds\");\n\n let control_stdin = unsafe {\n let fd = BorrowedFd::borrow_raw(0);\n dup(fd)\n }\n .map_err(|e| io::Error::other(format!(\"dup(0) failed: {}\", e)))?;\n\n let control_stdout = unsafe {\n let fd = BorrowedFd::borrow_raw(1);\n dup(fd)\n }\n .map_err(|e| io::Error::other(format!(\"dup(1) failed: {}\", e)))?;\n\n let worker_stderr = unsafe {\n let fd = BorrowedFd::borrow_raw(2);\n dup(fd)\n }\n .map_err(|e| io::Error::other(format!(\"dup(2) failed: {}\", e)))?;\n\n tracing::trace!(\n control_stdin = control_stdin.as_raw_fd(),\n control_stdout = control_stdout.as_raw_fd(),\n worker_stderr = worker_stderr.as_raw_fd(),\n \"Duped original fds\"\n );\n\n let mut target_stdin = unsafe { OwnedFd::from_raw_fd(CONTROL_STDIN_FD) };\n dup2(&control_stdin, &mut target_stdin)\n .map_err(|e| io::Error::other(format!(\"dup2 stdin failed: {}\", e)))?;\n std::mem::forget(target_stdin); // Don't close, we'll use it later\n\n let mut target_stdout = unsafe { OwnedFd::from_raw_fd(CONTROL_STDOUT_FD) };\n dup2(&control_stdout, &mut target_stdout)\n .map_err(|e| io::Error::other(format!(\"dup2 stdout failed: {}\", e)))?;\n std::mem::forget(target_stdout); // Don't close, we'll use it later\n\n let mut target_stderr = unsafe { OwnedFd::from_raw_fd(WORKER_STDERR_FD) };\n dup2(&worker_stderr, &mut target_stderr)\n .map_err(|e| io::Error::other(format!(\"dup2 stderr failed: {}\", e)))?;\n std::mem::forget(target_stderr); // Don't close, we'll use it later\n\n tracing::trace!(\n stdin_fd = CONTROL_STDIN_FD,\n stdout_fd = CONTROL_STDOUT_FD,\n stderr_fd = WORKER_STDERR_FD,\n \"Moved control channel to high fds\"\n );\n\n // Temps are now duplicated at high positions, safe to close\n drop(control_stdin);\n drop(control_stdout);\n drop(worker_stderr);\n\n tracing::debug!(\"Creating capture pipes for stdout/stderr\");\n\n let (stdout_read, stdout_write) =\n pipe().map_err(|e| io::Error::other(format!(\"pipe failed: {}\", e)))?;\n let (stderr_read, stderr_write) =\n pipe().map_err(|e| io::Error::other(format!(\"pipe failed: {}\", e)))?;\n\n tracing::trace!(\n stdout_read = stdout_read.as_raw_fd(),\n stdout_write = stdout_write.as_raw_fd(),\n stderr_read = stderr_read.as_raw_fd(),\n stderr_write = stderr_write.as_raw_fd(),\n \"Created capture pipes\"\n );\n\n let mut target_fd1 = unsafe { OwnedFd::from_raw_fd(1) };\n dup2(&stdout_write, &mut target_fd1)\n .map_err(|e| io::Error::other(format!(\"dup2(stdout) failed: {}\", e)))?;\n std::mem::forget(target_fd1); // Don't close fd 1\n\n let mut target_fd2 = unsafe { OwnedFd::from_raw_fd(2) };\n dup2(&stderr_write, &mut target_fd2)\n .map_err(|e| io::Error::other(format!(\"dup2(stderr) failed: {}\", e)))?;\n std::mem::forget(target_fd2); // Don't close fd 2\n\n tracing::trace!(\"Replaced fd 1/2 with capture pipes\");\n\n // Write ends are duped to 1/2, close originals\n drop(stdout_write);\n drop(stderr_write);\n\n tracing::debug!(\"Spawning capture threads\");\n\n // Capture both stdout and stderr from subprocesses. Rust tracing was initialized before\n // redirection, so its output also flows through the stderr pipe. All captured output\n // routes to coglet::user target. Bounded channel (500 messages) provides backpressure\n // if subprocess output exceeds processing rate.\n\n let stdout_tx = setup_log_tx.clone();\n let stdout_read_raw = stdout_read.as_raw_fd();\n std::thread::spawn(move || {\n // NOTE: No tracing in capture threads - would create feedback loop (stderr is captured)\n // Safety: We own stdout_read (moved into this thread)\n let mut file = unsafe { std::fs::File::from_raw_fd(stdout_read_raw) };\n let mut buf = [0u8; 4096];\n\n loop {\n match std::io::Read::read(&mut file, &mut buf) {\n Ok(0) => break,\n Ok(n) => {\n let data = String::from_utf8_lossy(&buf[..n]).to_string();\n if stdout_tx\n .blocking_send(ControlResponse::Log {\n source: LogSource::Stdout,\n data,\n })\n .is_err()\n {\n break;\n }\n }\n Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,\n Err(_) => break,\n }\n }\n });\n std::mem::forget(stdout_read); // Ownership transferred to thread\n\n let stderr_tx = setup_log_tx;\n let stderr_read_raw = stderr_read.as_raw_fd();\n std::thread::spawn(move || {\n // NOTE: No tracing in capture threads - would create feedback loop (stderr is captured)\n // Safety: We own stderr_read (moved into this thread)\n let mut file = unsafe { std::fs::File::from_raw_fd(stderr_read_raw) };\n let mut buf = [0u8; 4096];\n\n loop {\n match std::io::Read::read(&mut file, &mut buf) {\n Ok(0) => break,\n Ok(n) => {\n let data = String::from_utf8_lossy(&buf[..n]).to_string();\n if stderr_tx\n .blocking_send(ControlResponse::Log {\n source: LogSource::Stderr,\n data,\n })\n .is_err()\n {\n break;\n }\n }\n Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,\n Err(_) => break,\n }\n }\n });\n std::mem::forget(stderr_read); // Ownership transferred to thread\n\n // Note: Both stdout and stderr now point to capture pipes. Rust tracing was initialized\n // before fd redirection to write to stderr, so its output will be captured along with\n // subprocess stderr. Both will be routed to coglet::user target. The original stderr\n // is still available at fd 101 but unused after redirection.\n\n tracing::info!(\"File descriptor redirection complete\");\n\n // Safety: We own these fds\n Ok(ControlChannelFds {\n stdin_fd: unsafe { OwnedFd::from_raw_fd(CONTROL_STDIN_FD) },\n stdout_fd: unsafe { OwnedFd::from_raw_fd(CONTROL_STDOUT_FD) },\n })\n}\n\n#[cfg(not(unix))]\npub struct ControlChannelFds {\n pub stdin_fd: std::io::Stdin,\n pub stdout_fd: std::io::Stdout,\n}\n\n#[cfg(not(unix))]\npub fn redirect_fds_for_subprocess_isolation(\n _setup_log_tx: tokio::sync::mpsc::Sender,\n) -> io::Result {\n // No fd redirection on non-Unix - subprocesses will pollute control channel\n Ok(ControlChannelFds {\n stdin_fd: std::io::stdin(),\n stdout_fd: std::io::stdout(),\n })\n}\n" }, { "path": "crates/coglet/src/health.rs", "content": "//! Health status types for coglet runtime.\n\nuse serde::{Deserialize, Serialize};\n\n/// Health status of the coglet runtime.\n#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]\n#[serde(rename_all = \"SCREAMING_SNAKE_CASE\")]\npub enum Health {\n /// Just started, status unknown\n #[default]\n Unknown,\n /// Running setup()\n Starting,\n /// Ready to accept predictions\n Ready,\n /// At capacity (all slots in use)\n Busy,\n /// setup() failed\n SetupFailed,\n /// Unrecoverable error\n Defunct,\n}\n\n/// Response-only health status (includes transient states like UNHEALTHY).\n/// Used in HTTP responses but not stored as internal state.\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]\n#[serde(rename_all = \"SCREAMING_SNAKE_CASE\")]\npub enum HealthResponse {\n Unknown,\n Starting,\n Ready,\n Busy,\n SetupFailed,\n Defunct,\n /// User-defined healthcheck failed (transient - not stored)\n Unhealthy,\n}\n\nimpl From for HealthResponse {\n fn from(health: Health) -> Self {\n match health {\n Health::Unknown => HealthResponse::Unknown,\n Health::Starting => HealthResponse::Starting,\n Health::Ready => HealthResponse::Ready,\n Health::Busy => HealthResponse::Busy,\n Health::SetupFailed => HealthResponse::SetupFailed,\n Health::Defunct => HealthResponse::Defunct,\n }\n }\n}\n\n/// Status of the setup phase.\n#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]\n#[serde(rename_all = \"lowercase\")]\npub enum SetupStatus {\n Starting,\n Succeeded,\n Failed,\n}\n\n/// Result of the setup phase.\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub struct SetupResult {\n /// When setup started (ISO 8601 format).\n pub started_at: String,\n /// When setup completed (ISO 8601 format), if finished.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub completed_at: Option,\n /// Status of setup.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub status: Option,\n /// Captured logs during setup.\n #[serde(default, skip_serializing_if = \"String::is_empty\")]\n pub logs: String,\n}\n\nimpl SetupResult {\n /// Create a new SetupResult with the current time as started_at.\n pub fn starting() -> Self {\n Self {\n started_at: chrono::Utc::now().to_rfc3339(),\n completed_at: None,\n status: Some(SetupStatus::Starting),\n logs: String::new(),\n }\n }\n\n /// Mark setup as succeeded with accumulated logs.\n pub fn succeeded(mut self, logs: String) -> Self {\n self.completed_at = Some(chrono::Utc::now().to_rfc3339());\n self.status = Some(SetupStatus::Succeeded);\n self.logs = logs;\n self\n }\n\n /// Mark setup as failed with error logs.\n pub fn failed(mut self, logs: String) -> Self {\n self.completed_at = Some(chrono::Utc::now().to_rfc3339());\n self.status = Some(SetupStatus::Failed);\n self.logs = logs;\n self\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn health_default_is_unknown() {\n assert_eq!(Health::default(), Health::Unknown);\n }\n\n #[test]\n fn health_serializes_screaming_snake_case() {\n insta::assert_json_snapshot!(\n \"health_all_variants\",\n [\n Health::Unknown,\n Health::Starting,\n Health::Ready,\n Health::Busy,\n Health::SetupFailed,\n Health::Defunct,\n ]\n );\n }\n\n #[test]\n fn health_response_serializes_screaming_snake_case() {\n insta::assert_json_snapshot!(\n \"health_response_all_variants\",\n [\n HealthResponse::Unknown,\n HealthResponse::Starting,\n HealthResponse::Ready,\n HealthResponse::Busy,\n HealthResponse::SetupFailed,\n HealthResponse::Defunct,\n HealthResponse::Unhealthy,\n ]\n );\n }\n\n #[test]\n fn health_deserializes_screaming_snake_case() {\n assert_eq!(\n serde_json::from_str::(\"\\\"READY\\\"\").unwrap(),\n Health::Ready\n );\n assert_eq!(\n serde_json::from_str::(\"\\\"SETUP_FAILED\\\"\").unwrap(),\n Health::SetupFailed\n );\n }\n\n #[test]\n fn setup_status_serializes_lowercase() {\n insta::assert_json_snapshot!(\n \"setup_status_all_variants\",\n [\n SetupStatus::Starting,\n SetupStatus::Succeeded,\n SetupStatus::Failed,\n ]\n );\n }\n\n #[test]\n fn setup_status_deserializes_lowercase() {\n assert_eq!(\n serde_json::from_str::(\"\\\"succeeded\\\"\").unwrap(),\n SetupStatus::Succeeded\n );\n }\n}\n" }, { "path": "crates/coglet/src/input_validation.rs", "content": "//! Input validation against the OpenAPI schema.\n//!\n//! Validates prediction inputs before dispatching to the Python worker,\n//! catching missing required fields and unknown fields early with clear\n//! error messages (matching the format users expect from pydantic).\n\nuse std::collections::HashSet;\n\nuse serde_json::Value;\n\n/// A single validation error for one field.\n#[derive(Debug)]\npub struct ValidationError {\n /// Field name (used as loc[2] in the pydantic-compatible response).\n pub field: String,\n /// Human-readable error message.\n pub msg: String,\n /// Error type string (e.g. \"value_error.missing\").\n pub error_type: String,\n}\n\n/// Compiled input validator built from the OpenAPI schema's Input component.\npub struct InputValidator {\n validator: jsonschema::Validator,\n /// Known property names from the schema.\n properties: HashSet,\n /// Required field names from the schema.\n required: Vec,\n}\n\nimpl InputValidator {\n /// Build a validator from a full OpenAPI schema document.\n ///\n /// Extracts `components.schemas.Input`, injects `additionalProperties: false`\n /// (for pydantic parity), and compiles a JSON Schema validator.\n ///\n /// Returns None if the schema doesn't contain an Input component.\n pub fn from_openapi_schema(schema: &Value) -> Option {\n Self::from_openapi_schema_key(schema, \"Input\")\n }\n\n /// Build a validator from a full OpenAPI schema document using a custom\n /// schema key (e.g. \"TrainingInput\" for train endpoints).\n ///\n /// Returns None if the schema doesn't contain the specified component.\n pub fn from_openapi_schema_key(schema: &Value, key: &str) -> Option {\n let input_schema = schema.get(\"components\")?.get(\"schemas\")?.get(key)?;\n\n let properties: HashSet = input_schema\n .get(\"properties\")\n .and_then(|p| p.as_object())\n .map(|obj| obj.keys().cloned().collect())\n .unwrap_or_default();\n\n let required: Vec = input_schema\n .get(\"required\")\n .and_then(|r| r.as_array())\n .map(|a| {\n a.iter()\n .filter_map(|v| v.as_str().map(String::from))\n .collect()\n })\n .unwrap_or_default();\n\n // Clone and inject additionalProperties: false for pydantic parity\n let mut resolved = input_schema.clone();\n if let Some(obj) = resolved.as_object_mut() {\n obj.insert(\"additionalProperties\".to_string(), Value::Bool(false));\n }\n\n // Inline $ref pointers so the validator can resolve them without\n // the full OpenAPI document context. cog-schema-gen emits $ref for\n // enum choices (e.g. \"#/components/schemas/Color\").\n let all_schemas = schema.get(\"components\").and_then(|c| c.get(\"schemas\"));\n inline_refs(&mut resolved, all_schemas);\n\n let validator = jsonschema::validator_for(&resolved)\n .inspect_err(|e| {\n tracing::warn!(error = %e, \"Failed to compile input schema validator\");\n })\n .ok()?;\n\n Some(Self {\n validator,\n properties,\n required,\n })\n }\n\n pub fn required_count(&self) -> usize {\n self.required.len()\n }\n\n /// Validate an input value against the schema.\n ///\n /// Returns Ok(()) on success, or a list of per-field validation errors\n /// formatted for the pydantic-compatible `detail` response.\n pub fn validate(&self, input: &Value) -> Result<(), Vec> {\n if self.validator.validate(input).is_ok() {\n return Ok(());\n }\n\n let mut errors = Vec::new();\n let mut seen_required = false;\n let mut seen_additional = false;\n\n for error in self.validator.iter_errors(input) {\n let msg = error.to_string();\n\n // \"required\" errors: emit one entry per missing field\n if msg.contains(\"is a required property\") && !seen_required {\n seen_required = true;\n let input_obj = input.as_object();\n for field in &self.required {\n let present = input_obj\n .map(|obj| obj.contains_key(field))\n .unwrap_or(false);\n if !present {\n errors.push(ValidationError {\n field: field.clone(),\n msg: \"Field required\".to_string(),\n error_type: \"value_error.missing\".to_string(),\n });\n }\n }\n continue;\n }\n\n // \"additionalProperties\" errors: emit one entry per unknown field\n if msg.contains(\"Additional properties\") && !seen_additional {\n seen_additional = true;\n if let Some(input_obj) = input.as_object() {\n for key in input_obj.keys() {\n if !self.properties.contains(key) {\n errors.push(ValidationError {\n field: key.clone(),\n msg: format!(\"Unexpected field '{key}'\"),\n error_type: \"value_error.extra\".to_string(),\n });\n }\n }\n }\n continue;\n }\n\n // Skip duplicate required/additional messages\n if seen_required && msg.contains(\"is a required property\") {\n continue;\n }\n if seen_additional && msg.contains(\"Additional properties\") {\n continue;\n }\n\n // Type/constraint errors on specific fields\n let path = error.instance_path.to_string();\n let field = path.trim_start_matches('/');\n let field_name = if field.is_empty() {\n \"__root__\".to_string()\n } else {\n field.to_string()\n };\n errors.push(ValidationError {\n field: field_name,\n msg,\n error_type: \"value_error\".to_string(),\n });\n }\n\n if errors.is_empty() {\n Ok(())\n } else {\n Err(errors)\n }\n }\n}\n\n/// Recursively inline `$ref` pointers in a JSON Schema value.\n///\n/// Resolves `{\"$ref\": \"#/components/schemas/Foo\"}` by looking up `Foo` in the\n/// provided schemas map and replacing the `$ref` object with the referenced\n/// content. This allows the validator to work on an extracted subschema without\n/// needing the full OpenAPI document.\nfn inline_refs(value: &mut Value, all_schemas: Option<&Value>) {\n match value {\n Value::Object(obj) => {\n // If this object is a $ref, resolve it\n if let Some(Value::String(ref_str)) = obj.get(\"$ref\")\n && let Some(resolved) = resolve_ref(ref_str, all_schemas)\n {\n *value = resolved;\n // Recurse into the resolved value (it may contain more $refs)\n inline_refs(value, all_schemas);\n return;\n }\n // Recurse into all values\n for v in obj.values_mut() {\n inline_refs(v, all_schemas);\n }\n }\n Value::Array(arr) => {\n for v in arr.iter_mut() {\n inline_refs(v, all_schemas);\n }\n }\n _ => {}\n }\n}\n\n/// Resolve a `$ref` string like `#/components/schemas/Foo` against the schemas map.\nfn resolve_ref(ref_str: &str, all_schemas: Option<&Value>) -> Option {\n let name = ref_str.strip_prefix(\"#/components/schemas/\")?;\n all_schemas?.get(name).cloned()\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use serde_json::json;\n\n fn make_schema(input_schema: Value) -> Value {\n json!({\n \"components\": {\n \"schemas\": {\n \"Input\": input_schema\n }\n }\n })\n }\n\n #[test]\n fn validates_required_fields() {\n let schema = make_schema(json!({\n \"type\": \"object\",\n \"properties\": {\n \"s\": {\"type\": \"string\", \"title\": \"S\"}\n },\n \"required\": [\"s\"]\n }));\n\n let validator = InputValidator::from_openapi_schema(&schema).unwrap();\n\n // Valid input\n assert!(validator.validate(&json!({\"s\": \"hello\"})).is_ok());\n\n // Missing required field\n let errs = validator.validate(&json!({})).unwrap_err();\n assert_eq!(errs.len(), 1);\n assert_eq!(errs[0].field, \"s\");\n assert_eq!(errs[0].msg, \"Field required\");\n }\n\n #[test]\n fn rejects_additional_properties() {\n let schema = make_schema(json!({\n \"type\": \"object\",\n \"properties\": {\n \"s\": {\"type\": \"string\", \"title\": \"S\"}\n },\n \"required\": [\"s\"]\n }));\n\n let validator = InputValidator::from_openapi_schema(&schema).unwrap();\n\n // Extra field should fail\n let errs = validator\n .validate(&json!({\"s\": \"hello\", \"extra\": \"bad\"}))\n .unwrap_err();\n assert_eq!(errs.len(), 1);\n assert_eq!(errs[0].field, \"extra\");\n assert!(errs[0].msg.contains(\"Unexpected\"));\n }\n\n #[test]\n fn missing_and_extra_fields() {\n let schema = make_schema(json!({\n \"type\": \"object\",\n \"properties\": {\n \"s\": {\"type\": \"string\", \"title\": \"S\"}\n },\n \"required\": [\"s\"]\n }));\n\n let validator = InputValidator::from_openapi_schema(&schema).unwrap();\n\n // wrong=value with missing s\n let errs = validator.validate(&json!({\"wrong\": \"value\"})).unwrap_err();\n assert!(errs.len() >= 2);\n let fields: Vec<&str> = errs.iter().map(|e| e.field.as_str()).collect();\n assert!(fields.contains(&\"s\"), \"should report missing s: {fields:?}\");\n assert!(\n fields.contains(&\"wrong\"),\n \"should report extra wrong: {fields:?}\"\n );\n }\n\n #[test]\n fn validates_types() {\n let schema = make_schema(json!({\n \"type\": \"object\",\n \"properties\": {\n \"count\": {\"type\": \"integer\", \"title\": \"Count\"}\n },\n \"required\": [\"count\"]\n }));\n\n let validator = InputValidator::from_openapi_schema(&schema).unwrap();\n\n assert!(validator.validate(&json!({\"count\": 5})).is_ok());\n\n let errs = validator\n .validate(&json!({\"count\": \"not_a_number\"}))\n .unwrap_err();\n assert_eq!(errs[0].field, \"count\");\n }\n\n #[test]\n fn no_schema_returns_none() {\n let schema = json!({\"components\": {\"schemas\": {}}});\n assert!(InputValidator::from_openapi_schema(&schema).is_none());\n }\n\n #[test]\n fn resolves_ref_for_choices() {\n let schema = json!({\n \"components\": {\n \"schemas\": {\n \"Input\": {\n \"type\": \"object\",\n \"properties\": {\n \"color\": {\n \"allOf\": [{\"$ref\": \"#/components/schemas/Color\"}],\n \"x-order\": 0\n }\n },\n \"required\": [\"color\"]\n },\n \"Color\": {\n \"title\": \"Color\",\n \"description\": \"An enumeration.\",\n \"enum\": [\"red\", \"green\", \"blue\"],\n \"type\": \"string\"\n }\n }\n }\n });\n\n let validator = InputValidator::from_openapi_schema(&schema);\n assert!(validator.is_some(), \"validator should compile with $ref\");\n\n let validator = validator.unwrap();\n assert!(validator.validate(&json!({\"color\": \"red\"})).is_ok());\n assert!(validator.validate(&json!({\"color\": \"purple\"})).is_err());\n }\n\n #[test]\n fn optional_fields_work() {\n let schema = make_schema(json!({\n \"type\": \"object\",\n \"properties\": {\n \"s\": {\"type\": \"string\"},\n \"n\": {\"type\": \"integer\"}\n },\n \"required\": [\"s\"]\n }));\n\n let validator = InputValidator::from_openapi_schema(&schema).unwrap();\n\n assert!(validator.validate(&json!({\"s\": \"hello\"})).is_ok());\n assert!(validator.validate(&json!({\"s\": \"hello\", \"n\": 42})).is_ok());\n }\n}\n" }, { "path": "crates/coglet/src/lib.rs", "content": "//! coglet: Rust execution engine for cog models.\n\nmod health;\npub mod input_validation;\nmod prediction;\nmod predictor;\nmod version;\n\npub mod bridge;\nmod fd_redirect;\npub mod orchestrator;\npub mod permit;\npub mod service;\nmod setup_log_accumulator;\npub mod transport;\npub mod webhook;\npub mod worker;\nmod worker_tracing_layer;\n\npub use orchestrator::Orchestrator;\n\npub use service::{PredictionHandle, SyncPredictionGuard};\n\npub use health::{Health, SetupResult, SetupStatus};\npub use input_validation::InputValidator;\npub use prediction::{CancellationToken, Prediction, PredictionOutput, PredictionStatus};\npub use predictor::{PredictionError, PredictionGuard, PredictionMetrics, PredictionResult};\npub use service::{CreatePredictionError, HealthSnapshot, PredictionService};\npub use setup_log_accumulator::{SetupLogAccumulator, drain_accumulated_logs};\npub use version::{COGLET_VERSION, VersionInfo};\npub use worker::{\n PredictHandler, PredictResult, SetupError, SetupLogHook, SlotSender, WorkerConfig, run_worker,\n};\n" }, { "path": "crates/coglet/src/orchestrator.rs", "content": "//! Orchestrator - manages worker subprocess lifecycle and event loop.\n//!\n//! Flow:\n//! 1. Spawn worker subprocess\n//! 2. Send Init message, wait for Ready\n//! 3. Populate PermitPool with slot sockets\n//! 4. Run event loop routing responses to predictions\n//! 5. On worker crash: fail all predictions, shut down\n\nuse std::collections::HashMap;\nuse std::process::Stdio;\nuse std::sync::Arc;\nuse std::sync::Mutex as StdMutex;\nuse std::time::Duration;\n\nuse async_trait::async_trait;\nuse futures::{SinkExt, StreamExt};\nuse tokio::process::{Child, Command};\nuse tokio::sync::mpsc;\nuse tokio_util::codec::{FramedRead, FramedWrite};\n\nuse crate::PredictionOutput;\nuse crate::bridge::codec::JsonCodec;\nuse crate::bridge::protocol::{\n ControlRequest, ControlResponse, FileOutputKind, HealthcheckStatus, SlotId, SlotRequest,\n SlotResponse,\n};\nuse crate::bridge::transport::create_transport;\nuse crate::permit::{InactiveSlotIdleToken, PermitPool, SlotIdleToken};\nuse crate::prediction::Prediction;\n\n/// Upload a file to a signed endpoint, returning the final URL.\n///\n/// Matches the behavior of Python cog's `put_file_to_signed_endpoint`:\n/// PUT to `{endpoint}{filename}` with Content-Type header, then extract\n/// the final URL from the Location header (falling back to response URL),\n/// stripping query parameters. Follows redirects automatically.\nasync fn upload_file(\n endpoint: &str,\n filename: &str,\n data: &[u8],\n content_type: &str,\n) -> Result {\n let url = format!(\"{endpoint}{filename}\");\n let client = reqwest::Client::new();\n let resp = client\n .put(&url)\n .header(\"Content-Type\", content_type)\n .body(data.to_vec())\n .timeout(std::time::Duration::from_secs(25))\n .send()\n .await\n .map_err(|e| format!(\"upload request failed: {e}\"))?;\n\n if !resp.status().is_success() {\n return Err(format!(\"upload returned status {}\", resp.status()));\n }\n\n // Prefer Location header, fall back to final request URL (after redirects)\n let final_url = resp\n .headers()\n .get(\"location\")\n .and_then(|v| v.to_str().ok())\n .map(|s| s.to_string())\n .unwrap_or_else(|| resp.url().to_string());\n\n // Strip query parameters (signing gubbins)\n match reqwest::Url::parse(&final_url) {\n Ok(mut parsed) => {\n parsed.set_query(None);\n Ok(parsed.to_string())\n }\n Err(_) => Ok(final_url),\n }\n}\n\nfn ensure_trailing_slash(s: &str) -> String {\n if s.ends_with('/') {\n s.to_string()\n } else {\n format!(\"{s}/\")\n }\n}\n\n/// Try to lock a prediction mutex.\n/// On poison: logs error, recovers to fail the prediction, returns None.\n/// Caller should remove the prediction from tracking if None is returned.\nfn try_lock_prediction(\n pred: &Arc>,\n) -> Option> {\n match pred.lock() {\n Ok(guard) => Some(guard),\n Err(poisoned) => {\n tracing::error!(\"Prediction mutex poisoned - failing prediction\");\n let mut guard = poisoned.into_inner();\n if !guard.is_terminal() {\n guard.set_failed(\"Internal error: mutex poisoned\".to_string());\n }\n None\n }\n }\n}\n\n/// Wrap collected output items into the correct `PredictionOutput` variant.\n///\n/// Priority:\n/// 1. Schema says `\"type\": \"array\"` (`output_is_array = true`) → always `Stream`\n/// 2. Predictor signals `is_stream` (list/generator return) → always `Stream`\n/// 3. Otherwise → `Single` for one item, `Stream` for multiple\n///\n/// This ensures `List[Path]` with a single element returns `[\"url\"]` not `\"url\"`.\nfn wrap_outputs(\n outputs: Vec,\n output_is_array: bool,\n is_stream: bool,\n) -> PredictionOutput {\n let should_stream = output_is_array || is_stream;\n\n match outputs.as_slice() {\n [] => {\n if should_stream {\n PredictionOutput::Stream(vec![])\n } else {\n PredictionOutput::Single(serde_json::Value::Null)\n }\n }\n _ if should_stream => PredictionOutput::Stream(outputs),\n [single] => PredictionOutput::Single(single.clone()),\n _ => PredictionOutput::Stream(outputs),\n }\n}\n\nfn emit_worker_log(target: &str, level: &str, msg: &str) {\n use std::collections::HashMap;\n use std::sync::OnceLock;\n use tracing::{\n Level, Metadata,\n callsite::{Callsite, Identifier},\n field::FieldSet,\n };\n\n struct DummyCallsite;\n impl Callsite for DummyCallsite {\n fn set_interest(&self, _: tracing::subscriber::Interest) {}\n fn metadata(&self) -> &Metadata<'static> {\n unreachable!()\n }\n }\n\n static DUMMY: DummyCallsite = DummyCallsite;\n static CALLSITES: OnceLock<\n std::sync::Mutex>>,\n > = OnceLock::new();\n static FIELDS: &[&str] = &[\"message\"];\n\n let lvl = match level {\n \"error\" => Level::ERROR,\n \"warn\" => Level::WARN,\n \"info\" => Level::INFO,\n \"debug\" => Level::DEBUG,\n \"trace\" => Level::TRACE,\n _ => Level::INFO,\n };\n\n let target_static: &'static str = Box::leak(target.to_string().into_boxed_str());\n\n let callsites = CALLSITES.get_or_init(|| std::sync::Mutex::new(HashMap::new()));\n let mut map = match callsites.lock() {\n Ok(guard) => guard,\n Err(_poisoned) => {\n tracing::error!(\"Worker log callsite cache poisoned\");\n return;\n }\n };\n\n let meta = map.entry((target_static, lvl)).or_insert_with(|| {\n Metadata::new(\n \"worker_log\",\n target_static,\n lvl,\n Some(file!()),\n Some(line!()),\n Some(module_path!()),\n FieldSet::new(FIELDS, Identifier(&DUMMY)),\n tracing::metadata::Kind::EVENT,\n )\n });\n\n let meta_ref = meta as *const Metadata<'static>;\n drop(map);\n\n let meta = unsafe { &*meta_ref };\n\n tracing::dispatcher::get_default(|dispatch| {\n if dispatch.enabled(meta) {\n let fields = meta.fields();\n if let Some(field) = fields.field(\"message\") {\n let value_array = [(&field, Some(&msg as &dyn tracing::Value))];\n let values = fields.value_set(&value_array);\n dispatch.event(&tracing::Event::new(meta, &values));\n }\n }\n });\n}\n\n/// Result of a user-defined healthcheck.\n#[derive(Debug, Clone)]\npub struct HealthcheckResult {\n pub status: HealthcheckStatus,\n pub error: Option,\n}\n\nimpl HealthcheckResult {\n pub fn healthy() -> Self {\n Self {\n status: HealthcheckStatus::Healthy,\n error: None,\n }\n }\n\n pub fn unhealthy(error: impl Into) -> Self {\n Self {\n status: HealthcheckStatus::Unhealthy,\n error: Some(error.into()),\n }\n }\n\n pub fn is_healthy(&self) -> bool {\n self.status == HealthcheckStatus::Healthy\n }\n}\n\n/// Trait for prediction registration with the orchestrator.\n///\n/// This abstraction enables testing the service layer without a real worker subprocess.\n/// The service only needs to register predictions for response routing - all other\n/// orchestrator operations happen outside the predict path.\n#[async_trait]\npub trait Orchestrator: Send + Sync {\n /// Register a prediction for response routing in the event loop.\n async fn register_prediction(\n &self,\n slot_id: SlotId,\n prediction: Arc>,\n idle_sender: tokio::sync::oneshot::Sender,\n );\n\n /// Cancel a prediction by its prediction ID.\n ///\n /// The orchestrator resolves the prediction ID to a slot ID and sends\n /// a cancel request to the worker over the control socket.\n async fn cancel_by_prediction_id(&self, prediction_id: &str) -> Result<(), OrchestratorError>;\n\n /// Run user-defined healthcheck if available.\n async fn healthcheck(&self) -> Result;\n\n /// Shutdown the orchestrator and worker gracefully.\n async fn shutdown(&self) -> Result<(), OrchestratorError>;\n}\n\n#[derive(Debug, Clone)]\npub struct WorkerSpawnConfig {\n pub num_slots: usize,\n}\n\n#[derive(Debug, thiserror::Error)]\npub enum SpawnError {\n #[error(\"failed to spawn process: {0}\")]\n Spawn(#[from] std::io::Error),\n #[error(\"spawn failed: {0}\")]\n Other(String),\n}\n\n/// Extension point for different worker spawn strategies.\npub trait WorkerSpawner: Send + Sync {\n fn spawn(&self, config: &WorkerSpawnConfig) -> Result;\n}\n\n/// Simple spawner using Python subprocess.\npub struct SimpleSpawner;\n\nimpl WorkerSpawner for SimpleSpawner {\n fn spawn(&self, _config: &WorkerSpawnConfig) -> Result {\n let child = Command::new(\"python\")\n .args([\"-c\", \"import coglet; coglet.server._run_worker()\"])\n .stdin(Stdio::piped())\n .stdout(Stdio::piped())\n .stderr(Stdio::inherit())\n .spawn()?;\n Ok(child)\n }\n}\n\npub struct OrchestratorConfig {\n pub predictor_ref: String,\n pub num_slots: usize,\n pub is_train: bool,\n pub is_async: bool,\n pub setup_timeout: Option,\n pub spawner: Arc,\n /// Upload URL prefix for file outputs (from --upload-url CLI arg).\n pub upload_url: Option,\n}\n\nimpl OrchestratorConfig {\n pub fn new(predictor_ref: impl Into) -> Self {\n Self {\n predictor_ref: predictor_ref.into(),\n num_slots: 1,\n is_train: false,\n is_async: false,\n setup_timeout: None,\n spawner: Arc::new(SimpleSpawner),\n upload_url: None,\n }\n }\n\n pub fn with_upload_url(mut self, upload_url: Option) -> Self {\n self.upload_url = upload_url;\n self\n }\n\n pub fn with_num_slots(mut self, n: usize) -> Self {\n self.num_slots = n;\n self\n }\n\n pub fn with_train(mut self, is_train: bool) -> Self {\n self.is_train = is_train;\n self\n }\n\n pub fn with_async(mut self, is_async: bool) -> Self {\n self.is_async = is_async;\n self\n }\n\n pub fn with_setup_timeout(mut self, timeout: Option) -> Self {\n self.setup_timeout = timeout;\n self\n }\n\n pub fn with_spawner(mut self, spawner: Arc) -> Self {\n self.spawner = spawner;\n self\n }\n}\n\npub struct OrchestratorReady {\n pub pool: Arc,\n pub schema: Option,\n pub handle: OrchestratorHandle,\n pub setup_logs: String,\n}\n\npub struct OrchestratorHandle {\n child: Child,\n ctrl_writer:\n Arc>>>,\n register_tx: mpsc::Sender<(\n SlotId,\n Arc>,\n tokio::sync::oneshot::Sender,\n )>,\n healthcheck_tx: mpsc::Sender>,\n cancel_tx: mpsc::Sender,\n slot_ids: Vec,\n}\n\n#[async_trait]\nimpl Orchestrator for OrchestratorHandle {\n async fn register_prediction(\n &self,\n slot_id: SlotId,\n prediction: Arc>,\n idle_sender: tokio::sync::oneshot::Sender,\n ) {\n let _ = self\n .register_tx\n .send((slot_id, prediction, idle_sender))\n .await;\n }\n\n async fn cancel_by_prediction_id(&self, prediction_id: &str) -> Result<(), OrchestratorError> {\n self.cancel_tx\n .send(prediction_id.to_string())\n .await\n .map_err(|_| OrchestratorError::Protocol(\"cancel channel closed\".to_string()))\n }\n\n async fn healthcheck(&self) -> Result {\n tracing::trace!(\"Healthcheck requested via orchestrator handle\");\n let (response_tx, response_rx) = tokio::sync::oneshot::channel();\n\n // Send our channel to the event loop. If a healthcheck is already\n // in-flight, the event loop coalesces — we get the same result as\n // all other waiters when it comes back.\n self.healthcheck_tx\n .send(response_tx)\n .await\n .map_err(|_| OrchestratorError::Protocol(\"healthcheck channel closed\".to_string()))?;\n\n // Wait for the response with a timeout (worker has 5s, we give 10s total).\n // If we time out, the healthcheck keeps running — our sender just gets a\n // silent failure when the event loop eventually broadcasts.\n match tokio::time::timeout(Duration::from_secs(10), response_rx).await {\n Ok(Ok(result)) => {\n tracing::trace!(healthy = result.is_healthy(), \"Healthcheck completed\");\n Ok(result)\n }\n Ok(Err(_)) => {\n tracing::debug!(\"Healthcheck response channel dropped\");\n Err(OrchestratorError::Protocol(\n \"healthcheck response channel dropped\".to_string(),\n ))\n }\n Err(_) => {\n tracing::debug!(\"Healthcheck timed out after 10s\");\n Ok(HealthcheckResult::unhealthy(\"healthcheck timed out\"))\n }\n }\n }\n\n async fn shutdown(&self) -> Result<(), OrchestratorError> {\n let mut writer = self.ctrl_writer.lock().await;\n writer\n .send(ControlRequest::Shutdown)\n .await\n .map_err(|e| OrchestratorError::Protocol(format!(\"failed to send shutdown: {}\", e)))\n }\n}\n\nimpl OrchestratorHandle {\n pub async fn cancel(&self, slot_id: SlotId) -> Result<(), OrchestratorError> {\n let mut writer = self.ctrl_writer.lock().await;\n writer\n .send(ControlRequest::Cancel { slot: slot_id })\n .await\n .map_err(|e| OrchestratorError::Protocol(format!(\"failed to send cancel: {}\", e)))\n }\n\n pub fn slot_ids(&self) -> &[SlotId] {\n &self.slot_ids\n }\n\n pub async fn wait(&mut self) -> Result<(), OrchestratorError> {\n self.child.wait().await.map_err(|e| {\n OrchestratorError::Protocol(format!(\"failed to wait for worker: {}\", e))\n })?;\n Ok(())\n }\n}\n\n#[derive(Debug, thiserror::Error)]\npub enum OrchestratorError {\n #[error(\"failed to spawn worker: {0}\")]\n Spawn(String),\n #[error(\"worker setup failed: {0}\")]\n Setup(String),\n #[error(\"worker setup timed out\")]\n SetupTimeout,\n #[error(\"protocol error: {0}\")]\n Protocol(String),\n #[error(\"worker crashed\")]\n WorkerCrashed,\n}\n\npub async fn spawn_worker(\n config: OrchestratorConfig,\n setup_log_rx: &mut tokio::sync::mpsc::UnboundedReceiver,\n) -> Result {\n let num_slots = config.num_slots;\n\n tracing::info!(num_slots, \"Creating slot transport\");\n let (mut transport, child_transport_info) = create_transport(num_slots)\n .await\n .map_err(|e| OrchestratorError::Spawn(format!(\"failed to create transport: {}\", e)))?;\n\n tracing::info!(\"Spawning worker subprocess\");\n\n let spawn_config = WorkerSpawnConfig { num_slots };\n let mut child = config\n .spawner\n .spawn(&spawn_config)\n .map_err(|e| OrchestratorError::Spawn(format!(\"spawner failed: {}\", e)))?;\n\n let stdin = child\n .stdin\n .take()\n .ok_or_else(|| OrchestratorError::Spawn(\"stdin not captured\".to_string()))?;\n let stdout = child\n .stdout\n .take()\n .ok_or_else(|| OrchestratorError::Spawn(\"stdout not captured\".to_string()))?;\n\n let mut ctrl_writer = FramedWrite::new(stdin, JsonCodec::::new());\n let mut ctrl_reader = FramedRead::new(stdout, JsonCodec::::new());\n\n tracing::debug!(\"Sending Init to worker\");\n ctrl_writer\n .send(ControlRequest::Init {\n predictor_ref: config.predictor_ref.clone(),\n num_slots,\n transport_info: child_transport_info,\n is_train: config.is_train,\n is_async: config.is_async,\n })\n .await\n .map_err(|e| OrchestratorError::Protocol(format!(\"failed to send Init: {}\", e)))?;\n\n tracing::debug!(\"Waiting for worker to connect to slot sockets\");\n transport\n .accept_connections(num_slots)\n .await\n .map_err(|e| OrchestratorError::Spawn(format!(\"failed to accept connections: {}\", e)))?;\n\n tracing::debug!(\"Waiting for Ready from worker\");\n let setup_fut = async {\n loop {\n match ctrl_reader.next().await {\n Some(Ok(ControlResponse::Ready { slots, schema })) => {\n return Ok((slots, schema));\n }\n Some(Ok(ControlResponse::Log { source, data })) => {\n for line in data.lines() {\n tracing::info!(target: \"coglet::setup\", source = ?source, \"{}\", line);\n }\n }\n Some(Ok(ControlResponse::WorkerLog {\n target,\n level,\n message,\n })) => {\n emit_worker_log(&target, &level, &message);\n }\n Some(Ok(ControlResponse::DroppedLogs {\n count,\n interval_millis,\n })) => {\n tracing::trace!(count, interval_millis, \"Received DroppedLogs during setup\");\n let interval_secs = interval_millis as f64 / 1000.0;\n tracing::warn!(\n \"Log production exceeds consumption rate during setup. {} logs dropped in last {:.1}s\",\n count,\n interval_secs\n );\n }\n Some(Ok(ControlResponse::Failed { slot, error })) => {\n return Err(OrchestratorError::Setup(format!(\n \"worker setup failed (slot {}): {}\",\n slot, error\n )));\n }\n Some(Ok(ControlResponse::Fatal { reason })) => {\n return Err(OrchestratorError::Setup(format!(\n \"worker fatal: {}\",\n reason\n )));\n }\n Some(Ok(other)) => {\n tracing::warn!(?other, \"Unexpected message during setup\");\n }\n Some(Err(e)) => {\n return Err(OrchestratorError::Protocol(format!(\n \"control channel error: {}\",\n e\n )));\n }\n None => {\n return Err(OrchestratorError::WorkerCrashed);\n }\n }\n }\n };\n\n let (slot_ids, schema) = match config.setup_timeout {\n Some(timeout) => {\n tracing::debug!(\n timeout_secs = timeout.as_secs(),\n \"Waiting for setup with timeout\"\n );\n match tokio::time::timeout(timeout, setup_fut).await {\n Ok(Ok((slots, schema))) => {\n tracing::debug!(num_slots = slots.len(), \"Setup completed within timeout\");\n (slots, schema)\n }\n Ok(Err(e)) => {\n tracing::debug!(error = %e, \"Setup failed\");\n return Err(e);\n }\n Err(_) => {\n tracing::debug!(timeout_secs = timeout.as_secs(), \"Setup timed out\");\n return Err(OrchestratorError::SetupTimeout);\n }\n }\n }\n None => {\n tracing::debug!(\"Waiting for setup with no timeout\");\n setup_fut.await?\n }\n };\n\n let setup_logs = crate::setup_log_accumulator::drain_accumulated_logs(setup_log_rx);\n tracing::debug!(\n setup_logs_len = setup_logs.len(),\n \"Drained accumulated setup logs\"\n );\n\n tracing::debug!(num_slots = slot_ids.len(), \"Worker ready\");\n\n if let Some(ref s) = schema\n && let Ok(json) = serde_json::to_string_pretty(s)\n {\n tracing::trace!(target: \"coglet::schema\", schema = %json, \"OpenAPI schema\");\n }\n\n // Determine whether the output type is an array from the schema so the\n // event loop can correctly wrap single-element list returns as Stream\n // instead of collapsing them to Single.\n let output_is_array = schema\n .as_ref()\n .and_then(|s| s.get(\"components\"))\n .and_then(|c| c.get(\"schemas\"))\n .and_then(|schemas| {\n let key = if config.is_train {\n \"TrainingOutput\"\n } else {\n \"Output\"\n };\n schemas.get(key)\n })\n .and_then(|output| output.get(\"type\"))\n .and_then(|t| t.as_str())\n .is_some_and(|t| t == \"array\");\n\n let pool = Arc::new(PermitPool::new(num_slots));\n let sockets = transport.drain_sockets();\n\n let mut slot_readers = Vec::with_capacity(num_slots);\n for (slot_id, socket) in slot_ids.iter().zip(sockets) {\n let (read_half, write_half) = socket.into_split();\n\n let writer = FramedWrite::new(write_half, JsonCodec::::new());\n pool.add_permit(*slot_id, writer);\n\n let reader = FramedRead::new(read_half, JsonCodec::::new());\n slot_readers.push((*slot_id, reader));\n }\n\n let (register_tx, register_rx) = mpsc::channel(num_slots);\n let (healthcheck_tx, healthcheck_rx) = mpsc::channel(1);\n let (cancel_tx, cancel_rx) = mpsc::channel(16);\n\n let ctrl_writer = Arc::new(tokio::sync::Mutex::new(ctrl_writer));\n\n let handle = OrchestratorHandle {\n child,\n ctrl_writer: Arc::clone(&ctrl_writer),\n register_tx,\n healthcheck_tx,\n cancel_tx,\n slot_ids: slot_ids.clone(),\n };\n\n let pool_for_loop = Arc::clone(&pool);\n let ctrl_writer_for_loop = Arc::clone(&ctrl_writer);\n let upload_url = config.upload_url.clone();\n tokio::spawn(async move {\n run_event_loop(\n ctrl_reader,\n ctrl_writer_for_loop,\n slot_readers,\n register_rx,\n healthcheck_rx,\n cancel_rx,\n pool_for_loop,\n upload_url,\n output_is_array,\n )\n .await;\n });\n\n Ok(OrchestratorReady {\n pool,\n schema,\n handle,\n setup_logs,\n })\n}\n\n#[allow(clippy::too_many_arguments)]\nasync fn run_event_loop(\n mut ctrl_reader: FramedRead>,\n ctrl_writer: Arc<\n tokio::sync::Mutex>>,\n >,\n slot_readers: Vec<(\n SlotId,\n FramedRead>,\n )>,\n mut register_rx: mpsc::Receiver<(\n SlotId,\n Arc>,\n tokio::sync::oneshot::Sender,\n )>,\n mut healthcheck_rx: mpsc::Receiver>,\n mut cancel_rx: mpsc::Receiver,\n pool: Arc,\n upload_url: Option,\n // Schema says Output is \"type\": \"array\" — always wrap as Stream.\n // When false, the schema was unavailable or Output type is Any; fall\n // back to the predictor's is_stream flag on the Done message.\n output_is_array: bool,\n) {\n let mut predictions: HashMap>> = HashMap::new();\n let mut idle_senders: HashMap> =\n HashMap::new();\n let mut pending_healthchecks: Vec> = Vec::new();\n let mut healthcheck_counter: u64 = 0;\n let mut pending_uploads: HashMap>> = HashMap::new();\n\n let (slot_msg_tx, mut slot_msg_rx) =\n mpsc::channel::<(SlotId, Result)>(100);\n\n for (slot_id, mut reader) in slot_readers {\n let tx = slot_msg_tx.clone();\n tokio::spawn(async move {\n loop {\n let msg = reader.next().await;\n match msg {\n Some(Ok(response)) => {\n if tx.send((slot_id, Ok(response))).await.is_err() {\n break;\n }\n }\n Some(Err(e)) => {\n let _ = tx.send((slot_id, Err(e))).await;\n break;\n }\n None => {\n break;\n }\n }\n }\n tracing::debug!(%slot_id, \"Slot reader task exiting\");\n });\n }\n drop(slot_msg_tx);\n\n loop {\n tokio::select! {\n biased;\n\n ctrl_msg = ctrl_reader.next() => {\n match ctrl_msg {\n Some(Ok(ControlResponse::Idle { slot })) => {\n tracing::debug!(%slot, \"Slot idle notification received (control channel)\");\n match idle_senders.remove(&slot) {\n Some(sender) => {\n let token = InactiveSlotIdleToken::new(slot);\n if sender.send(token.activate()).is_err() {\n tracing::warn!(%slot, \"Idle token receiver dropped before idle confirmation\");\n }\n }\n None => {\n tracing::warn!(%slot, \"Received Idle for slot with no pending idle confirmation\");\n }\n\n }\n }\n Some(Ok(ControlResponse::Cancelled { slot })) => {\n tracing::debug!(%slot, \"Slot cancelled (control channel)\");\n }\n Some(Ok(ControlResponse::Failed { slot, error })) => {\n tracing::warn!(%slot, %error, \"Slot poisoned\");\n pool.poison(slot);\n if let Some(pred) = predictions.remove(&slot)\n && let Some(mut p) = try_lock_prediction(&pred)\n && !p.is_terminal()\n {\n p.set_failed(error);\n }\n }\n Some(Ok(ControlResponse::Fatal { reason })) => {\n tracing::error!(%reason, \"Worker fatal\");\n for (slot, pred) in predictions.drain() {\n tracing::warn!(%slot, \"Failing prediction due to worker fatal error\");\n pool.poison(slot);\n if let Some(mut p) = try_lock_prediction(&pred)\n && !p.is_terminal()\n {\n p.set_failed(reason.clone());\n }\n }\n let result = HealthcheckResult::unhealthy(&reason);\n for tx in pending_healthchecks.drain(..) {\n let _ = tx.send(result.clone());\n }\n break;\n }\n Some(Ok(ControlResponse::Ready { .. })) => {\n tracing::warn!(\"Unexpected Ready in event loop\");\n }\n Some(Ok(ControlResponse::Log { source: _, data })) => {\n for line in data.lines() {\n tracing::info!(target: \"coglet::user\", \"{}\", line);\n }\n }\n Some(Ok(ControlResponse::WorkerLog { target, level, message })) => {\n emit_worker_log(&target, &level, &message);\n }\n Some(Ok(ControlResponse::DroppedLogs { count, interval_millis })) => {\n tracing::trace!(count, interval_millis, \"Received DroppedLogs message\");\n let interval_secs = interval_millis as f64 / 1000.0;\n tracing::warn!(\n \"Log production exceeds consumption rate. {} logs dropped in last {:.1}s\",\n count, interval_secs\n );\n }\n Some(Ok(ControlResponse::HealthcheckResult { id: _, status, error })) => {\n tracing::trace!(\n ?status,\n ?error,\n pending_count = pending_healthchecks.len(),\n \"Received healthcheck result from worker\"\n );\n if pending_healthchecks.is_empty() {\n tracing::warn!(\"Received healthcheck result but no pending requests\");\n } else {\n let result = match status {\n HealthcheckStatus::Healthy => HealthcheckResult::healthy(),\n HealthcheckStatus::Unhealthy => {\n HealthcheckResult::unhealthy(error.unwrap_or_else(|| \"unhealthy\".to_string()))\n }\n };\n tracing::trace!(\n pending_count = pending_healthchecks.len(),\n \"Distributing healthcheck result to pending callers\"\n );\n for tx in pending_healthchecks.drain(..) {\n let _ = tx.send(result.clone());\n }\n }\n }\n Some(Ok(ControlResponse::ShuttingDown)) => {\n tracing::info!(\"Worker shutting down\");\n break;\n }\n Some(Err(e)) => {\n tracing::error!(error = %e, \"Control channel error\");\n break;\n }\n None => {\n tracing::warn!(\"Control channel closed (worker crashed?)\");\n for (slot, pred) in predictions.drain() {\n tracing::warn!(%slot, \"Failing prediction due to worker crash\");\n if let Some(mut p) = try_lock_prediction(&pred) {\n p.set_failed(\"Worker crashed\".to_string());\n }\n }\n // Fail any pending healthchecks\n for tx in pending_healthchecks.drain(..) {\n let _ = tx.send(HealthcheckResult::unhealthy(\"Worker crashed\"));\n }\n break;\n }\n }\n }\n\n Some(response_tx) = healthcheck_rx.recv() => {\n let in_flight = !pending_healthchecks.is_empty();\n pending_healthchecks.push(response_tx);\n\n // Only send to worker if no healthcheck is already in-flight.\n // Otherwise this caller just waits for the same result.\n if !in_flight {\n healthcheck_counter += 1;\n let hc_id = format!(\"hc_{}\", healthcheck_counter);\n tracing::trace!(%hc_id, \"Sending healthcheck request to worker\");\n\n let mut writer = ctrl_writer.lock().await;\n if let Err(e) = writer.send(ControlRequest::Healthcheck { id: hc_id }).await {\n tracing::error!(error = %e, \"Failed to send healthcheck request\");\n let result = HealthcheckResult::unhealthy(format!(\"Failed to send: {}\", e));\n for tx in pending_healthchecks.drain(..) {\n let _ = tx.send(result.clone());\n }\n }\n } else {\n tracing::trace!(\n pending_count = pending_healthchecks.len(),\n \"Healthcheck already in-flight, coalescing request\"\n );\n }\n }\n\n Some(prediction_id) = cancel_rx.recv() => {\n // Resolve prediction_id → slot_id by iterating (fine for small concurrency)\n let slot = predictions.iter().find_map(|(sid, pred)| {\n try_lock_prediction(pred)\n .filter(|p| p.id() == prediction_id)\n .map(|_| *sid)\n });\n match slot {\n Some(slot_id) => {\n tracing::info!(\n target: \"coglet::prediction\",\n %prediction_id,\n %slot_id,\n \"Cancelling prediction\"\n );\n let mut writer = ctrl_writer.lock().await;\n if let Err(e) = writer.send(ControlRequest::Cancel { slot: slot_id }).await {\n tracing::error!(\n %slot_id,\n error = %e,\n \"Failed to send cancel request to worker\"\n );\n }\n // Also abort any pending upload tasks for this slot\n if let Some(handles) = pending_uploads.remove(&slot_id) {\n for h in handles { h.abort(); }\n }\n }\n None => {\n tracing::debug!(%prediction_id, \"Cancel requested for unknown prediction (may have already completed)\");\n }\n }\n }\n\n Some((slot_id, prediction, idle_sender)) = register_rx.recv() => {\n let prediction_id = match try_lock_prediction(&prediction) {\n Some(p) => p.id().to_string(),\n None => {\n // Mutex poisoned during registration - prediction already failed\n tracing::error!(%slot_id, \"Prediction mutex poisoned during registration\");\n continue;\n }\n };\n // NOTE: we insert the idle sender, and idle senders are only removed on consumption of the\n // `tokio::sync::oneshot::Sender`, this means the only time we'll leak memory here is if the\n // slot is poisoned or otherwise in a bad state. It is intentional that we don't remove idle\n // senders in any other case.\n idle_senders.insert(slot_id, idle_sender);\n tracing::info!(\n target: \"coglet::prediction\",\n %prediction_id,\n \"Starting prediction\"\n );\n tracing::debug!(%slot_id, %prediction_id, \"Registered prediction\");\n predictions.insert(slot_id, prediction);\n }\n\n Some((slot_id, result)) = slot_msg_rx.recv() => {\n match result {\n Ok(SlotResponse::Log { source, data }) => {\n let (prediction_id, poisoned) = if let Some(pred) = predictions.get(&slot_id) {\n if let Some(mut p) = try_lock_prediction(pred) {\n p.append_log(&data);\n (Some(p.id().to_string()), false)\n } else {\n (None, true)\n }\n } else {\n (None, false)\n };\n // Remove poisoned predictions outside the borrow\n if poisoned {\n predictions.remove(&slot_id);\n }\n\n let trimmed = data.trim();\n if !trimmed.is_empty() {\n if let Some(id) = prediction_id {\n tracing::info!(\n target: \"coglet::prediction\",\n prediction_id = %id,\n source = ?source,\n \"{}\",\n trimmed\n );\n } else {\n tracing::warn!(\n target: \"coglet::prediction\",\n prediction_id = \"NO_ACTIVE_PREDICTION\",\n source = ?source,\n \"{}\",\n trimmed\n );\n }\n }\n }\n Ok(SlotResponse::Metric { name, value, mode }) => {\n let poisoned = if let Some(pred) = predictions.get(&slot_id) {\n if let Some(mut p) = try_lock_prediction(pred) {\n p.set_metric(name, value, mode);\n false\n } else {\n true\n }\n } else {\n false\n };\n if poisoned {\n predictions.remove(&slot_id);\n }\n }\n Ok(SlotResponse::Output { output }) => {\n let poisoned = if let Some(pred) = predictions.get(&slot_id) {\n if let Some(mut p) = try_lock_prediction(pred) {\n p.append_output(output);\n false\n } else {\n true\n }\n } else {\n false\n };\n // Remove poisoned predictions outside the borrow\n if poisoned {\n predictions.remove(&slot_id);\n }\n }\n Ok(SlotResponse::FileOutput { filename, kind, mime_type }) => {\n tracing::debug!(%slot_id, %filename, ?kind, \"FileOutput received\");\n let bytes = match std::fs::read(&filename) {\n Ok(b) => b,\n Err(e) => {\n tracing::error!(%slot_id, %filename, error = %e, \"Failed to read FileOutput\");\n continue;\n }\n };\n match kind {\n FileOutputKind::Oversized => {\n let output: serde_json::Value = match serde_json::from_slice(&bytes) {\n Ok(val) => val,\n Err(e) => {\n tracing::error!(%slot_id, %filename, error = %e, \"Failed to parse oversized JSON\");\n continue;\n }\n };\n let poisoned = if let Some(pred) = predictions.get(&slot_id) {\n if let Some(mut p) = try_lock_prediction(pred) {\n p.append_output(output);\n false\n } else {\n true\n }\n } else {\n false\n };\n if poisoned {\n predictions.remove(&slot_id);\n }\n }\n FileOutputKind::FileType => {\n let mime = mime_type.unwrap_or_else(|| {\n mime_guess::from_path(&filename)\n .first_or_octet_stream()\n .to_string()\n });\n if let Some(ref url) = upload_url {\n // Spawn upload task so we don't block the event loop\n let pred = predictions.get(&slot_id).cloned();\n let endpoint = ensure_trailing_slash(url);\n let basename = std::path::Path::new(&filename)\n .file_name()\n .and_then(|n| n.to_str())\n .unwrap_or(\"output\")\n .to_string();\n let handle = tokio::spawn(async move {\n match upload_file(&endpoint, &basename, &bytes, &mime).await {\n Ok(url) => {\n if let Some(pred) = pred\n && let Some(mut p) = try_lock_prediction(&pred)\n {\n p.append_output(serde_json::Value::String(url));\n }\n }\n Err(e) => {\n tracing::error!(error = %e, \"Failed to upload file output\");\n }\n }\n });\n pending_uploads.entry(slot_id).or_default().push(handle);\n } else {\n // No upload URL — base64-encode as data URI\n use base64::Engine;\n let encoded = base64::engine::general_purpose::STANDARD\n .encode(&bytes);\n let output = serde_json::Value::String(format!(\n \"data:{mime};base64,{encoded}\"\n ));\n let poisoned = if let Some(pred) = predictions.get(&slot_id) {\n if let Some(mut p) = try_lock_prediction(pred) {\n p.append_output(output);\n false\n } else {\n true\n }\n } else {\n false\n };\n if poisoned {\n predictions.remove(&slot_id);\n }\n }\n }\n }\n }\n Ok(SlotResponse::Done { id, output: _, predict_time, is_stream }) => {\n tracing::info!(\n target: \"coglet::prediction\",\n prediction_id = %id,\n predict_time,\n is_stream,\n output_is_array,\n \"Prediction succeeded\"\n );\n let uploads = pending_uploads.remove(&slot_id).unwrap_or_default();\n if let Some(pred) = predictions.remove(&slot_id) {\n if uploads.is_empty() {\n // No pending uploads — complete synchronously to avoid\n // a race between tokio::spawn and Notify::notified() in\n // service.rs. notify_waiters() only wakes already-\n // registered waiters; spawning a task can fire the\n // notification before the service registers its waiter.\n if let Some(mut p) = try_lock_prediction(&pred) {\n let pred_output = wrap_outputs(\n p.take_outputs(),\n output_is_array,\n is_stream,\n );\n p.set_succeeded(pred_output);\n }\n } else {\n // Has pending uploads — must spawn to await them.\n // Clone the cancel token so we can abort uploads if\n // the prediction is cancelled while uploads are in flight.\n let (cancel_token, upload_pred_id) = match try_lock_prediction(&pred) {\n Some(p) => (Some(p.cancel_token()), p.id().to_string()),\n None => (None, id.clone()),\n };\n tokio::spawn(async move {\n if let Some(token) = cancel_token {\n let upload_fut = futures::future::join_all(uploads);\n tokio::pin!(upload_fut);\n tokio::select! {\n _ = &mut upload_fut => {}\n _ = token.cancelled() => {\n tracing::info!(\n target: \"coglet::prediction\",\n prediction_id = %upload_pred_id,\n \"Aborting in-flight uploads due to cancellation\"\n );\n if let Some(mut p) = try_lock_prediction(&pred) {\n p.set_canceled();\n }\n return;\n }\n }\n } else {\n for h in uploads {\n let _ = h.await;\n }\n }\n if let Some(mut p) = try_lock_prediction(&pred) {\n let pred_output = wrap_outputs(\n p.take_outputs(),\n output_is_array,\n is_stream,\n );\n p.set_succeeded(pred_output);\n }\n });\n }\n } else {\n tracing::warn!(%slot_id, %id, \"Prediction not found for Done message\");\n }\n }\n Ok(SlotResponse::Failed { id, error }) => {\n tracing::info!(\n target: \"coglet::prediction\",\n prediction_id = %id,\n %error,\n \"Prediction failed\"\n );\n // Abort any pending uploads — prediction is terminal\n if let Some(handles) = pending_uploads.remove(&slot_id) {\n for h in handles { h.abort(); }\n }\n if let Some(pred) = predictions.remove(&slot_id)\n && let Some(mut p) = try_lock_prediction(&pred)\n {\n p.set_failed(error);\n }\n }\n Ok(SlotResponse::Cancelled { id }) => {\n tracing::info!(\n target: \"coglet::prediction\",\n prediction_id = %id,\n \"Prediction cancelled\"\n );\n // Abort any pending uploads — prediction is terminal\n if let Some(handles) = pending_uploads.remove(&slot_id) {\n for h in handles { h.abort(); }\n }\n if let Some(pred) = predictions.remove(&slot_id)\n && let Some(mut p) = try_lock_prediction(&pred)\n {\n p.set_canceled();\n }\n }\n Err(e) => {\n tracing::error!(%slot_id, error = %e, \"Slot socket error\");\n if let Some(handles) = pending_uploads.remove(&slot_id) {\n for h in handles { h.abort(); }\n }\n if let Some(pred) = predictions.remove(&slot_id)\n && let Some(mut p) = try_lock_prediction(&pred)\n {\n p.set_failed(format!(\"Slot socket error: {}\", e));\n }\n }\n }\n }\n }\n }\n\n tracing::info!(\"Event loop exiting\");\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use serde_json::json;\n\n // ── wrap_outputs: schema says array (output_is_array = true) ──\n\n #[test]\n fn wrap_outputs_schema_array_empty() {\n // List[Path] that returned no items → empty array\n let result = wrap_outputs(vec![], true, true);\n assert!(result.is_stream());\n assert_eq!(result.into_values(), Vec::::new());\n }\n\n #[test]\n fn wrap_outputs_schema_array_single_item() {\n // List[Path] with num_outputs=1 → [\"url\"] not \"url\"\n let result = wrap_outputs(vec![json!(\"https://example.com/img.png\")], true, true);\n assert!(result.is_stream());\n assert_eq!(\n result.into_values(),\n vec![json!(\"https://example.com/img.png\")]\n );\n }\n\n #[test]\n fn wrap_outputs_schema_array_multiple_items() {\n // List[Path] with num_outputs=4\n let items = vec![\n json!(\"https://example.com/1.png\"),\n json!(\"https://example.com/2.png\"),\n json!(\"https://example.com/3.png\"),\n json!(\"https://example.com/4.png\"),\n ];\n let result = wrap_outputs(items.clone(), true, true);\n assert!(result.is_stream());\n assert_eq!(result.into_values(), items);\n }\n\n #[test]\n fn wrap_outputs_schema_array_overrides_is_stream_false() {\n // Schema says array but predictor didn't set is_stream (shouldn't happen,\n // but schema is authoritative)\n let result = wrap_outputs(vec![json!(\"url\")], true, false);\n assert!(result.is_stream());\n }\n\n // ── wrap_outputs: predictor signal (is_stream = true, no schema) ──\n\n #[test]\n fn wrap_outputs_predictor_stream_empty() {\n // Generator that yielded nothing, no schema\n let result = wrap_outputs(vec![], false, true);\n assert!(result.is_stream());\n assert_eq!(result.into_values(), Vec::::new());\n }\n\n #[test]\n fn wrap_outputs_predictor_stream_single_item() {\n // Any-typed list with one element, no schema\n let result = wrap_outputs(vec![json!(\"only_item\")], false, true);\n assert!(result.is_stream());\n assert_eq!(result.into_values(), vec![json!(\"only_item\")]);\n }\n\n #[test]\n fn wrap_outputs_predictor_stream_multiple_items() {\n // Generator yielding multiple, no schema\n let items = vec![json!(\"a\"), json!(\"b\"), json!(\"c\")];\n let result = wrap_outputs(items.clone(), false, true);\n assert!(result.is_stream());\n assert_eq!(result.into_values(), items);\n }\n\n // ── wrap_outputs: scalar output (neither schema array nor predictor stream) ──\n\n #[test]\n fn wrap_outputs_scalar_empty() {\n // Single output that was null (e.g. Path sent via FileOutput, not yet resolved?)\n let result = wrap_outputs(vec![], false, false);\n assert!(!result.is_stream());\n assert_eq!(result.final_value(), &json!(null));\n }\n\n #[test]\n fn wrap_outputs_scalar_single() {\n // return Path(\"output.png\") → single string\n let result = wrap_outputs(vec![json!(\"https://example.com/output.png\")], false, false);\n assert!(!result.is_stream());\n assert_eq!(\n result.final_value(),\n &json!(\"https://example.com/output.png\")\n );\n }\n\n #[test]\n fn wrap_outputs_scalar_multiple_falls_back_to_stream() {\n // Shouldn't happen for scalar returns, but if multiple items arrive\n // with neither flag set, Stream is the safe choice\n let items = vec![json!(\"a\"), json!(\"b\")];\n let result = wrap_outputs(items.clone(), false, false);\n assert!(result.is_stream());\n assert_eq!(result.into_values(), items);\n }\n\n // ── Serialization: is_stream field on Done message ──\n\n #[test]\n fn done_is_stream_false_omitted_from_json() {\n let msg = SlotResponse::Done {\n id: \"p1\".into(),\n output: None,\n predict_time: 1.0,\n is_stream: false,\n };\n let json = serde_json::to_value(&msg).unwrap();\n assert!(\n json.get(\"is_stream\").is_none(),\n \"is_stream=false should be omitted\"\n );\n }\n\n #[test]\n fn done_is_stream_true_present_in_json() {\n let msg = SlotResponse::Done {\n id: \"p1\".into(),\n output: None,\n predict_time: 1.0,\n is_stream: true,\n };\n let json = serde_json::to_value(&msg).unwrap();\n assert_eq!(json.get(\"is_stream\"), Some(&json!(true)));\n }\n\n #[test]\n fn done_without_is_stream_deserializes_as_false() {\n // Backward compat: old workers won't send is_stream\n let json = json!({\n \"type\": \"done\",\n \"id\": \"p1\",\n \"predict_time\": 1.0\n });\n let msg: SlotResponse = serde_json::from_value(json).unwrap();\n match msg {\n SlotResponse::Done { is_stream, .. } => assert!(!is_stream),\n _ => panic!(\"wrong variant\"),\n }\n }\n}\n" }, { "path": "crates/coglet/src/permit/mod.rs", "content": "//! Permit pool for concurrent slot management.\n//!\n//! The permit system uses typestate to enforce valid state transitions at compile time:\n//! - `PermitInUse` → `PermitIdle` via `into_idle()` (returns to pool on drop)\n//! - `PermitInUse` → `PermitPoisoned` via `into_poisoned()` (orphaned on drop)\n//! - `PermitPoisoned` → `PermitIdle`: NOT POSSIBLE (no method exists)\n//!\n//! Slot poisoning is a pool-level property (`PermitPool::poison()`). A poisoned slot\n//! is permanently removed from the pool regardless of whether a prediction was active.\n//! `PermitIdle::drop` checks the pool-level poison flag and skips returning the permit.\n\nmod pool;\nmod slot;\n\npub use pool::{\n AnyPermit, InactiveSlotIdleToken, PermitError, PermitIdle, PermitInUse, PermitPoisoned,\n PermitPool, SlotIdleToken,\n};\npub use slot::{PredictionSlot, UnregisteredPredictionSlot};\n" }, { "path": "crates/coglet/src/permit/pool.rs", "content": "//! Permit pool implementation with typestate for compile-time state transition safety.\n//!\n//! Slot poisoning is a pool-level property: a poisoned slot is permanently removed\n//! from the pool regardless of whether a prediction was active on it.\n\nuse std::sync::Arc;\nuse std::sync::Mutex as StdMutex;\nuse std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};\n\nuse futures::SinkExt;\nuse tokio::net::unix::OwnedWriteHalf;\nuse tokio::sync::{Mutex, mpsc};\nuse tokio_util::codec::FramedWrite;\n\nuse crate::bridge::codec::JsonCodec;\nuse crate::bridge::protocol::{SlotId, SlotRequest};\n\npub(crate) struct PermitInner {\n pub slot_id: SlotId,\n pub writer: FramedWrite>,\n pub idle_flag: Arc,\n pub poisoned: Arc,\n}\n\nstruct PoolConnection {\n pool_tx: mpsc::Sender,\n pool_available: Arc,\n}\n\nimpl Clone for PoolConnection {\n fn clone(&self) -> Self {\n Self {\n pool_tx: self.pool_tx.clone(),\n pool_available: Arc::clone(&self.pool_available),\n }\n }\n}\n\n/// A permit actively running a prediction.\npub struct PermitInUse {\n slot_id: SlotId,\n writer: Option>>,\n idle_flag: Arc,\n poisoned: Arc,\n pool: PoolConnection,\n}\n\nimpl PermitInUse {\n pub(crate) fn new(\n inner: PermitInner,\n pool_tx: mpsc::Sender,\n pool_available: Arc,\n ) -> Self {\n inner.idle_flag.store(false, Ordering::Release);\n\n Self {\n slot_id: inner.slot_id,\n writer: Some(inner.writer),\n idle_flag: inner.idle_flag,\n poisoned: inner.poisoned,\n pool: PoolConnection {\n pool_tx,\n pool_available,\n },\n }\n }\n\n pub fn slot_id(&self) -> SlotId {\n self.slot_id\n }\n\n /// Transition to idle state - permit will return to pool on drop\n /// (unless the slot has been poisoned at the pool level).\n pub fn into_idle(mut self) -> PermitIdle {\n self.idle_flag.store(true, Ordering::Release);\n PermitIdle {\n slot_id: self.slot_id,\n writer: self.writer.take(),\n idle_flag: Arc::clone(&self.idle_flag),\n poisoned: Arc::clone(&self.poisoned),\n pool: self.pool.clone(),\n }\n }\n\n /// Transition to poisoned state - permit will NOT return to pool.\n ///\n /// Also sets the pool-level poison flag so the slot is never reused.\n pub fn into_poisoned(mut self) -> PermitPoisoned {\n self.poisoned.store(true, Ordering::Release);\n PermitPoisoned {\n slot_id: self.slot_id,\n _writer: self.writer.take(),\n }\n }\n\n pub async fn send(&mut self, request: SlotRequest) -> Result<(), PermitError> {\n let writer = self.writer.as_mut().ok_or(PermitError::Consumed)?;\n writer\n .send(request)\n .await\n .map_err(|e| PermitError::Send(e.to_string()))\n }\n}\n\nimpl Drop for PermitInUse {\n fn drop(&mut self) {\n if self.writer.is_some() && !self.poisoned.load(Ordering::Acquire) {\n tracing::error!(slot = %self.slot_id, \"PermitInUse dropped without state transition\");\n }\n }\n}\n\n/// A permit that completed successfully - returns to pool on drop\n/// (unless the slot has been poisoned at the pool level).\npub struct PermitIdle {\n slot_id: SlotId,\n writer: Option>>,\n idle_flag: Arc,\n poisoned: Arc,\n pool: PoolConnection,\n}\n\nimpl PermitIdle {\n pub fn slot_id(&self) -> SlotId {\n self.slot_id\n }\n}\n\nimpl Drop for PermitIdle {\n fn drop(&mut self) {\n // If the slot was poisoned at the pool level, don't return it.\n if self.poisoned.load(Ordering::Acquire) {\n tracing::warn!(slot = %self.slot_id, \"Slot poisoned - not returning to pool\");\n return;\n }\n\n if let Some(writer) = self.writer.take() {\n let inner = PermitInner {\n slot_id: self.slot_id,\n writer,\n idle_flag: Arc::clone(&self.idle_flag),\n poisoned: Arc::clone(&self.poisoned),\n };\n\n if self.pool.pool_tx.try_send(inner).is_ok() {\n self.pool.pool_available.fetch_add(1, Ordering::Release);\n }\n }\n }\n}\n\n/// A poisoned permit - slot permanently failed, will NOT return to pool.\npub struct PermitPoisoned {\n slot_id: SlotId,\n _writer: Option>>,\n}\n\nimpl PermitPoisoned {\n pub fn slot_id(&self) -> SlotId {\n self.slot_id\n }\n}\n\nimpl Drop for PermitPoisoned {\n fn drop(&mut self) {\n tracing::warn!(slot = %self.slot_id, \"Slot poisoned - capacity reduced\");\n }\n}\n\n/// A permit in any state (for containers needing dynamic state).\npub enum AnyPermit {\n InUse(PermitInUse),\n Idle(PermitIdle),\n Poisoned(PermitPoisoned),\n}\n\nimpl AnyPermit {\n pub fn slot_id(&self) -> SlotId {\n match self {\n AnyPermit::InUse(p) => p.slot_id(),\n AnyPermit::Idle(p) => p.slot_id(),\n AnyPermit::Poisoned(p) => p.slot_id(),\n }\n }\n\n pub fn is_idle(&self) -> bool {\n matches!(self, AnyPermit::Idle(_))\n }\n\n pub fn is_poisoned(&self) -> bool {\n matches!(self, AnyPermit::Poisoned(_))\n }\n\n pub fn is_in_use(&self) -> bool {\n matches!(self, AnyPermit::InUse(_))\n }\n}\n\n#[must_use = \"must be activated to enable slot idle transition\"]\n#[derive(Debug)]\npub struct InactiveSlotIdleToken {\n slot_id: SlotId,\n}\n\nimpl InactiveSlotIdleToken {\n pub fn new(slot_id: SlotId) -> Self {\n Self { slot_id }\n }\n\n pub fn slot_id(&self) -> SlotId {\n self.slot_id\n }\n\n pub fn activate(self) -> SlotIdleToken {\n SlotIdleToken {\n slot_id: self.slot_id,\n create_time: std::time::Instant::now(),\n alarm_handle: tokio::spawn(async move {\n // This task exists solely to alert if the token isn't consumed within a reasonable time.\n // If we see this alert, it means the slot won't return to the pool until the process exits.\n tokio::time::sleep(SlotIdleToken::ALERT_THRESHOLD).await;\n tracing::error!(slot = %self.slot_id, \"IdleToken not consumed after 5s - slot will not return to pool\");\n }),\n }\n }\n}\n\n/// Token confirming the worker has marked the slot as idle, allowing the permit to return to the pool on drop.\n#[must_use = \"IdleToken confirms the worker has marked the slot as idle\"]\n#[derive(Debug)]\npub struct SlotIdleToken {\n pub(crate) slot_id: SlotId,\n pub(crate) create_time: std::time::Instant,\n pub(crate) alarm_handle: tokio::task::JoinHandle<()>,\n}\n\nimpl SlotIdleToken {\n const ALERT_THRESHOLD: std::time::Duration = std::time::Duration::from_secs(5);\n\n pub fn slot_id(&self) -> SlotId {\n self.slot_id\n }\n\n pub fn consume(self) {\n let elapsed = self.create_time.elapsed();\n if elapsed > Self::ALERT_THRESHOLD {\n tracing::warn!(slot = %self.slot_id, latency = ?elapsed, \"Delayed IdleToken Consumption\");\n }\n tracing::debug!(slot = %self.slot_id, \"IdleToken consumed\");\n }\n}\n\nimpl Drop for SlotIdleToken {\n fn drop(&mut self) {\n self.alarm_handle.abort();\n }\n}\n\n#[derive(Debug, Clone, thiserror::Error)]\npub enum PermitError {\n #[error(\"Permit already consumed\")]\n Consumed,\n #[error(\"Failed to send on slot socket: {0}\")]\n Send(String),\n}\n\n/// Pool of prediction slot permits.\n///\n/// Slot poisoning is tracked here. A poisoned slot is permanently removed\n/// from the pool — its permit will not be returned or acquired again.\npub struct PermitPool {\n available_rx: Mutex>,\n available_tx: mpsc::Sender,\n num_slots: usize,\n available_count: Arc,\n /// Per-slot poison flags, shared with permits for fast checking.\n poison_flags: StdMutex)>>,\n}\n\nimpl PermitPool {\n pub fn new(num_slots: usize) -> Self {\n let (tx, rx) = mpsc::channel(num_slots);\n\n Self {\n available_rx: Mutex::new(rx),\n available_tx: tx,\n num_slots,\n available_count: Arc::new(AtomicUsize::new(0)),\n poison_flags: StdMutex::new(Vec::with_capacity(num_slots)),\n }\n }\n\n pub fn add_permit(\n &self,\n slot_id: SlotId,\n writer: FramedWrite>,\n ) {\n let poisoned = Arc::new(AtomicBool::new(false));\n\n // Store the flag for external poisoning.\n if let Ok(mut flags) = self.poison_flags.lock() {\n flags.push((slot_id, Arc::clone(&poisoned)));\n }\n\n let inner = PermitInner {\n slot_id,\n writer,\n idle_flag: Arc::new(AtomicBool::new(true)),\n poisoned,\n };\n\n if let Err(e) = self.available_tx.try_send(inner) {\n tracing::error!(slot = %slot_id, error = %e, \"Failed to add permit to pool\");\n } else {\n self.available_count.fetch_add(1, Ordering::Release);\n }\n }\n\n /// Poison a slot. The permit will not be returned to the pool.\n ///\n /// This works whether the slot is idle (in the pool) or in use (held by a prediction).\n /// - Idle: the permit will be discarded on next `acquire`/`try_acquire`.\n /// - In use: `PermitIdle::drop` will see the flag and not return it.\n pub fn poison(&self, slot_id: SlotId) {\n if let Ok(flags) = self.poison_flags.lock() {\n for (id, flag) in flags.iter() {\n if *id == slot_id {\n if !flag.swap(true, Ordering::AcqRel) {\n tracing::warn!(slot = %slot_id, \"Slot poisoned - capacity permanently reduced\");\n }\n return;\n }\n }\n }\n tracing::warn!(slot = %slot_id, \"Attempted to poison unknown slot\");\n }\n\n /// Check if a slot is poisoned.\n pub fn is_poisoned(&self, slot_id: SlotId) -> bool {\n if let Ok(flags) = self.poison_flags.lock() {\n for (id, flag) in flags.iter() {\n if *id == slot_id {\n return flag.load(Ordering::Acquire);\n }\n }\n }\n false\n }\n\n pub fn try_acquire(&self) -> Option {\n let mut rx = self.available_rx.try_lock().ok()?;\n loop {\n let inner = rx.try_recv().ok()?;\n self.available_count.fetch_sub(1, Ordering::Release);\n\n // Skip poisoned permits — they're permanently dead.\n if inner.poisoned.load(Ordering::Acquire) {\n tracing::debug!(slot = %inner.slot_id, \"Discarding poisoned permit from pool\");\n continue;\n }\n\n return Some(PermitInUse::new(\n inner,\n self.available_tx.clone(),\n Arc::clone(&self.available_count),\n ));\n }\n }\n\n pub async fn acquire(&self) -> Option {\n let mut rx = self.available_rx.lock().await;\n loop {\n let inner = rx.recv().await?;\n self.available_count.fetch_sub(1, Ordering::Release);\n\n // Skip poisoned permits — they're permanently dead.\n if inner.poisoned.load(Ordering::Acquire) {\n tracing::debug!(slot = %inner.slot_id, \"Discarding poisoned permit from pool\");\n continue;\n }\n\n return Some(PermitInUse::new(\n inner,\n self.available_tx.clone(),\n Arc::clone(&self.available_count),\n ));\n }\n }\n\n pub fn num_slots(&self) -> usize {\n self.num_slots\n }\n\n pub fn available(&self) -> usize {\n self.available_count.load(Ordering::Acquire)\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use tokio::net::UnixStream;\n\n async fn make_socket_pair() -> (OwnedWriteHalf, tokio::net::unix::OwnedReadHalf) {\n let (a, b) = UnixStream::pair().unwrap();\n let (read, write) = a.into_split();\n let _ = b;\n (write, read)\n }\n\n #[tokio::test]\n async fn pool_add_and_acquire() {\n let pool = PermitPool::new(2);\n\n let (write1, _read1) = make_socket_pair().await;\n let (write2, _read2) = make_socket_pair().await;\n\n let slot1 = SlotId::new();\n let slot2 = SlotId::new();\n\n pool.add_permit(slot1, FramedWrite::new(write1, JsonCodec::new()));\n pool.add_permit(slot2, FramedWrite::new(write2, JsonCodec::new()));\n\n let p1 = pool.try_acquire();\n assert!(p1.is_some());\n\n let p2 = pool.try_acquire();\n assert!(p2.is_some());\n\n let p3 = pool.try_acquire();\n assert!(p3.is_none());\n }\n\n #[tokio::test]\n async fn permit_returns_to_pool_when_idle() {\n let pool = PermitPool::new(1);\n\n let (write, _read) = make_socket_pair().await;\n let slot = SlotId::new();\n\n pool.add_permit(slot, FramedWrite::new(write, JsonCodec::new()));\n\n {\n let permit = pool.try_acquire().unwrap();\n let _idle_permit = permit.into_idle();\n }\n\n let permit = pool.try_acquire();\n assert!(permit.is_some());\n }\n\n #[tokio::test]\n async fn permit_orphaned_when_poisoned() {\n let pool = PermitPool::new(1);\n\n let (write, _read) = make_socket_pair().await;\n let slot = SlotId::new();\n\n pool.add_permit(slot, FramedWrite::new(write, JsonCodec::new()));\n\n {\n let permit = pool.try_acquire().unwrap();\n let _poisoned_permit = permit.into_poisoned();\n }\n\n let permit = pool.try_acquire();\n assert!(permit.is_none());\n }\n\n #[tokio::test]\n async fn pool_poison_idle_slot() {\n // Poison a slot while it's idle in the pool — acquire should skip it.\n let pool = PermitPool::new(2);\n\n let (write1, _read1) = make_socket_pair().await;\n let (write2, _read2) = make_socket_pair().await;\n\n let slot1 = SlotId::new();\n let slot2 = SlotId::new();\n\n pool.add_permit(slot1, FramedWrite::new(write1, JsonCodec::new()));\n pool.add_permit(slot2, FramedWrite::new(write2, JsonCodec::new()));\n\n assert!(!pool.is_poisoned(slot1));\n pool.poison(slot1);\n assert!(pool.is_poisoned(slot1));\n assert!(!pool.is_poisoned(slot2));\n\n // First acquire should skip poisoned slot1, return slot2.\n let permit = pool.try_acquire().unwrap();\n assert_eq!(permit.slot_id(), slot2);\n\n // No more permits available.\n assert!(pool.try_acquire().is_none());\n }\n\n #[tokio::test]\n async fn pool_poison_in_use_slot_prevents_return() {\n // Poison a slot while a prediction holds it — into_idle + drop should NOT return it.\n let pool = PermitPool::new(1);\n\n let (write, _read) = make_socket_pair().await;\n let slot = SlotId::new();\n\n pool.add_permit(slot, FramedWrite::new(write, JsonCodec::new()));\n\n {\n let permit = pool.try_acquire().unwrap();\n // Poison while in use.\n pool.poison(slot);\n // Transition to idle — drop should see the poison flag.\n let _idle = permit.into_idle();\n }\n\n // Permit should NOT have returned to the pool.\n assert!(pool.try_acquire().is_none());\n }\n\n #[tokio::test]\n async fn pool_poison_is_idempotent() {\n let pool = PermitPool::new(1);\n\n let (write, _read) = make_socket_pair().await;\n let slot = SlotId::new();\n\n pool.add_permit(slot, FramedWrite::new(write, JsonCodec::new()));\n\n pool.poison(slot);\n pool.poison(slot); // Should not panic or double-count.\n assert!(pool.is_poisoned(slot));\n }\n}\n" }, { "path": "crates/coglet/src/permit/slot.rs", "content": "//! PredictionSlot - holds Prediction and Permit side-by-side.\n//!\n//! This separation allows the prediction to be behind Mutex for concurrent\n//! updates while the permit's idle_flag can be set without holding the lock.\n//!\n//! Slot poisoning is NOT managed here — it's a pool-level property.\n//! The slot always transitions to idle when done; `PermitIdle::drop` checks\n//! the pool-level poison flag to decide whether to return the permit.\n\nuse std::sync::{Arc, Mutex};\n\nuse super::{AnyPermit, PermitInUse, SlotIdleToken};\nuse crate::bridge::protocol::SlotId;\nuse crate::prediction::Prediction;\n\n#[derive(Debug, Clone, thiserror::Error)]\npub enum SlotError {\n #[error(\"receive error while waiting for idle token\")]\n IdleTokenReceiveError(#[from] tokio::sync::oneshot::error::RecvError),\n #[error(\"permit already consumed\")]\n PermitAlreadyConsumed,\n}\n\n/// Pre-registration slot state - holds prediction while permit is being\n/// acquired and slot is being registered with orchestrator.\npub struct UnregisteredPredictionSlot {\n prediction_slot: PredictionSlot,\n idle_tx: tokio::sync::oneshot::Sender,\n}\n\nimpl UnregisteredPredictionSlot {\n pub fn new(\n prediction_slot: PredictionSlot,\n idle_tx: tokio::sync::oneshot::Sender,\n ) -> Self {\n Self {\n prediction_slot,\n idle_tx,\n }\n }\n\n /// Consumes the unregistered slot and returns its components for registration.\n pub fn into_parts(self) -> (tokio::sync::oneshot::Sender, PredictionSlot) {\n (self.idle_tx, self.prediction_slot)\n }\n\n pub fn prediction(&self) -> Arc> {\n self.prediction_slot.prediction()\n }\n}\n\n/// Holds a prediction and its permit side-by-side.\n///\n/// On drop: Permit returns to pool (if idle and not poisoned at pool level).\npub struct PredictionSlot {\n prediction: Arc>,\n slot_id: SlotId,\n permit: Option,\n idle_rx: Option>,\n}\n\nimpl PredictionSlot {\n pub fn new(\n prediction: Prediction,\n permit: PermitInUse,\n idle_rx: tokio::sync::oneshot::Receiver,\n ) -> Self {\n let slot_id = permit.slot_id();\n Self {\n prediction: Arc::new(Mutex::new(prediction)),\n slot_id,\n permit: Some(AnyPermit::InUse(permit)),\n idle_rx: Some(idle_rx),\n }\n }\n\n pub fn prediction(&self) -> Arc> {\n Arc::clone(&self.prediction)\n }\n\n pub fn permit_mut(&mut self) -> Option<&mut PermitInUse> {\n match &mut self.permit {\n Some(AnyPermit::InUse(p)) => Some(p),\n _ => None,\n }\n }\n\n pub fn slot_id(&self) -> SlotId {\n self.slot_id\n }\n\n /// Mark the slot as idle - permit will return to pool on drop (unless the slot has\n /// been poisoned at the pool level). Awaits until the idle token is received, which\n /// ensures the slot has been confirmed idle by the worker. If the idle token is not\n /// received, the permit is not returned to the pool,\n #[must_use = \"into_idle confirms the slot is idle and allows the permit to return to the pool on drop\"]\n pub async fn into_idle(mut self) -> Result<(), SlotError> {\n if let Some(receiver) = self.idle_rx.take() {\n let idle_token = receiver.await?;\n debug_assert_eq!(\n idle_token.slot_id(),\n self.slot_id,\n \"IdleToken slot_id mismatch\"\n );\n idle_token.consume();\n }\n\n let permit = self.permit.take();\n debug_assert!(\n permit.is_some(),\n \"Attempted to mark slot as idle but permit was already consumed\"\n );\n match permit {\n Some(AnyPermit::InUse(p)) => {\n let idle = p.into_idle();\n self.permit = Some(AnyPermit::Idle(idle));\n Ok(())\n }\n Some(AnyPermit::Idle(p)) => {\n self.permit = Some(AnyPermit::Idle(p));\n Ok(())\n }\n Some(AnyPermit::Poisoned(p)) => {\n // Permit was explicitly poisoned (legacy path) — keep it.\n debug_assert!(false, \"Cannot mark poisoned slot as idle\");\n tracing::error!(slot = %p.slot_id(), \"Bug: attempted to mark poisoned slot as idle\");\n self.permit = Some(AnyPermit::Poisoned(p));\n Ok(())\n }\n None => {\n // Permit was already consumed (bug) — log and do nothing.\n tracing::error!(slot = %self.slot_id(), \"Bug: attempted to mark slot as idle but permit was already consumed\");\n Err(SlotError::PermitAlreadyConsumed)\n }\n }\n }\n\n pub fn is_idle(&self) -> bool {\n self.permit.as_ref().is_some_and(|p| p.is_idle())\n }\n\n pub fn id(&self) -> String {\n self.prediction\n .try_lock()\n .map(|p| p.id().to_string())\n .unwrap_or_default()\n }\n}\n\nimpl Drop for PredictionSlot {\n fn drop(&mut self) {\n if let Some(AnyPermit::InUse(_)) = &self.permit\n && let Ok(mut prediction) = self.prediction.try_lock()\n && !prediction.is_terminal()\n {\n tracing::error!(\n slot = %self.slot_id(),\n prediction_id = %prediction.id(),\n \"Slot dropped while InUse with non-terminal prediction\"\n );\n prediction.set_failed(\"Slot dropped unexpectedly\".to_string());\n }\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use crate::bridge::codec::JsonCodec;\n use crate::permit::{InactiveSlotIdleToken, PermitPool};\n use tokio::net::UnixStream;\n use tokio_util::codec::FramedWrite;\n\n #[tokio::test]\n async fn slot_creation() {\n let pool = PermitPool::new(1);\n\n let (a, _b) = UnixStream::pair().unwrap();\n let (_, write) = a.into_split();\n let slot_id = SlotId::new();\n\n pool.add_permit(slot_id, FramedWrite::new(write, JsonCodec::new()));\n\n let permit = pool.try_acquire().unwrap();\n let prediction = Prediction::new(\"test_123\".to_string(), None);\n\n let (_idle_tx, idle_rx) = tokio::sync::oneshot::channel();\n let slot = PredictionSlot::new(prediction, permit, idle_rx);\n assert_eq!(slot.slot_id(), slot_id);\n }\n\n #[tokio::test]\n async fn slot_mark_idle_returns_permit() {\n let pool = PermitPool::new(1);\n\n let (a, _b) = UnixStream::pair().unwrap();\n let (_, write) = a.into_split();\n let slot_id = SlotId::new();\n\n pool.add_permit(slot_id, FramedWrite::new(write, JsonCodec::new()));\n\n {\n let permit = pool.try_acquire().unwrap();\n let prediction = Prediction::new(\"test_123\".to_string(), None);\n let (idle_tx, idle_rx) = tokio::sync::oneshot::channel();\n let slot = PredictionSlot::new(prediction, permit, idle_rx);\n idle_tx\n .send(InactiveSlotIdleToken::new(slot_id).activate())\n .unwrap();\n slot.into_idle().await.unwrap();\n }\n\n assert!(pool.try_acquire().is_some());\n }\n\n #[tokio::test]\n async fn slot_not_idle_orphans_permit() {\n let pool = PermitPool::new(1);\n\n let (a, _b) = UnixStream::pair().unwrap();\n let (_, write) = a.into_split();\n let slot_id = SlotId::new();\n\n pool.add_permit(slot_id, FramedWrite::new(write, JsonCodec::new()));\n\n {\n let permit = pool.try_acquire().unwrap();\n let prediction = Prediction::new(\"test_123\".to_string(), None);\n let (_idle_tx, idle_rx) = tokio::sync::oneshot::channel();\n let _slot = PredictionSlot::new(prediction, permit, idle_rx);\n }\n\n assert!(pool.try_acquire().is_none());\n }\n\n #[tokio::test]\n async fn slot_idle_channel_closed_does_not_return_permit() {\n let pool = PermitPool::new(1);\n\n let (a, _b) = UnixStream::pair().unwrap();\n let (_, write) = a.into_split();\n let slot_id = SlotId::new();\n\n pool.add_permit(slot_id, FramedWrite::new(write, JsonCodec::new()));\n\n let permit = pool.try_acquire().unwrap();\n let prediction = Prediction::new(\"test_123\".to_string(), None);\n let (idle_tx, idle_rx) = tokio::sync::oneshot::channel::();\n let slot = PredictionSlot::new(prediction, permit, idle_rx);\n drop(idle_tx);\n\n let result = slot.into_idle().await;\n assert!(matches!(result, Err(SlotError::IdleTokenReceiveError(_))));\n assert!(pool.try_acquire().is_none());\n }\n}\n" }, { "path": "crates/coglet/src/prediction.rs", "content": "//! Prediction state tracking.\n\nuse std::collections::HashMap;\nuse std::sync::Arc;\nuse std::time::Instant;\n\nuse tokio::sync::Notify;\npub use tokio_util::sync::CancellationToken;\n\nuse crate::bridge::protocol::MetricMode;\nuse crate::webhook::{WebhookEventType, WebhookSender};\n\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum PredictionStatus {\n Starting,\n Processing,\n Succeeded,\n Failed,\n Canceled,\n}\n\nimpl PredictionStatus {\n pub fn is_terminal(&self) -> bool {\n matches!(self, Self::Succeeded | Self::Failed | Self::Canceled)\n }\n\n pub fn as_str(&self) -> &'static str {\n match self {\n Self::Starting => \"starting\",\n Self::Processing => \"processing\",\n Self::Succeeded => \"succeeded\",\n Self::Failed => \"failed\",\n Self::Canceled => \"canceled\",\n }\n }\n}\n\n/// Prediction output - single value or streamed chunks.\n#[derive(Debug, Clone, serde::Serialize)]\n#[serde(untagged)]\npub enum PredictionOutput {\n Single(serde_json::Value),\n Stream(Vec),\n}\n\nimpl PredictionOutput {\n pub fn is_stream(&self) -> bool {\n matches!(self, PredictionOutput::Stream(_))\n }\n\n pub fn into_values(self) -> Vec {\n match self {\n PredictionOutput::Single(v) => vec![v],\n PredictionOutput::Stream(v) => v,\n }\n }\n\n /// Get the final/only output value (last for stream, the value for single).\n pub fn final_value(&self) -> &serde_json::Value {\n match self {\n PredictionOutput::Single(v) => v,\n PredictionOutput::Stream(v) => v.last().unwrap_or(&serde_json::Value::Null),\n }\n }\n}\n\n/// Prediction lifecycle state.\npub struct Prediction {\n id: String,\n cancel_token: CancellationToken,\n started_at: Instant,\n status: PredictionStatus,\n logs: String,\n outputs: Vec,\n output: Option,\n error: Option,\n webhook: Option,\n completion: Arc,\n /// User-emitted metrics. Merged with system metrics (predict_time) in terminal response.\n metrics: HashMap,\n}\n\nimpl Prediction {\n pub fn new(id: String, webhook: Option) -> Self {\n Self {\n id,\n cancel_token: CancellationToken::new(),\n started_at: Instant::now(),\n status: PredictionStatus::Starting,\n logs: String::new(),\n outputs: Vec::new(),\n output: None,\n error: None,\n webhook,\n completion: Arc::new(Notify::new()),\n metrics: HashMap::new(),\n }\n }\n\n pub fn id(&self) -> &str {\n &self.id\n }\n\n pub fn cancel_token(&self) -> CancellationToken {\n self.cancel_token.clone()\n }\n\n pub fn is_canceled(&self) -> bool {\n self.cancel_token.is_cancelled()\n }\n\n pub fn status(&self) -> PredictionStatus {\n self.status\n }\n\n pub fn is_terminal(&self) -> bool {\n self.status.is_terminal()\n }\n\n pub fn set_processing(&mut self) {\n self.status = PredictionStatus::Processing;\n self.fire_webhook(WebhookEventType::Start);\n }\n\n pub fn set_succeeded(&mut self, output: PredictionOutput) {\n if self.status.is_terminal() {\n return;\n }\n self.status = PredictionStatus::Succeeded;\n self.output = Some(output);\n self.fire_terminal_webhook();\n // notify_one stores a permit so a future .notified().await will\n // consume it immediately. notify_waiters only wakes currently-\n // registered waiters and would race with the service task that\n // checks is_terminal() then awaits — the notification can fire\n // in between. There is exactly one waiter per prediction\n // (service.rs predict()), so notify_one is semantically correct.\n self.completion.notify_one();\n }\n\n pub fn set_failed(&mut self, error: String) {\n if self.status.is_terminal() {\n return;\n }\n self.status = PredictionStatus::Failed;\n self.error = Some(error);\n self.fire_terminal_webhook();\n self.completion.notify_one();\n }\n\n pub fn set_canceled(&mut self) {\n if self.status.is_terminal() {\n return;\n }\n self.status = PredictionStatus::Canceled;\n self.fire_terminal_webhook();\n self.completion.notify_one();\n }\n\n pub fn elapsed(&self) -> std::time::Duration {\n self.started_at.elapsed()\n }\n\n pub fn append_log(&mut self, data: &str) {\n self.logs.push_str(data);\n self.fire_webhook(WebhookEventType::Logs);\n }\n\n pub fn logs(&self) -> &str {\n &self.logs\n }\n\n /// Set a user metric with the given accumulation mode.\n ///\n /// - `Replace`: overwrites the value (or deletes if null).\n /// - `Increment`: adds to an existing numeric value. Errors silently if types mismatch.\n /// - `Append`: pushes onto an existing array, creating one if needed.\n ///\n /// Dot-path keys (e.g., \"timing.preprocess\") are resolved into nested objects.\n pub fn set_metric(&mut self, name: String, value: serde_json::Value, mode: MetricMode) {\n // Dot-path resolution: \"a.b.c\" → nested objects\n let parts: Vec<&str> = name.split('.').collect();\n if parts.len() > 1 {\n self.set_metric_dotpath(&parts, value, mode);\n return;\n }\n\n match mode {\n MetricMode::Replace => {\n if value.is_null() {\n self.metrics.remove(&name);\n } else {\n self.metrics.insert(name, value);\n }\n }\n MetricMode::Increment => {\n let entry = self.metrics.entry(name).or_insert(serde_json::json!(0));\n if let (Some(existing), Some(delta)) = (entry.as_f64(), value.as_f64()) {\n // Preserve integer type if both are integers\n if entry.is_i64() && value.is_i64() {\n *entry = serde_json::json!(existing as i64 + delta as i64);\n } else if entry.is_u64() && value.is_u64() {\n *entry = serde_json::json!(existing as u64 + delta as u64);\n } else {\n *entry = serde_json::json!(existing + delta);\n }\n }\n // Non-numeric increment is silently ignored\n }\n MetricMode::Append => {\n let entry = self\n .metrics\n .entry(name)\n .or_insert(serde_json::Value::Array(vec![]));\n if let Some(arr) = entry.as_array_mut() {\n arr.push(value);\n } else {\n // Existing value is not an array — wrap it and append\n let existing = entry.take();\n *entry = serde_json::json!([existing, value]);\n }\n }\n }\n }\n\n /// Resolve a dot-path key into nested objects and apply the metric.\n fn set_metric_dotpath(&mut self, parts: &[&str], value: serde_json::Value, mode: MetricMode) {\n debug_assert!(parts.len() > 1);\n\n let root_key = parts[0].to_string();\n\n // Navigate/create nested structure\n let entry = self\n .metrics\n .entry(root_key)\n .or_insert_with(|| serde_json::json!({}));\n\n let mut current = entry;\n for &part in &parts[1..parts.len() - 1] {\n // Ensure intermediate nodes are objects\n if !current.is_object() {\n *current = serde_json::json!({});\n }\n current = current\n .as_object_mut()\n .unwrap()\n .entry(part)\n .or_insert_with(|| serde_json::json!({}));\n }\n\n let leaf_key = parts[parts.len() - 1];\n\n // Ensure the parent is an object\n if !current.is_object() {\n *current = serde_json::json!({});\n }\n let obj = current.as_object_mut().unwrap();\n\n match mode {\n MetricMode::Replace => {\n if value.is_null() {\n obj.remove(leaf_key);\n } else {\n obj.insert(leaf_key.to_string(), value);\n }\n }\n MetricMode::Increment => {\n let entry = obj.entry(leaf_key).or_insert(serde_json::json!(0));\n if let (Some(existing), Some(delta)) = (entry.as_f64(), value.as_f64()) {\n if entry.is_i64() && value.is_i64() {\n *entry = serde_json::json!(existing as i64 + delta as i64);\n } else if entry.is_u64() && value.is_u64() {\n *entry = serde_json::json!(existing as u64 + delta as u64);\n } else {\n *entry = serde_json::json!(existing + delta);\n }\n }\n }\n MetricMode::Append => {\n let entry = obj\n .entry(leaf_key)\n .or_insert(serde_json::Value::Array(vec![]));\n if let Some(arr) = entry.as_array_mut() {\n arr.push(value);\n } else {\n let existing = entry.take();\n *entry = serde_json::json!([existing, value]);\n }\n }\n }\n }\n\n pub fn metrics(&self) -> &HashMap {\n &self.metrics\n }\n\n pub fn append_output(&mut self, output: serde_json::Value) {\n self.outputs.push(output);\n self.fire_webhook(WebhookEventType::Output);\n }\n\n pub fn outputs(&self) -> &[serde_json::Value] {\n &self.outputs\n }\n\n pub fn take_outputs(&mut self) -> Vec {\n std::mem::take(&mut self.outputs)\n }\n\n pub fn output(&self) -> Option<&PredictionOutput> {\n self.output.as_ref()\n }\n\n pub fn error(&self) -> Option<&str> {\n self.error.as_deref()\n }\n\n pub async fn wait(&self) {\n if self.status.is_terminal() {\n return;\n }\n self.completion.notified().await;\n }\n\n pub fn completion(&self) -> Arc {\n Arc::clone(&self.completion)\n }\n\n /// Take the webhook sender (for sending on drop).\n pub fn take_webhook(&mut self) -> Option {\n self.webhook.take()\n }\n\n /// Fire a non-terminal webhook (throttled, fire-and-forget).\n ///\n /// Builds the current state as a JSON payload and sends it via the\n /// stored WebhookSender. Spawns a tokio task — does not block.\n fn fire_webhook(&self, event: WebhookEventType) {\n if let Some(ref webhook) = self.webhook {\n let payload = self.build_webhook_payload();\n webhook.send(event, &payload);\n }\n }\n\n /// Fire the terminal webhook and consume the WebhookSender.\n ///\n /// Takes ownership of the webhook sender so it can only fire once.\n /// Spawns a tokio task with retry logic for reliability.\n fn fire_terminal_webhook(&mut self) {\n if let Some(webhook) = self.webhook.take() {\n let payload = self.build_webhook_payload();\n tokio::spawn(async move {\n webhook\n .send_terminal(WebhookEventType::Completed, &payload)\n .await;\n });\n }\n }\n\n /// Build a JSON snapshot of the current prediction state.\n ///\n /// This is the single source of truth for prediction JSON. Used by\n /// webhook payloads, GET responses, and terminal responses. Callers\n /// can merge additional fields (e.g. `input`) into the result.\n pub fn build_state_snapshot(&self) -> serde_json::Value {\n let mut payload = serde_json::json!({\n \"id\": self.id,\n \"status\": self.status.as_str(),\n \"logs\": self.logs,\n });\n\n // Include output: use final output if set (terminal), otherwise\n // include accumulated streaming outputs for intermediate states.\n if let Some(ref output) = self.output {\n payload[\"output\"] = serde_json::json!(output);\n } else if !self.outputs.is_empty() {\n payload[\"output\"] = serde_json::json!(self.outputs);\n }\n\n if let Some(ref error) = self.error {\n payload[\"error\"] = serde_json::Value::String(error.clone());\n }\n\n // Include metrics: always include user metrics, add predict_time on terminal\n if !self.metrics.is_empty() || self.status.is_terminal() {\n let mut metrics_obj = serde_json::Map::new();\n for (k, v) in &self.metrics {\n metrics_obj.insert(k.clone(), v.clone());\n }\n if self.status.is_terminal() {\n let predict_time = self.elapsed().as_secs_f64();\n metrics_obj.insert(\"predict_time\".to_string(), serde_json::json!(predict_time));\n }\n payload[\"metrics\"] = serde_json::Value::Object(metrics_obj);\n }\n\n payload\n }\n\n /// Build webhook payload (delegates to build_state_snapshot).\n fn build_webhook_payload(&self) -> serde_json::Value {\n self.build_state_snapshot()\n }\n\n pub fn build_terminal_response(&self) -> serde_json::Value {\n self.build_state_snapshot()\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn status_is_terminal() {\n assert!(!PredictionStatus::Starting.is_terminal());\n assert!(!PredictionStatus::Processing.is_terminal());\n assert!(PredictionStatus::Succeeded.is_terminal());\n assert!(PredictionStatus::Failed.is_terminal());\n assert!(PredictionStatus::Canceled.is_terminal());\n }\n\n #[test]\n fn new_starts_in_starting_status() {\n let pred = Prediction::new(\"test\".to_string(), None);\n assert_eq!(pred.status(), PredictionStatus::Starting);\n assert_eq!(pred.id(), \"test\");\n }\n\n #[test]\n fn set_succeeded() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_succeeded(PredictionOutput::Single(serde_json::json!(\"hello\")));\n assert_eq!(pred.status(), PredictionStatus::Succeeded);\n }\n\n #[test]\n fn set_failed() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_failed(\"something went wrong\".to_string());\n assert_eq!(pred.status(), PredictionStatus::Failed);\n }\n\n #[test]\n fn set_canceled() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_canceled();\n assert_eq!(pred.status(), PredictionStatus::Canceled);\n }\n\n #[test]\n fn cancel_token_works() {\n let pred = Prediction::new(\"test\".to_string(), None);\n let token = pred.cancel_token();\n\n assert!(!pred.is_canceled());\n token.cancel();\n assert!(pred.is_canceled());\n }\n\n #[test]\n fn elapsed_time_increases() {\n let pred = Prediction::new(\"test\".to_string(), None);\n let t1 = pred.elapsed();\n std::thread::sleep(std::time::Duration::from_millis(10));\n let t2 = pred.elapsed();\n assert!(t2 > t1);\n }\n\n #[test]\n fn append_log() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.append_log(\"line 1\\n\");\n pred.append_log(\"line 2\\n\");\n assert_eq!(pred.logs(), \"line 1\\nline 2\\n\");\n }\n\n #[test]\n fn append_output() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.append_output(serde_json::json!(\"chunk1\"));\n pred.append_output(serde_json::json!(\"chunk2\"));\n assert_eq!(pred.outputs().len(), 2);\n }\n\n #[tokio::test]\n async fn wait_returns_immediately_if_terminal() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_succeeded(PredictionOutput::Single(serde_json::json!(\"done\")));\n\n pred.wait().await;\n assert_eq!(pred.status(), PredictionStatus::Succeeded);\n }\n\n #[test]\n fn prediction_output_single() {\n let output = PredictionOutput::Single(serde_json::json!(\"hello\"));\n assert!(!output.is_stream());\n assert_eq!(output.into_values(), vec![serde_json::json!(\"hello\")]);\n }\n\n #[test]\n fn prediction_output_stream() {\n let output = PredictionOutput::Stream(vec![serde_json::json!(\"a\"), serde_json::json!(\"b\")]);\n assert!(output.is_stream());\n }\n\n // ====================================================================\n // Metric tests\n // ====================================================================\n\n #[test]\n fn metric_replace_sets_value() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\"temp\".into(), serde_json::json!(0.7), MetricMode::Replace);\n assert_eq!(pred.metrics()[\"temp\"], serde_json::json!(0.7));\n }\n\n #[test]\n fn metric_replace_overwrites() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\"temp\".into(), serde_json::json!(0.7), MetricMode::Replace);\n pred.set_metric(\"temp\".into(), serde_json::json!(0.9), MetricMode::Replace);\n assert_eq!(pred.metrics()[\"temp\"], serde_json::json!(0.9));\n }\n\n #[test]\n fn metric_replace_null_deletes() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\"temp\".into(), serde_json::json!(0.7), MetricMode::Replace);\n pred.set_metric(\"temp\".into(), serde_json::Value::Null, MetricMode::Replace);\n assert!(!pred.metrics().contains_key(\"temp\"));\n }\n\n #[test]\n fn metric_increment_integers() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\"count\".into(), serde_json::json!(1), MetricMode::Increment);\n pred.set_metric(\"count\".into(), serde_json::json!(3), MetricMode::Increment);\n assert_eq!(pred.metrics()[\"count\"], serde_json::json!(4));\n }\n\n #[test]\n fn metric_increment_floats() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\n \"score\".into(),\n serde_json::json!(1.5),\n MetricMode::Increment,\n );\n pred.set_metric(\n \"score\".into(),\n serde_json::json!(2.5),\n MetricMode::Increment,\n );\n assert_eq!(pred.metrics()[\"score\"], serde_json::json!(4.0));\n }\n\n #[test]\n fn metric_increment_creates_from_zero() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\"count\".into(), serde_json::json!(5), MetricMode::Increment);\n assert_eq!(pred.metrics()[\"count\"], serde_json::json!(5));\n }\n\n #[test]\n fn metric_append_creates_array() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\n \"logprobs\".into(),\n serde_json::json!(-1.2),\n MetricMode::Append,\n );\n pred.set_metric(\n \"logprobs\".into(),\n serde_json::json!(-0.3),\n MetricMode::Append,\n );\n assert_eq!(pred.metrics()[\"logprobs\"], serde_json::json!([-1.2, -0.3]));\n }\n\n #[test]\n fn metric_append_to_non_array_wraps() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\"val\".into(), serde_json::json!(1), MetricMode::Replace);\n pred.set_metric(\"val\".into(), serde_json::json!(2), MetricMode::Append);\n assert_eq!(pred.metrics()[\"val\"], serde_json::json!([1, 2]));\n }\n\n #[test]\n fn metric_dotpath_creates_nested() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\n \"timing.preprocess\".into(),\n serde_json::json!(0.1),\n MetricMode::Replace,\n );\n assert_eq!(\n pred.metrics()[\"timing\"],\n serde_json::json!({\"preprocess\": 0.1})\n );\n }\n\n #[test]\n fn metric_dotpath_deep() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\"a.b.c\".into(), serde_json::json!(42), MetricMode::Replace);\n assert_eq!(pred.metrics()[\"a\"], serde_json::json!({\"b\": {\"c\": 42}}));\n }\n\n #[test]\n fn metric_dotpath_multiple_leaves() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\n \"timing.preprocess\".into(),\n serde_json::json!(0.1),\n MetricMode::Replace,\n );\n pred.set_metric(\n \"timing.inference\".into(),\n serde_json::json!(0.8),\n MetricMode::Replace,\n );\n assert_eq!(\n pred.metrics()[\"timing\"],\n serde_json::json!({\"preprocess\": 0.1, \"inference\": 0.8})\n );\n }\n\n #[test]\n fn metric_dotpath_delete_leaf() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\n \"timing.preprocess\".into(),\n serde_json::json!(0.1),\n MetricMode::Replace,\n );\n pred.set_metric(\n \"timing.preprocess\".into(),\n serde_json::Value::Null,\n MetricMode::Replace,\n );\n // Parent object should still exist but be empty\n assert_eq!(pred.metrics()[\"timing\"], serde_json::json!({}));\n }\n\n #[test]\n fn metric_dotpath_increment() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\n \"stats.tokens\".into(),\n serde_json::json!(10),\n MetricMode::Increment,\n );\n pred.set_metric(\n \"stats.tokens\".into(),\n serde_json::json!(5),\n MetricMode::Increment,\n );\n assert_eq!(pred.metrics()[\"stats\"], serde_json::json!({\"tokens\": 15}));\n }\n\n #[test]\n fn metric_complex_values() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\n \"config\".into(),\n serde_json::json!({\"layers\": 12, \"heads\": 8}),\n MetricMode::Replace,\n );\n pred.set_metric(\n \"scores\".into(),\n serde_json::json!([0.9, 0.8, 0.7]),\n MetricMode::Replace,\n );\n assert_eq!(\n pred.metrics()[\"config\"],\n serde_json::json!({\"layers\": 12, \"heads\": 8})\n );\n assert_eq!(pred.metrics()[\"scores\"], serde_json::json!([0.9, 0.8, 0.7]));\n }\n\n #[test]\n fn terminal_snapshot_merges_metrics_with_predict_time() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_metric(\"temp\".into(), serde_json::json!(0.7), MetricMode::Replace);\n pred.set_metric(\"count\".into(), serde_json::json!(42), MetricMode::Replace);\n pred.set_succeeded(PredictionOutput::Single(serde_json::json!(\"ok\")));\n\n let snapshot = pred.build_state_snapshot();\n let metrics = snapshot[\"metrics\"].as_object().unwrap();\n assert_eq!(metrics[\"temp\"], serde_json::json!(0.7));\n assert_eq!(metrics[\"count\"], serde_json::json!(42));\n assert!(metrics.contains_key(\"predict_time\"));\n }\n\n #[test]\n fn terminal_snapshot_predict_time_overrides_user() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n // User tries to set predict_time - system should override\n pred.set_metric(\n \"predict_time\".into(),\n serde_json::json!(999.0),\n MetricMode::Replace,\n );\n pred.set_succeeded(PredictionOutput::Single(serde_json::json!(\"ok\")));\n\n let snapshot = pred.build_state_snapshot();\n let metrics = snapshot[\"metrics\"].as_object().unwrap();\n // predict_time should be the actual elapsed, not 999.0\n assert_ne!(metrics[\"predict_time\"], serde_json::json!(999.0));\n }\n\n #[test]\n fn terminal_state_guard_set_failed_after_succeeded() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_succeeded(PredictionOutput::Single(serde_json::json!(\"ok\")));\n pred.set_failed(\"Slot dropped unexpectedly\".to_string());\n // Must stay succeeded, not overwritten to failed\n assert_eq!(pred.status(), PredictionStatus::Succeeded);\n assert!(pred.error().is_none());\n }\n\n #[test]\n fn terminal_state_guard_set_succeeded_after_failed() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_failed(\"error\".to_string());\n pred.set_succeeded(PredictionOutput::Single(serde_json::json!(\"late\")));\n assert_eq!(pred.status(), PredictionStatus::Failed);\n assert_eq!(pred.error(), Some(\"error\"));\n }\n\n #[test]\n fn terminal_state_guard_set_canceled_after_succeeded() {\n let mut pred = Prediction::new(\"test\".to_string(), None);\n pred.set_succeeded(PredictionOutput::Single(serde_json::json!(\"done\")));\n pred.set_canceled();\n assert_eq!(pred.status(), PredictionStatus::Succeeded);\n }\n}\n" }, { "path": "crates/coglet/src/predictor.rs", "content": "//! Predictor traits and prediction lifecycle types.\n\nuse std::collections::HashMap;\nuse std::time::{Duration, Instant};\n\npub use crate::prediction::{CancellationToken, PredictionOutput};\n\n/// Result of a completed prediction.\n#[derive(Debug, Clone)]\npub struct PredictionResult {\n pub output: PredictionOutput,\n pub predict_time: Option,\n pub logs: String,\n /// User-emitted metrics from the prediction.\n pub metrics: HashMap,\n}\n\n/// Metrics collected during prediction.\n#[derive(Debug, Clone, Default)]\npub struct PredictionMetrics {\n pub predict_time: Option,\n}\n\n/// RAII guard for prediction lifecycle timing.\npub struct PredictionGuard {\n start_time: Instant,\n metrics: PredictionMetrics,\n cancel_token: CancellationToken,\n}\n\nimpl PredictionGuard {\n pub fn new() -> Self {\n Self {\n start_time: Instant::now(),\n metrics: PredictionMetrics::default(),\n cancel_token: CancellationToken::new(),\n }\n }\n\n pub fn cancel_token(&self) -> CancellationToken {\n self.cancel_token.clone()\n }\n\n pub fn is_cancelled(&self) -> bool {\n self.cancel_token.is_cancelled()\n }\n\n pub fn cancel(&self) {\n self.cancel_token.cancel();\n }\n\n pub fn finish(mut self) -> PredictionMetrics {\n self.metrics.predict_time = Some(self.start_time.elapsed());\n self.metrics\n }\n}\n\nimpl Default for PredictionGuard {\n fn default() -> Self {\n Self::new()\n }\n}\n\n#[derive(Debug, thiserror::Error)]\npub enum PredictionError {\n #[error(\"Prediction failed: {0}\")]\n Failed(String),\n\n #[error(\"Input validation error: {0}\")]\n InvalidInput(String),\n\n #[error(\n \"Setup has not finished yet. Wait until it has finished, or GET /health-check for status.\"\n )]\n NotReady,\n\n #[error(\"Prediction was cancelled\")]\n Cancelled,\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use serde_json::json;\n\n #[test]\n fn prediction_output_single_is_not_stream() {\n let output = PredictionOutput::Single(json!(\"hello\"));\n assert!(!output.is_stream());\n }\n\n #[test]\n fn prediction_output_stream_is_stream() {\n let output = PredictionOutput::Stream(vec![json!(\"a\"), json!(\"b\")]);\n assert!(output.is_stream());\n }\n\n #[test]\n fn prediction_output_serializes_untagged() {\n let single = PredictionOutput::Single(json!(\"hello\"));\n insta::assert_json_snapshot!(\"output_single\", single);\n\n let stream = PredictionOutput::Stream(vec![json!(1), json!(2)]);\n insta::assert_json_snapshot!(\"output_stream\", stream);\n }\n\n #[test]\n fn prediction_guard_tracks_time() {\n let guard = PredictionGuard::new();\n std::thread::sleep(std::time::Duration::from_millis(10));\n let metrics = guard.finish();\n\n assert!(metrics.predict_time.is_some());\n let time = metrics.predict_time.unwrap();\n assert!(time.as_millis() >= 10);\n assert!(time.as_secs() < 1);\n }\n\n #[test]\n fn prediction_error_display() {\n let err = PredictionError::Failed(\"something broke\".to_string());\n assert_eq!(format!(\"{}\", err), \"Prediction failed: something broke\");\n\n let err = PredictionError::InvalidInput(\"bad json\".to_string());\n assert_eq!(format!(\"{}\", err), \"Input validation error: bad json\");\n\n let err = PredictionError::NotReady;\n assert_eq!(\n format!(\"{}\", err),\n \"Setup has not finished yet. Wait until it has finished, or GET /health-check for status.\"\n );\n }\n}\n" }, { "path": "crates/coglet/src/service.rs", "content": "//! PredictionService: Transport-agnostic prediction lifecycle management.\n//!\n//! This service is the single owner of prediction state. It manages:\n//! - Slot management (PermitPool for concurrency control)\n//! - Prediction lifecycle (DashMap of active predictions)\n//! - Cancellation (cancel tokens + orchestrator delegation)\n//! - Health tracking (state, setup result)\n//! - Shutdown coordination (bidirectional)\n//!\n//! Webhooks are fired from Prediction mutation methods (set_processing,\n//! set_succeeded, etc.) where the real state lives — no dual state tracking.\n\nuse std::sync::{Arc, Mutex as StdMutex};\n\nuse dashmap::DashMap;\nuse tokio::sync::{RwLock, watch};\n\nuse crate::bridge::protocol::{MAX_INLINE_IPC_SIZE, SlotRequest};\nuse crate::health::{Health, SetupResult};\nuse crate::input_validation::InputValidator;\nuse crate::orchestrator::{HealthcheckResult, Orchestrator};\nuse crate::permit::{PermitPool, PredictionSlot, UnregisteredPredictionSlot};\nuse crate::prediction::{CancellationToken, Prediction, PredictionStatus};\nuse crate::predictor::{PredictionError, PredictionOutput, PredictionResult};\nuse crate::version::VersionInfo;\nuse crate::webhook::WebhookSender;\n\n/// Try to lock a prediction mutex. On poison, fail the prediction and return None.\nfn try_lock_prediction(\n pred: &Arc>,\n) -> Option> {\n match pred.lock() {\n Ok(guard) => Some(guard),\n Err(poisoned) => {\n tracing::error!(\"Prediction mutex poisoned - failing prediction\");\n let mut guard = poisoned.into_inner();\n if !guard.is_terminal() {\n guard.set_failed(\"Internal error: mutex poisoned\".to_string());\n }\n None\n }\n }\n}\n\n#[derive(Debug, thiserror::Error)]\npub enum CreatePredictionError {\n #[error(\"Service not ready\")]\n NotReady,\n #[error(\"At capacity (no slots available)\")]\n AtCapacity,\n}\n\n/// Snapshot of service health for transports to query.\n#[derive(Debug, Clone)]\npub struct HealthSnapshot {\n pub state: Health,\n pub available_slots: usize,\n pub total_slots: usize,\n pub setup_result: Option,\n pub version: VersionInfo,\n}\n\nimpl HealthSnapshot {\n pub fn is_ready(&self) -> bool {\n self.state == Health::Ready\n }\n\n /// BUSY state: ready but all slots in use.\n pub fn is_busy(&self) -> bool {\n self.state == Health::Ready && self.available_slots == 0\n }\n}\n\n/// Entry in the predictions DashMap.\n///\n/// Holds the real prediction (via Arc), cancel token, and input\n/// (for API responses — Prediction doesn't store input).\nstruct PredictionEntry {\n prediction: Arc>,\n cancel_token: CancellationToken,\n input: serde_json::Value,\n}\n\n/// Handle to a submitted prediction for cancellation on disconnect.\npub struct PredictionHandle {\n id: String,\n cancel_token: CancellationToken,\n}\n\nimpl PredictionHandle {\n pub fn id(&self) -> &str {\n &self.id\n }\n\n pub fn cancel_token(&self) -> CancellationToken {\n self.cancel_token.clone()\n }\n\n /// Create a guard that cancels on drop (for sync predictions).\n ///\n /// On drop (e.g. HTTP connection closed), the guard calls\n /// `service.cancel(id)` which fires the CancellationToken AND\n /// delegates to the orchestrator to cancel the worker subprocess.\n pub fn sync_guard(&self, service: Arc) -> SyncPredictionGuard {\n SyncPredictionGuard::new(self.id.clone(), service)\n }\n}\n\n/// Guard for sync predictions - cancels on drop unless disarmed.\n///\n/// When the HTTP connection drops (client disconnect), axum drops the\n/// response future which drops this guard. The guard calls\n/// `service.cancel(id)` to trigger both the CancellationToken\n/// (Rust-side observers) and the orchestrator (worker subprocess cancel).\npub struct SyncPredictionGuard {\n prediction_id: Option,\n service: Arc,\n}\n\nimpl SyncPredictionGuard {\n pub fn new(prediction_id: String, service: Arc) -> Self {\n Self {\n prediction_id: Some(prediction_id),\n service,\n }\n }\n\n pub fn disarm(&mut self) {\n self.prediction_id = None;\n }\n}\n\nimpl Drop for SyncPredictionGuard {\n fn drop(&mut self) {\n if let Some(ref id) = self.prediction_id {\n self.service.cancel(id);\n }\n }\n}\n\n/// Orchestrator runtime state - pool and orchestrator together.\n///\n/// Ensures pool and orchestrator are always set atomically.\npub struct OrchestratorState {\n pub pool: Arc,\n pub orchestrator: Arc,\n}\n\nimpl Clone for OrchestratorState {\n fn clone(&self) -> Self {\n Self {\n pool: Arc::clone(&self.pool),\n orchestrator: Arc::clone(&self.orchestrator),\n }\n }\n}\n\n/// Transport-agnostic prediction service.\n///\n/// Created with `new_no_pool()`, then configured with `set_orchestrator()` once\n/// the worker subprocess is ready.\npub struct PredictionService {\n /// Orchestrator state (pool + handle together).\n orchestrator: RwLock>,\n\n health: RwLock,\n setup_result: RwLock>,\n\n /// Active predictions — single source of truth for prediction state.\n predictions: DashMap,\n\n shutdown_tx: watch::Sender,\n shutdown_rx: watch::Receiver,\n\n version: VersionInfo,\n\n schema: RwLock>,\n input_validator: RwLock>,\n train_validator: RwLock>,\n}\n\nimpl PredictionService {\n /// Create without configuration (for early HTTP start).\n ///\n /// Health check returns STARTING until `set_orchestrator()` is called.\n pub fn new_no_pool() -> Self {\n let (shutdown_tx, shutdown_rx) = watch::channel(false);\n Self {\n orchestrator: RwLock::new(None),\n health: RwLock::new(Health::Unknown),\n setup_result: RwLock::new(None),\n predictions: DashMap::new(),\n shutdown_tx,\n shutdown_rx,\n version: VersionInfo::new(),\n schema: RwLock::new(None),\n input_validator: RwLock::new(None),\n train_validator: RwLock::new(None),\n }\n }\n\n /// Configure orchestrator mode atomically.\n pub async fn set_orchestrator(\n &self,\n pool: Arc,\n orchestrator: Arc,\n ) {\n *self.orchestrator.write().await = Some(OrchestratorState { pool, orchestrator });\n }\n\n pub async fn has_orchestrator(&self) -> bool {\n self.orchestrator.read().await.is_some()\n }\n\n /// Shutdown the orchestrator gracefully.\n ///\n /// Sends a shutdown message to the worker process and waits for it to exit.\n /// If no orchestrator is configured, this is a no-op.\n pub async fn shutdown(&self) {\n if let Some(ref state) = *self.orchestrator.read().await\n && let Err(e) = state.orchestrator.shutdown().await\n {\n tracing::warn!(error = %e, \"Error during orchestrator shutdown\");\n }\n }\n\n /// Set initial health state (for non-Ready states only).\n ///\n /// READY requires an orchestrator, so use `set_health()` after `set_orchestrator()`.\n /// Silently ignores attempts to set READY here.\n pub fn with_health(mut self, health: Health) -> Self {\n if health != Health::Ready {\n self.health = RwLock::new(health);\n }\n self\n }\n\n pub fn with_version(mut self, version: VersionInfo) -> Self {\n self.version = version;\n self\n }\n\n /// Get the runtime version info.\n pub fn version(&self) -> &VersionInfo {\n &self.version\n }\n\n /// Whether the model supports training (has a TrainingInput schema).\n pub async fn supports_training(&self) -> bool {\n self.train_validator.read().await.is_some()\n }\n\n /// Get the permit pool from orchestrator.\n pub async fn pool(&self) -> Option> {\n if let Some(ref state) = *self.orchestrator.read().await {\n Some(Arc::clone(&state.pool))\n } else {\n None\n }\n }\n\n pub async fn health(&self) -> HealthSnapshot {\n let state = *self.health.read().await;\n let setup_result = self.setup_result.read().await.clone();\n let pool = self.pool().await;\n let (available_slots, total_slots) = match pool.as_ref() {\n Some(p) => (p.available(), p.num_slots()),\n None => (0, 0),\n };\n\n tracing::trace!(\n ?state,\n available_slots,\n total_slots,\n setup_status = ?setup_result.as_ref().map(|r| r.status),\n \"Building health snapshot\"\n );\n\n HealthSnapshot {\n state,\n available_slots,\n total_slots,\n setup_result,\n version: self.version.clone(),\n }\n }\n\n /// Set health state. Setting READY requires orchestrator to be configured.\n ///\n /// Silently ignores attempts to set READY without orchestrator.\n pub async fn set_health(&self, health: Health) {\n if health == Health::Ready && self.orchestrator.read().await.is_none() {\n tracing::warn!(\"Attempted to set READY without orchestrator, ignoring\");\n return;\n }\n let previous = *self.health.read().await;\n tracing::debug!(from = ?previous, to = ?health, \"Health state transition\");\n *self.health.write().await = health;\n }\n\n pub async fn set_setup_result(&self, result: SetupResult) {\n tracing::debug!(\n status = ?result.status,\n started_at = %result.started_at,\n completed_at = ?result.completed_at,\n logs_len = result.logs.len(),\n \"Setting setup result\"\n );\n *self.setup_result.write().await = Some(result);\n }\n\n pub async fn set_schema(&self, schema: serde_json::Value) {\n // Compile input validators from the schema components\n let validator = InputValidator::from_openapi_schema(&schema);\n if let Some(v) = &validator {\n tracing::info!(\n \"Input validation enabled ({} required fields)\",\n v.required_count()\n );\n }\n *self.input_validator.write().await = validator;\n\n // Compile a separate validator for training inputs (TrainingInput)\n let train_val = InputValidator::from_openapi_schema_key(&schema, \"TrainingInput\");\n if let Some(v) = &train_val {\n tracing::info!(\n \"Training input validation enabled ({} required fields)\",\n v.required_count()\n );\n }\n *self.train_validator.write().await = train_val;\n\n *self.schema.write().await = Some(schema);\n }\n\n pub async fn schema(&self) -> Option {\n self.schema.read().await.clone()\n }\n\n /// Validate prediction input against the OpenAPI schema.\n ///\n /// Returns Ok(()) if no schema is loaded or if validation passes.\n /// Returns Err with per-field validation errors on failure.\n pub async fn validate_input(\n &self,\n input: &serde_json::Value,\n ) -> Result<(), Vec> {\n let guard = self.input_validator.read().await;\n if let Some(ref validator) = *guard {\n validator.validate(input)\n } else {\n Ok(())\n }\n }\n\n /// Validate training input against the TrainingInput schema.\n ///\n /// Falls back to the predict validator if no training schema is present.\n pub async fn validate_train_input(\n &self,\n input: &serde_json::Value,\n ) -> Result<(), Vec> {\n let guard = self.train_validator.read().await;\n if let Some(ref validator) = *guard {\n return validator.validate(input);\n }\n drop(guard);\n // Fallback: no TrainingInput schema — use predict validator (legacy compat)\n self.validate_input(input).await\n }\n\n /// Run user-defined healthcheck via orchestrator.\n ///\n /// Returns healthy if no orchestrator is configured (not ready yet).\n pub async fn healthcheck(\n &self,\n ) -> Result {\n if let Some(ref state) = *self.orchestrator.read().await {\n tracing::trace!(\"Dispatching healthcheck to orchestrator\");\n let result = state.orchestrator.healthcheck().await;\n tracing::trace!(\n healthy = result.as_ref().map(|r| r.is_healthy()).unwrap_or(false),\n error = ?result.as_ref().ok().and_then(|r| r.error.as_ref()),\n \"Healthcheck result from orchestrator\"\n );\n result\n } else {\n tracing::debug!(\"No orchestrator configured, returning default healthy\");\n Ok(HealthcheckResult::healthy())\n }\n }\n\n /// Submit a new prediction: create Prediction, acquire slot, register in DashMap.\n ///\n /// Returns a PredictionHandle (for cancel-on-disconnect) and the\n /// UnregisteredPredictionSlot (for running the prediction).\n pub async fn submit_prediction(\n &self,\n id: String,\n input: serde_json::Value,\n webhook: Option,\n ) -> Result<(PredictionHandle, UnregisteredPredictionSlot), CreatePredictionError> {\n let health = *self.health.read().await;\n if health != Health::Ready {\n return Err(CreatePredictionError::NotReady);\n }\n\n // Pool must exist if health is Ready\n let pool = self.pool().await;\n let pool = pool.as_ref().ok_or(CreatePredictionError::NotReady)?;\n\n let permit = pool\n .try_acquire()\n .ok_or(CreatePredictionError::AtCapacity)?;\n\n let prediction = Prediction::new(id.clone(), webhook);\n let cancel_token = prediction.cancel_token();\n let (idle_tx, idle_rx) = tokio::sync::oneshot::channel();\n let slot = PredictionSlot::new(prediction, permit, idle_rx);\n let prediction_arc = slot.prediction();\n\n // Register in DashMap — this is the single source of truth\n self.predictions.insert(\n id.clone(),\n PredictionEntry {\n prediction: prediction_arc,\n cancel_token: cancel_token.clone(),\n input,\n },\n );\n\n let handle = PredictionHandle { id, cancel_token };\n\n Ok((handle, UnregisteredPredictionSlot::new(slot, idle_tx)))\n }\n\n /// Check if a prediction with this ID is already in-flight.\n pub fn prediction_exists(&self, id: &str) -> bool {\n self.predictions.contains_key(id)\n }\n\n /// Get a snapshot of prediction state for API responses.\n ///\n /// Locks the real Prediction to read current state — no stale copies.\n /// Adds `input` from the PredictionEntry on top of the shared snapshot.\n pub fn get_prediction_response(&self, id: &str) -> Option {\n let entry = self.predictions.get(id)?;\n let pred = entry.prediction.lock().ok()?;\n\n let mut response = pred.build_state_snapshot();\n response[\"input\"] = entry.input.clone();\n\n Some(response)\n }\n\n /// Run a prediction to completion via orchestrator.\n pub async fn predict(\n &self,\n unregistered_slot: UnregisteredPredictionSlot,\n input: serde_json::Value,\n context: std::collections::HashMap,\n ) -> Result {\n let state = self.orchestrator.read().await.clone();\n let state = state\n .ok_or_else(|| PredictionError::Failed(\"No orchestrator configured\".to_string()))?;\n\n let (idle_tx, mut slot) = unregistered_slot.into_parts();\n let prediction_id = slot.id();\n let slot_id = slot.slot_id();\n\n {\n let prediction = slot.prediction();\n let Some(mut pred) = try_lock_prediction(&prediction) else {\n return Err(PredictionError::Failed(\n \"Prediction mutex poisoned\".to_string(),\n ));\n };\n pred.set_processing();\n }\n\n // Register for response routing in event loop\n let prediction_arc = slot.prediction();\n state\n .orchestrator\n .register_prediction(slot_id, Arc::clone(&prediction_arc), idle_tx)\n .await;\n\n // Create per-prediction dirs for file-based inputs/outputs\n let prediction_dir =\n std::path::PathBuf::from(\"/tmp/coglet/predictions\").join(&prediction_id);\n let output_dir = prediction_dir.join(\"outputs\");\n let input_dir = prediction_dir.join(\"inputs\");\n std::fs::create_dir_all(&output_dir)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to create output dir: {}\", e)))?;\n std::fs::create_dir_all(&input_dir)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to create input dir: {}\", e)))?;\n\n let request = build_slot_request(\n prediction_id.clone(),\n input,\n output_dir\n .to_str()\n .expect(\"output dir path is valid UTF-8\")\n .to_string(),\n &input_dir,\n context,\n )\n .map_err(|e| PredictionError::Failed(format!(\"Failed to build slot request: {}\", e)))?;\n\n // permit_mut returns None if permit isn't InUse (shouldn't happen here)\n let permit = slot\n .permit_mut()\n .ok_or_else(|| PredictionError::Failed(\"Permit not in use\".to_string()))?;\n\n if let Err(e) = permit.send(request).await {\n tracing::error!(%slot_id, error = %e, \"Failed to send prediction request\");\n // Broken socket means the slot is dead — poison it at the pool level.\n state.pool.poison(slot_id);\n if let Some(mut pred) = try_lock_prediction(&prediction_arc) {\n pred.set_failed(format!(\"Failed to send request: {}\", e));\n }\n return Err(PredictionError::Failed(format!(\n \"Failed to send request: {}\",\n e\n )));\n }\n\n // Wait for prediction to complete\n // Check if already terminal first to avoid race with fast completions\n let (already_terminal, completion) = {\n let Some(pred) = try_lock_prediction(&prediction_arc) else {\n return Err(PredictionError::Failed(\n \"Prediction mutex poisoned\".to_string(),\n ));\n };\n (pred.is_terminal(), pred.completion())\n };\n if !already_terminal {\n completion.notified().await;\n }\n\n let (status, output, error, logs, predict_time, metrics) = {\n let Some(pred) = try_lock_prediction(&prediction_arc) else {\n return Err(PredictionError::Failed(\n \"Prediction mutex poisoned\".to_string(),\n ));\n };\n (\n pred.status(),\n pred.output().cloned(),\n pred.error().map(|s| s.to_string()),\n pred.logs().to_string(),\n pred.elapsed(),\n pred.metrics().clone(),\n )\n };\n\n // If `into_idle()` fails, it does not necessarily mean the prediction failed,\n // so we return the result if available, but log the error and poison the slot to prevent reuse.\n // This is performed asynchronously to avoid blocking the prediction response to the caller.\n tokio::spawn(async move {\n if let Err(e) = slot.into_idle().await {\n tracing::error!(%slot_id, error = %e, \"Failed to transition slot to idle, poisoning slot\");\n state.pool.poison(slot_id);\n }\n });\n\n match status {\n PredictionStatus::Succeeded => Ok(PredictionResult {\n output: output.unwrap_or(PredictionOutput::Single(serde_json::Value::Null)),\n predict_time: Some(predict_time),\n logs,\n metrics,\n }),\n PredictionStatus::Failed => Err(PredictionError::Failed(\n error.unwrap_or_else(|| \"Unknown error\".to_string()),\n )),\n PredictionStatus::Canceled => Err(PredictionError::Cancelled),\n _ => Err(PredictionError::Failed(format!(\n \"Prediction ended in unexpected state: {:?}\",\n status\n ))),\n }\n }\n\n /// Cancel a prediction by ID. Returns true if found and cancelled.\n ///\n /// Fires the CancellationToken (for Rust-side observers like upload tasks)\n /// and delegates to the orchestrator to send `ControlRequest::Cancel` to the worker.\n pub fn cancel(&self, id: &str) -> bool {\n if let Some(entry) = self.predictions.get(id) {\n entry.cancel_token.cancel();\n\n // Delegate to orchestrator to actually cancel the worker-side prediction.\n // This must be non-blocking since cancel() is sync, so we spawn a task.\n let id_owned = id.to_string();\n let orchestrator = self\n .orchestrator\n .try_read()\n .ok()\n .and_then(|guard| guard.as_ref().map(|s| Arc::clone(&s.orchestrator)));\n if let Some(orch) = orchestrator {\n tokio::spawn(async move {\n if let Err(e) = orch.cancel_by_prediction_id(&id_owned).await {\n tracing::error!(\n prediction_id = %id_owned,\n error = %e,\n \"Failed to send cancel to orchestrator\"\n );\n }\n });\n }\n true\n } else {\n false\n }\n }\n\n /// Remove a prediction from the DashMap after completion.\n pub fn remove_prediction(&self, id: &str) {\n self.predictions.remove(id);\n }\n\n pub fn trigger_shutdown(&self) {\n let _ = self.shutdown_tx.send(true);\n }\n\n pub fn shutdown_rx(&self) -> watch::Receiver {\n self.shutdown_rx.clone()\n }\n}\n\n/// Build a `SlotRequest::Predict`, spilling the input to disk if it exceeds\n/// `MAX_INLINE_IPC_SIZE`. This prevents IPC frame overflow on the slot socket.\n///\n/// NOTE: The input is serialized here to check its size against the threshold.\n/// For the inline path the original `Value` is kept and will be serialized again\n/// by `JsonCodec` — a double-serialize trade-off we accept to keep the codec\n/// generic. The spill path writes the pre-serialized bytes directly.\nfn build_slot_request(\n id: String,\n input: serde_json::Value,\n output_dir: String,\n input_dir: &std::path::Path,\n context: std::collections::HashMap,\n) -> std::io::Result {\n let serialized = serde_json::to_vec(&input)\n .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;\n\n if serialized.len() > MAX_INLINE_IPC_SIZE {\n let path = input_dir.join(format!(\"spill_{}.json\", uuid::Uuid::new_v4()));\n std::fs::write(&path, &serialized)?;\n let input_file = path\n .to_str()\n .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, \"non-UTF-8 path\"))?\n .to_string();\n Ok(SlotRequest::Predict {\n id,\n input: None,\n input_file: Some(input_file),\n output_dir,\n context,\n })\n } else {\n Ok(SlotRequest::Predict {\n id,\n input: Some(input),\n input_file: None,\n output_dir,\n context,\n })\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use crate::bridge::protocol::SlotId;\n use crate::permit::{InactiveSlotIdleToken, SlotIdleToken};\n use std::sync::atomic::{AtomicUsize, Ordering};\n use std::time::Duration;\n\n /// Mock orchestrator that immediately completes predictions.\n struct MockOrchestrator {\n register_count: AtomicUsize,\n complete_immediately: bool,\n send_idle_ack: bool,\n }\n\n impl MockOrchestrator {\n fn new() -> Self {\n Self {\n register_count: AtomicUsize::new(0),\n complete_immediately: true,\n send_idle_ack: false,\n }\n }\n\n fn register_count(&self) -> usize {\n self.register_count.load(Ordering::SeqCst)\n }\n\n fn with_idle_ack(mut self) -> Self {\n self.send_idle_ack = true;\n self\n }\n }\n\n #[async_trait::async_trait]\n impl Orchestrator for MockOrchestrator {\n async fn register_prediction(\n &self,\n slot_id: SlotId,\n prediction: Arc>,\n idle_sender: tokio::sync::oneshot::Sender,\n ) {\n self.register_count.fetch_add(1, Ordering::SeqCst);\n if self.complete_immediately {\n let mut pred = prediction.lock().unwrap();\n pred.set_succeeded(crate::PredictionOutput::Single(serde_json::json!(\n \"mock result\"\n )));\n }\n if self.send_idle_ack {\n let _ = idle_sender.send(InactiveSlotIdleToken::new(slot_id).activate());\n }\n }\n\n async fn cancel_by_prediction_id(\n &self,\n _prediction_id: &str,\n ) -> Result<(), crate::orchestrator::OrchestratorError> {\n Ok(())\n }\n\n async fn healthcheck(\n &self,\n ) -> Result {\n Ok(HealthcheckResult::healthy())\n }\n\n async fn shutdown(&self) -> Result<(), crate::orchestrator::OrchestratorError> {\n Ok(())\n }\n }\n\n async fn create_test_pool(num_slots: usize) -> Arc {\n use crate::bridge::codec::JsonCodec;\n use crate::bridge::protocol::SlotRequest;\n use futures::StreamExt;\n use tokio::net::UnixStream;\n\n let pool = Arc::new(PermitPool::new(num_slots));\n for _ in 0..num_slots {\n let (a, b) = UnixStream::pair().unwrap();\n let (_read_a, write_a) = a.into_split();\n let (read_b, _write_b) = b.into_split();\n\n // Spawn a task to consume messages from the socket (prevents broken pipe)\n let mut reader =\n tokio_util::codec::FramedRead::new(read_b, JsonCodec::::new());\n tokio::spawn(async move { while reader.next().await.is_some() {} });\n\n let writer =\n tokio_util::codec::FramedWrite::new(write_a, JsonCodec::::new());\n pool.add_permit(SlotId::new(), writer);\n }\n pool\n }\n\n async fn create_test_pool_with_slots(num_slots: usize) -> (Arc, Vec) {\n use crate::bridge::codec::JsonCodec;\n use crate::bridge::protocol::SlotRequest;\n use futures::StreamExt;\n use tokio::net::UnixStream;\n\n let pool = Arc::new(PermitPool::new(num_slots));\n let mut slot_ids = Vec::with_capacity(num_slots);\n for _ in 0..num_slots {\n let (a, b) = UnixStream::pair().unwrap();\n let (_read_a, write_a) = a.into_split();\n let (read_b, _write_b) = b.into_split();\n\n let mut reader =\n tokio_util::codec::FramedRead::new(read_b, JsonCodec::::new());\n tokio::spawn(async move { while reader.next().await.is_some() {} });\n\n let writer =\n tokio_util::codec::FramedWrite::new(write_a, JsonCodec::::new());\n let slot_id = SlotId::new();\n pool.add_permit(slot_id, writer);\n slot_ids.push(slot_id);\n }\n (pool, slot_ids)\n }\n\n async fn create_broken_test_pool() -> (Arc, SlotId) {\n use crate::bridge::codec::JsonCodec;\n use crate::bridge::protocol::SlotRequest;\n use tokio::net::UnixStream;\n\n let pool = Arc::new(PermitPool::new(1));\n let (a, b) = UnixStream::pair().unwrap();\n let (_read_a, write_a) = a.into_split();\n drop(b);\n\n let writer = tokio_util::codec::FramedWrite::new(write_a, JsonCodec::::new());\n let slot_id = SlotId::new();\n pool.add_permit(slot_id, writer);\n (pool, slot_id)\n }\n\n #[tokio::test]\n async fn service_new_no_pool_works() {\n let svc = PredictionService::new_no_pool();\n let health = svc.health().await;\n\n assert_eq!(health.state, Health::Unknown);\n assert_eq!(health.total_slots, 0);\n assert_eq!(health.available_slots, 0);\n assert!(svc.pool().await.is_none());\n }\n\n #[tokio::test]\n async fn service_no_pool_initially() {\n let svc = PredictionService::new_no_pool();\n\n assert!(svc.pool().await.is_none());\n assert!(!svc.has_orchestrator().await);\n }\n\n #[tokio::test]\n async fn shutdown_signal_works() {\n let svc = PredictionService::new_no_pool();\n let mut rx = svc.shutdown_rx();\n\n assert!(!*rx.borrow());\n\n svc.trigger_shutdown();\n rx.changed().await.unwrap();\n\n assert!(*rx.borrow());\n }\n\n #[tokio::test]\n async fn submit_fails_when_not_ready() {\n let svc = PredictionService::new_no_pool();\n\n let result = svc\n .submit_prediction(\"test\".to_string(), serde_json::json!({}), None)\n .await;\n assert!(matches!(result, Err(CreatePredictionError::NotReady)));\n }\n\n #[tokio::test]\n async fn cannot_set_ready_without_orchestrator() {\n let svc = PredictionService::new_no_pool();\n\n // with_health silently ignores READY\n let svc2 = PredictionService::new_no_pool().with_health(Health::Ready);\n assert_eq!(svc2.health().await.state, Health::Unknown);\n\n // set_health also ignores READY without orchestrator\n svc.set_health(Health::Ready).await;\n assert_eq!(svc.health().await.state, Health::Unknown);\n }\n\n #[tokio::test]\n async fn set_orchestrator_enables_ready_health() {\n let svc = PredictionService::new_no_pool();\n let pool = create_test_pool(2).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(pool, orchestrator).await;\n assert!(svc.has_orchestrator().await);\n\n // Now we can set READY\n svc.set_health(Health::Ready).await;\n let health = svc.health().await;\n assert_eq!(health.state, Health::Ready);\n assert_eq!(health.total_slots, 2);\n assert_eq!(health.available_slots, 2);\n }\n\n #[tokio::test]\n async fn submit_prediction_succeeds_when_ready() {\n let svc = PredictionService::new_no_pool();\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(pool, orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (handle, _slot) = svc\n .submit_prediction(\"test-1\".to_string(), serde_json::json!({}), None)\n .await\n .unwrap();\n\n assert_eq!(handle.id(), \"test-1\");\n assert!(svc.prediction_exists(\"test-1\"));\n }\n\n #[tokio::test]\n async fn submit_returns_at_capacity_when_no_slots() {\n let svc = PredictionService::new_no_pool();\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(pool, orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n // First prediction takes the only slot\n let (_handle1, _slot1) = svc\n .submit_prediction(\"test-1\".to_string(), serde_json::json!({}), None)\n .await\n .unwrap();\n\n // Second should fail with AtCapacity\n let result = svc\n .submit_prediction(\"test-2\".to_string(), serde_json::json!({}), None)\n .await;\n assert!(matches!(result, Err(CreatePredictionError::AtCapacity)));\n }\n\n #[tokio::test]\n async fn predict_calls_orchestrator_register() {\n let svc = PredictionService::new_no_pool();\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n let orch_ref = Arc::clone(&orchestrator);\n\n svc.set_orchestrator(pool, orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (_handle, slot) = svc\n .submit_prediction(\n \"test-1\".to_string(),\n serde_json::json!({\"prompt\": \"hello\"}),\n None,\n )\n .await\n .unwrap();\n\n let result = svc\n .predict(\n slot,\n serde_json::json!({\"prompt\": \"hello\"}),\n Default::default(),\n )\n .await;\n\n // MockOrchestrator completes immediately with success\n assert!(result.is_ok(), \"predict failed: {:?}\", result.err());\n assert_eq!(orch_ref.register_count(), 1);\n }\n\n #[tokio::test]\n async fn health_shows_busy_when_all_slots_used() {\n let svc = PredictionService::new_no_pool();\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(pool, orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n // Before acquiring slot\n let health = svc.health().await;\n assert!(!health.is_busy());\n assert_eq!(health.available_slots, 1);\n\n // After acquiring slot\n let (_handle, _slot) = svc\n .submit_prediction(\"test-1\".to_string(), serde_json::json!({}), None)\n .await\n .unwrap();\n let health = svc.health().await;\n assert!(health.is_busy());\n assert_eq!(health.available_slots, 0);\n }\n\n #[tokio::test]\n async fn predict_idle_channel_closed_poison_slot_async() {\n let svc = PredictionService::new_no_pool();\n let (pool, slot_ids) = create_test_pool_with_slots(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n let slot_id = slot_ids[0];\n\n svc.set_orchestrator(Arc::clone(&pool), orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (_handle, slot) = svc\n .submit_prediction(\n \"test-1\".to_string(),\n serde_json::json!({\"prompt\": \"hello\"}),\n None,\n )\n .await\n .unwrap();\n\n let result = svc\n .predict(\n slot,\n serde_json::json!({\"prompt\": \"hello\"}),\n Default::default(),\n )\n .await;\n assert!(result.is_ok(), \"predict failed: {:?}\", result.err());\n\n tokio::time::timeout(Duration::from_secs(1), async {\n loop {\n if pool.is_poisoned(slot_id) {\n break;\n }\n tokio::time::sleep(Duration::from_millis(10)).await;\n }\n })\n .await\n .expect(\"slot was not poisoned after idle token channel closed\");\n }\n\n #[tokio::test]\n async fn predict_idle_ack_returns_capacity_async() {\n let svc = PredictionService::new_no_pool();\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new().with_idle_ack());\n\n svc.set_orchestrator(Arc::clone(&pool), orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (_handle, slot) = svc\n .submit_prediction(\n \"test-1\".to_string(),\n serde_json::json!({\"prompt\": \"hello\"}),\n None,\n )\n .await\n .unwrap();\n\n let result = svc\n .predict(\n slot,\n serde_json::json!({\"prompt\": \"hello\"}),\n Default::default(),\n )\n .await;\n assert!(result.is_ok(), \"predict failed: {:?}\", result.err());\n\n tokio::time::timeout(Duration::from_secs(1), async {\n loop {\n if pool.available() == 1 {\n break;\n }\n tokio::time::sleep(Duration::from_millis(10)).await;\n }\n })\n .await\n .expect(\"slot capacity was not returned after idle acknowledgement\");\n }\n\n #[tokio::test]\n async fn predict_send_failure_poison_slot() {\n let svc = PredictionService::new_no_pool();\n let (pool, slot_id) = create_broken_test_pool().await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(Arc::clone(&pool), orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (_handle, slot) = svc\n .submit_prediction(\n \"test-1\".to_string(),\n serde_json::json!({\"prompt\": \"hello\"}),\n None,\n )\n .await\n .unwrap();\n\n let result = svc\n .predict(\n slot,\n serde_json::json!({\"prompt\": \"hello\"}),\n Default::default(),\n )\n .await;\n assert!(matches!(result, Err(PredictionError::Failed(_))));\n assert!(pool.is_poisoned(slot_id));\n assert!(pool.try_acquire().is_none());\n }\n\n #[tokio::test]\n async fn cancel_prediction_works() {\n let svc = PredictionService::new_no_pool();\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(pool, orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (handle, _slot) = svc\n .submit_prediction(\"test-cancel\".to_string(), serde_json::json!({}), None)\n .await\n .unwrap();\n\n let cancel_token = handle.cancel_token();\n let cancelled = svc.cancel(\"test-cancel\");\n assert!(cancelled);\n assert!(cancel_token.is_cancelled());\n }\n\n #[tokio::test]\n async fn cancel_nonexistent_returns_false() {\n let svc = PredictionService::new_no_pool();\n assert!(!svc.cancel(\"nonexistent\"));\n }\n\n #[tokio::test]\n async fn sync_guard_cancels_on_drop() {\n let svc = Arc::new(PredictionService::new_no_pool());\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(pool, orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (handle, _slot) = svc\n .submit_prediction(\"test-guard\".to_string(), serde_json::json!({}), None)\n .await\n .unwrap();\n\n let cancel_token = handle.cancel_token();\n\n {\n let _guard = handle.sync_guard(Arc::clone(&svc));\n assert!(!cancel_token.is_cancelled());\n }\n\n assert!(cancel_token.is_cancelled());\n }\n\n #[tokio::test]\n async fn sync_guard_disarm_prevents_cancel() {\n let svc = Arc::new(PredictionService::new_no_pool());\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(pool, orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (handle, _slot) = svc\n .submit_prediction(\"test-disarm\".to_string(), serde_json::json!({}), None)\n .await\n .unwrap();\n\n let cancel_token = handle.cancel_token();\n\n {\n let mut guard = handle.sync_guard(Arc::clone(&svc));\n guard.disarm();\n }\n\n assert!(!cancel_token.is_cancelled());\n }\n\n #[tokio::test]\n async fn remove_prediction_cleans_up() {\n let svc = PredictionService::new_no_pool();\n let pool = create_test_pool(1).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n\n svc.set_orchestrator(pool, orchestrator).await;\n svc.set_health(Health::Ready).await;\n\n let (_handle, _slot) = svc\n .submit_prediction(\"test-remove\".to_string(), serde_json::json!({}), None)\n .await\n .unwrap();\n\n assert!(svc.prediction_exists(\"test-remove\"));\n svc.remove_prediction(\"test-remove\");\n assert!(!svc.prediction_exists(\"test-remove\"));\n }\n\n #[test]\n fn build_slot_request_small_input_inline() {\n let dir = tempfile::tempdir().unwrap();\n let input = serde_json::json!({\"text\": \"hello\"});\n let req = build_slot_request(\n \"p1\".into(),\n input.clone(),\n \"/tmp/out\".into(),\n dir.path(),\n Default::default(),\n )\n .unwrap();\n\n match req {\n SlotRequest::Predict {\n id,\n input: Some(v),\n input_file: None,\n output_dir,\n ..\n } => {\n assert_eq!(id, \"p1\");\n assert_eq!(v, input);\n assert_eq!(output_dir, \"/tmp/out\");\n }\n _ => panic!(\"expected inline input\"),\n }\n }\n\n #[test]\n fn build_slot_request_large_input_spills() {\n let dir = tempfile::tempdir().unwrap();\n // Create an input larger than 6 MiB\n let big = \"x\".repeat(7 * 1024 * 1024);\n let input = serde_json::json!({\"data\": big});\n let req = build_slot_request(\n \"p2\".into(),\n input.clone(),\n \"/tmp/out\".into(),\n dir.path(),\n Default::default(),\n )\n .unwrap();\n\n match req {\n SlotRequest::Predict {\n id,\n input: None,\n input_file: Some(ref path),\n output_dir,\n ..\n } => {\n assert_eq!(id, \"p2\");\n assert_eq!(output_dir, \"/tmp/out\");\n // Spill file should exist on disk\n assert!(std::path::Path::new(path).exists());\n // Content should be valid JSON matching the original input\n let bytes = std::fs::read(path).unwrap();\n let roundtrip: serde_json::Value = serde_json::from_slice(&bytes).unwrap();\n assert_eq!(roundtrip, input);\n }\n _ => panic!(\"expected file-backed input\"),\n }\n }\n\n #[test]\n fn build_slot_request_roundtrip() {\n let dir = tempfile::tempdir().unwrap();\n let big = \"y\".repeat(7 * 1024 * 1024);\n let input = serde_json::json!({\"payload\": big});\n let req = build_slot_request(\n \"p3\".into(),\n input.clone(),\n \"/tmp/out\".into(),\n dir.path(),\n Default::default(),\n )\n .unwrap();\n\n // Rehydrate and verify we get back the same value\n let (id, rehydrated, output_dir, _context) = req.rehydrate_input().unwrap();\n assert_eq!(id, \"p3\");\n assert_eq!(rehydrated, input);\n assert_eq!(output_dir, \"/tmp/out\");\n }\n}\n" }, { "path": "crates/coglet/src/setup_log_accumulator.rs", "content": "//! Tracing layer that accumulates all logs from coglet server during setup.\n//!\n//! Captures every tracing event from the moment the server starts until setup completes.\n//! This includes:\n//! - Initial server startup logs (\"coglet \")\n//! - Orchestrator logs (\"Spawning worker subprocess\")\n//! - Re-emitted worker logs (via emit_worker_log)\n//! - Transport logs, codec warnings, everything\n//!\n//! Uses unbounded mpsc channel for lock-free accumulation.\n\nuse tokio::sync::mpsc;\nuse tracing::Subscriber;\nuse tracing_subscriber::layer::{Context, Layer};\n\npub struct SetupLogAccumulator {\n tx: mpsc::UnboundedSender,\n}\n\nimpl SetupLogAccumulator {\n pub fn new(tx: mpsc::UnboundedSender) -> Self {\n Self { tx }\n }\n}\n\nimpl Layer for SetupLogAccumulator\nwhere\n S: Subscriber,\n{\n fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {\n if self.tx.is_closed() {\n return;\n }\n\n let metadata = event.metadata();\n let target = metadata.target();\n\n let mut visitor = MessageVisitor::default();\n event.record(&mut visitor);\n\n let log_line = format!(\"[{}] {}\", target, visitor.message);\n let _ = self.tx.send(log_line);\n }\n}\n\n#[derive(Default)]\nstruct MessageVisitor {\n message: String,\n}\n\nimpl tracing::field::Visit for MessageVisitor {\n fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {\n if field.name() == \"message\" {\n self.message = format!(\"{:?}\", value);\n if self.message.starts_with('\"') && self.message.ends_with('\"') {\n self.message = self.message[1..self.message.len() - 1].to_string();\n }\n }\n }\n\n fn record_str(&mut self, field: &tracing::field::Field, value: &str) {\n if field.name() == \"message\" {\n self.message = value.to_string();\n }\n }\n}\n\npub fn drain_accumulated_logs(rx: &mut mpsc::UnboundedReceiver) -> String {\n let mut lines = Vec::new();\n while let Ok(line) = rx.try_recv() {\n lines.push(line);\n }\n\n if lines.is_empty() {\n String::new()\n } else {\n let mut result = lines.join(\"\\n\");\n result.push('\\n');\n result\n }\n}\n" }, { "path": "crates/coglet/src/snapshots/coglet__health__tests__health_all_variants.snap", "content": "---\nsource: coglet/src/health.rs\nexpression: \"[Health::Unknown, Health::Starting, Health::Ready, Health::Busy,\\nHealth::SetupFailed, Health::Defunct,]\"\n---\n[\n \"UNKNOWN\",\n \"STARTING\",\n \"READY\",\n \"BUSY\",\n \"SETUP_FAILED\",\n \"DEFUNCT\"\n]\n" }, { "path": "crates/coglet/src/snapshots/coglet__health__tests__health_response_all_variants.snap", "content": "---\nsource: coglet/src/health.rs\nexpression: \"[HealthResponse::Unknown, HealthResponse::Starting, HealthResponse::Ready,\\nHealthResponse::Busy, HealthResponse::SetupFailed, HealthResponse::Defunct,\\nHealthResponse::Unhealthy,]\"\n---\n[\n \"UNKNOWN\",\n \"STARTING\",\n \"READY\",\n \"BUSY\",\n \"SETUP_FAILED\",\n \"DEFUNCT\",\n \"UNHEALTHY\"\n]\n" }, { "path": "crates/coglet/src/snapshots/coglet__health__tests__setup_status_all_variants.snap", "content": "---\nsource: coglet/src/health.rs\nexpression: \"[SetupStatus::Starting, SetupStatus::Succeeded, SetupStatus::Failed,]\"\n---\n[\n \"starting\",\n \"succeeded\",\n \"failed\"\n]\n" }, { "path": "crates/coglet/src/snapshots/coglet__predictor__tests__output_single.snap", "content": "---\nsource: coglet/src/predictor.rs\nexpression: single\n---\n\"hello\"\n" }, { "path": "crates/coglet/src/snapshots/coglet__predictor__tests__output_stream.snap", "content": "---\nsource: coglet/src/predictor.rs\nexpression: stream\n---\n[\n 1,\n 2\n]\n" }, { "path": "crates/coglet/src/snapshots/coglet__version__tests__version_full.snap", "content": "---\nsource: coglet/src/version.rs\nexpression: info\n---\n{\n \"coglet\": \"0.1.0\",\n \"git_sha\": \"abc1234-dirty\",\n \"build_time\": \"2026-03-12T18:00:00Z\",\n \"python_sdk\": \"0.9.0\",\n \"python\": \"3.11.0\"\n}\n" }, { "path": "crates/coglet/src/snapshots/coglet__version__tests__version_minimal.snap", "content": "---\nsource: coglet/src/version.rs\nexpression: info\n---\n{\n \"coglet\": \"0.1.0\"\n}\n" }, { "path": "crates/coglet/src/transport/http/mod.rs", "content": "//! HTTP transport for coglet using axum.\n\nmod routes;\nmod server;\n\npub use server::{ServerConfig, serve};\n" }, { "path": "crates/coglet/src/transport/http/routes.rs", "content": "//! HTTP route handlers.\n\nuse std::sync::Arc;\n\nuse axum::{\n Router,\n extract::{DefaultBodyLimit, Path, State},\n http::{HeaderMap, StatusCode},\n response::{IntoResponse, Json},\n routing::{get, post, put},\n};\nuse serde::{Deserialize, Serialize};\n\n#[cfg(test)]\nuse crate::health::Health;\nuse crate::health::{HealthResponse, SetupResult};\nuse crate::predictor::PredictionError;\nuse crate::service::{CreatePredictionError, HealthSnapshot, PredictionService};\nuse crate::version::VersionInfo;\nuse crate::webhook::{TraceContext, WebhookConfig, WebhookEventType, WebhookSender};\n\n#[derive(Debug, Serialize)]\npub struct HealthCheckResponse {\n pub status: HealthResponse,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub setup: Option,\n pub version: VersionInfo,\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub user_healthcheck_error: Option,\n}\n\nimpl HealthCheckResponse {\n pub fn from_snapshot(snapshot: HealthSnapshot, user_healthcheck_error: Option) -> Self {\n // Determine response status\n let status = if user_healthcheck_error.is_some() {\n HealthResponse::Unhealthy\n } else if snapshot.is_busy() {\n HealthResponse::Busy\n } else {\n snapshot.state.into()\n };\n\n Self {\n status,\n setup: snapshot.setup_result,\n version: snapshot.version,\n user_healthcheck_error,\n }\n }\n}\n\n#[derive(Debug, Deserialize)]\npub struct PredictionRequest {\n pub id: Option,\n #[serde(\n default = \"default_empty_input\",\n deserialize_with = \"deserialize_input\"\n )]\n pub input: serde_json::Value,\n /// Per-prediction context made available to predictors via `current_scope().context`.\n #[serde(default)]\n pub context: std::collections::HashMap,\n pub webhook: Option,\n #[serde(default = \"default_webhook_events_filter\")]\n pub webhook_events_filter: Vec,\n}\n\nfn default_empty_input() -> serde_json::Value {\n serde_json::json!({})\n}\n\nfn deserialize_input<'de, D>(deserializer: D) -> Result\nwhere\n D: serde::Deserializer<'de>,\n{\n let value = serde_json::Value::deserialize(deserializer)?;\n Ok(if value.is_null() {\n serde_json::json!({})\n } else {\n value\n })\n}\n\nfn default_webhook_events_filter() -> Vec {\n vec![\n WebhookEventType::Start,\n WebhookEventType::Output,\n WebhookEventType::Logs,\n WebhookEventType::Completed,\n ]\n}\n\nfn generate_prediction_id() -> String {\n use std::time::{SystemTime, UNIX_EPOCH};\n // SAFETY: SystemTime::now() is always after UNIX_EPOCH on any reasonable system.\n // This cannot fail unless the system clock is set before 1970.\n let timestamp = SystemTime::now()\n .duration_since(UNIX_EPOCH)\n .expect(\"system clock is before 1970\")\n .as_nanos();\n format!(\"pred_{:x}\", timestamp)\n}\n\n/// Root discovery endpoint — returns a map of available API endpoints.\n///\n/// Restores the `GET /` endpoint from cog <= 0.16.x for service discovery.\n/// `cog_version` reports the Python SDK version when available (matching the\n/// old Python server behaviour), falling back to the coglet runtime version.\nasync fn root(State(service): State>) -> Json {\n let version = service.version();\n let cog_version = version.python_sdk.as_deref().unwrap_or(version.coglet);\n let mut doc = serde_json::json!({\n \"cog_version\": cog_version,\n \"docs_url\": \"/docs\",\n \"openapi_url\": \"/openapi.json\",\n \"shutdown_url\": \"/shutdown\",\n \"healthcheck_url\": \"/health-check\",\n \"predictions_url\": \"/predictions\",\n \"predictions_idempotent_url\": \"/predictions/{prediction_id}\",\n \"predictions_cancel_url\": \"/predictions/{prediction_id}/cancel\",\n });\n\n if service.supports_training().await {\n let obj = doc.as_object_mut().expect(\"doc is an object\");\n obj.insert(\"trainings_url\".to_string(), serde_json::json!(\"/trainings\"));\n obj.insert(\n \"trainings_idempotent_url\".to_string(),\n serde_json::json!(\"/trainings/{training_id}\"),\n );\n obj.insert(\n \"trainings_cancel_url\".to_string(),\n serde_json::json!(\"/trainings/{training_id}/cancel\"),\n );\n }\n\n Json(doc)\n}\n\nasync fn health_check(State(service): State>) -> Json {\n tracing::trace!(\"Health check endpoint called\");\n let snapshot = service.health().await;\n tracing::trace!(\n state = ?snapshot.state,\n available_slots = snapshot.available_slots,\n total_slots = snapshot.total_slots,\n has_setup_result = snapshot.setup_result.is_some(),\n \"Health snapshot retrieved\"\n );\n\n // Run user healthcheck if ready (even when busy — healthcheck health\n // and slot availability are orthogonal concerns).\n let user_healthcheck_error = if snapshot.is_ready() {\n write_readiness_file();\n\n // Run user-defined healthcheck\n tracing::trace!(\"Running user-defined healthcheck\");\n match service.healthcheck().await {\n Ok(result) if result.is_healthy() => {\n tracing::trace!(\"User healthcheck passed\");\n None\n }\n Ok(result) => {\n tracing::debug!(error = ?result.error, \"User healthcheck reported unhealthy\");\n result.error\n }\n Err(e) => {\n tracing::debug!(error = %e, \"User healthcheck returned error\");\n Some(format!(\"Healthcheck error: {}\", e))\n }\n }\n } else {\n tracing::trace!(state = ?snapshot.state, \"Skipping user healthcheck (not ready)\");\n None\n };\n\n let response = HealthCheckResponse::from_snapshot(snapshot, user_healthcheck_error);\n tracing::trace!(status = ?response.status, \"Health check response\");\n Json(response)\n}\n\n/// Write /var/run/cog/ready for K8s readiness probe.\nfn write_readiness_file() {\n if std::env::var(\"KUBERNETES_SERVICE_HOST\").is_err() {\n return;\n }\n\n let dir = std::path::Path::new(\"/var/run/cog\");\n let file = dir.join(\"ready\");\n\n if file.exists() {\n return;\n }\n\n if let Err(e) = std::fs::create_dir_all(dir) {\n tracing::warn!(error = %e, \"Failed to create /var/run/cog directory\");\n return;\n }\n\n if let Err(e) = std::fs::write(&file, b\"\") {\n tracing::warn!(error = %e, \"Failed to write readiness file\");\n }\n}\n\nfn should_respond_async(headers: &HeaderMap) -> bool {\n headers\n .get(\"prefer\")\n .and_then(|v| v.to_str().ok())\n .map(|v| v == \"respond-async\")\n .unwrap_or(false)\n}\n\nfn extract_trace_context(headers: &HeaderMap) -> TraceContext {\n TraceContext {\n traceparent: headers\n .get(\"traceparent\")\n .and_then(|v| v.to_str().ok())\n .map(|s| s.to_string()),\n tracestate: headers\n .get(\"tracestate\")\n .and_then(|v| v.to_str().ok())\n .map(|s| s.to_string()),\n }\n}\n\nasync fn create_prediction(\n State(service): State>,\n headers: HeaderMap,\n body: Option>,\n) -> impl IntoResponse {\n let request = body.map(|Json(r)| r).unwrap_or_else(|| PredictionRequest {\n id: None,\n input: serde_json::json!({}),\n context: Default::default(),\n webhook: None,\n webhook_events_filter: default_webhook_events_filter(),\n });\n let prediction_id = request.id.unwrap_or_else(generate_prediction_id);\n let respond_async = should_respond_async(&headers);\n let trace_context = extract_trace_context(&headers);\n create_prediction_with_id(\n service,\n prediction_id,\n request.input,\n request.context,\n request.webhook,\n request.webhook_events_filter,\n respond_async,\n trace_context,\n false,\n )\n .await\n}\n\nasync fn create_prediction_idempotent(\n State(service): State>,\n Path(prediction_id): Path,\n headers: HeaderMap,\n body: Option>,\n) -> impl IntoResponse {\n let request = body.map(|Json(r)| r).unwrap_or_else(|| PredictionRequest {\n id: None,\n input: serde_json::json!({}),\n context: Default::default(),\n webhook: None,\n webhook_events_filter: default_webhook_events_filter(),\n });\n\n if let Some(ref req_id) = request.id\n && req_id != &prediction_id\n {\n return (\n StatusCode::UNPROCESSABLE_ENTITY,\n Json(serde_json::json!({\n \"detail\": [{\n \"loc\": [\"body\", \"id\"],\n \"msg\": \"prediction ID must match the ID supplied in the URL\",\n \"type\": \"value_error\"\n }]\n })),\n );\n }\n\n // Check if prediction with this ID is already in-flight\n if let Some(response) = service.get_prediction_response(&prediction_id) {\n return (StatusCode::ACCEPTED, Json(response));\n }\n\n let respond_async = should_respond_async(&headers);\n let trace_context = extract_trace_context(&headers);\n create_prediction_with_id(\n service,\n prediction_id,\n request.input,\n request.context,\n request.webhook,\n request.webhook_events_filter,\n respond_async,\n trace_context,\n false,\n )\n .await\n}\n\nfn build_webhook_sender(\n webhook: Option,\n events_filter: Vec,\n trace_context: TraceContext,\n) -> Option {\n let webhook_url = webhook?;\n let events: std::collections::HashSet<_> = events_filter.into_iter().collect();\n\n match WebhookSender::with_trace_context(\n webhook_url.clone(),\n WebhookConfig {\n events_filter: events,\n ..Default::default()\n },\n trace_context,\n ) {\n Ok(sender) => Some(sender),\n Err(e) => {\n tracing::error!(url = %webhook_url, error = %e, \"Failed to create webhook sender\");\n None\n }\n }\n}\n\n#[allow(clippy::too_many_arguments)]\nasync fn create_prediction_with_id(\n service: Arc,\n prediction_id: String,\n input: serde_json::Value,\n context: std::collections::HashMap,\n webhook: Option,\n webhook_events_filter: Vec,\n respond_async: bool,\n trace_context: TraceContext,\n is_training: bool,\n) -> (StatusCode, Json) {\n // Validate input against the appropriate schema\n let validation_result = if is_training {\n service.validate_train_input(&input).await\n } else {\n service.validate_input(&input).await\n };\n if let Err(errors) = validation_result {\n let detail: Vec = errors\n .into_iter()\n .map(|e| {\n serde_json::json!({\n \"loc\": [\"body\", \"input\", e.field],\n \"msg\": e.msg,\n \"type\": e.error_type\n })\n })\n .collect();\n return (\n StatusCode::UNPROCESSABLE_ENTITY,\n Json(serde_json::json!({ \"detail\": detail })),\n );\n }\n\n let webhook_sender = build_webhook_sender(\n webhook.clone(),\n webhook_events_filter.clone(),\n trace_context.clone(),\n );\n\n // Submit prediction: creates Prediction, acquires slot, registers in service\n let (handle, unregistered_slot) = match service\n .submit_prediction(prediction_id.clone(), input.clone(), webhook_sender)\n .await\n {\n Ok(r) => r,\n Err(CreatePredictionError::NotReady) => {\n let msg = PredictionError::NotReady.to_string();\n return (\n StatusCode::SERVICE_UNAVAILABLE,\n Json(serde_json::json!({\n \"error\": msg,\n \"status\": \"failed\"\n })),\n );\n }\n Err(CreatePredictionError::AtCapacity) => {\n return (\n StatusCode::CONFLICT,\n Json(serde_json::json!({\n \"error\": \"At capacity - all prediction slots busy\",\n \"status\": \"failed\"\n })),\n );\n }\n };\n\n let prediction = unregistered_slot.prediction();\n\n // Async mode: spawn background task, return immediately\n if respond_async {\n let service_clone = Arc::clone(&service);\n let id_for_cleanup = prediction_id.clone();\n let context_async = context.clone();\n tokio::spawn(async move {\n let _result = service_clone\n .predict(unregistered_slot, input, context_async)\n .await;\n // Prediction state is already updated by predict() internally\n // (set_succeeded/set_failed/set_canceled fire webhooks automatically)\n service_clone.remove_prediction(&id_for_cleanup);\n });\n\n return (\n StatusCode::ACCEPTED,\n Json(serde_json::json!({\n \"id\": prediction_id,\n \"status\": \"starting\"\n })),\n );\n }\n\n // Sync mode: spawn prediction into a background task so the slot lifetime\n // is NOT tied to the HTTP connection. If the client disconnects, the\n // SyncPredictionGuard fires cancel, but the slot/permit stays alive in the\n // spawned task until the worker acknowledges the cancel.\n let mut sync_guard = handle.sync_guard(Arc::clone(&service));\n\n let service_bg = Arc::clone(&service);\n let id_bg = prediction_id.clone();\n let result_rx = {\n let (tx, rx) = tokio::sync::oneshot::channel();\n tokio::spawn(async move {\n let result = service_bg.predict(unregistered_slot, input, context).await;\n // Prediction state is already updated by predict() internally\n service_bg.remove_prediction(&id_bg);\n let _ = tx.send(result);\n });\n rx\n };\n\n // Wait for the prediction to complete. If the connection drops, axum\n // cancels this future, dropping sync_guard which fires cancel.\n let result = match result_rx.await {\n Ok(r) => r,\n Err(_) => {\n // Background task panicked or was cancelled\n Err(PredictionError::Failed(\"prediction task lost\".to_string()))\n }\n };\n\n let predict_time = prediction\n .try_lock()\n .map(|p| p.elapsed())\n .unwrap_or(std::time::Duration::ZERO)\n .as_secs_f64();\n\n // Disarm guard - prediction completed normally (connection still alive)\n sync_guard.disarm();\n\n // Build metrics object: user metrics + predict_time\n let build_metrics = |user_metrics: &std::collections::HashMap| {\n let mut m = serde_json::Map::new();\n for (k, v) in user_metrics {\n m.insert(k.clone(), v.clone());\n }\n m.insert(\"predict_time\".to_string(), serde_json::json!(predict_time));\n serde_json::Value::Object(m)\n };\n\n match result {\n Ok(r) => {\n let metrics = build_metrics(&r.metrics);\n (\n StatusCode::OK,\n Json(serde_json::json!({\n \"id\": prediction_id,\n \"output\": r.output,\n \"logs\": r.logs,\n \"status\": \"succeeded\",\n \"metrics\": metrics\n })),\n )\n }\n Err(PredictionError::InvalidInput(msg)) => (\n StatusCode::UNPROCESSABLE_ENTITY,\n Json(serde_json::json!({\n \"id\": prediction_id,\n \"error\": msg,\n \"logs\": \"\",\n \"status\": \"failed\",\n \"metrics\": { \"predict_time\": predict_time }\n })),\n ),\n Err(PredictionError::NotReady) => {\n let msg = PredictionError::NotReady.to_string();\n (\n StatusCode::SERVICE_UNAVAILABLE,\n Json(serde_json::json!({\n \"id\": prediction_id,\n \"error\": msg,\n \"logs\": \"\",\n \"status\": \"failed\"\n })),\n )\n }\n Err(PredictionError::Failed(msg)) => (\n // 200 for parity with Python - prediction failure is data, not HTTP error\n StatusCode::OK,\n Json(serde_json::json!({\n \"id\": prediction_id,\n \"error\": msg,\n \"logs\": \"\",\n \"status\": \"failed\",\n \"metrics\": { \"predict_time\": predict_time }\n })),\n ),\n Err(PredictionError::Cancelled) => (\n StatusCode::OK,\n Json(serde_json::json!({\n \"id\": prediction_id,\n \"logs\": \"\",\n \"status\": \"canceled\",\n \"metrics\": { \"predict_time\": predict_time }\n })),\n ),\n }\n}\n\nasync fn cancel_prediction(\n State(service): State>,\n Path(prediction_id): Path,\n) -> impl IntoResponse {\n let cancelled = service.cancel(&prediction_id);\n\n if cancelled {\n (StatusCode::OK, Json(serde_json::json!({})))\n } else {\n (StatusCode::NOT_FOUND, Json(serde_json::json!({})))\n }\n}\n\nasync fn shutdown(State(service): State>) -> impl IntoResponse {\n tracing::info!(\"Shutdown requested via HTTP\");\n service.trigger_shutdown();\n (StatusCode::OK, Json(serde_json::json!({})))\n}\n\nasync fn openapi_schema(State(service): State>) -> impl IntoResponse {\n match service.schema().await {\n Some(schema) => (StatusCode::OK, Json(schema)),\n None => (\n StatusCode::SERVICE_UNAVAILABLE,\n Json(serde_json::json!({\n \"error\": \"OpenAPI schema not available\"\n })),\n ),\n }\n}\n\n// Training routes — same dispatch as predictions but validated against\n// TrainingInput schema instead of Input.\n\nasync fn create_training(\n State(service): State>,\n headers: HeaderMap,\n body: Option>,\n) -> impl IntoResponse {\n let request = body.map(|Json(r)| r).unwrap_or_else(|| PredictionRequest {\n id: None,\n input: serde_json::json!({}),\n context: Default::default(),\n webhook: None,\n webhook_events_filter: default_webhook_events_filter(),\n });\n let prediction_id = request.id.unwrap_or_else(generate_prediction_id);\n let respond_async = should_respond_async(&headers);\n let trace_context = extract_trace_context(&headers);\n create_prediction_with_id(\n service,\n prediction_id,\n request.input,\n request.context,\n request.webhook,\n request.webhook_events_filter,\n respond_async,\n trace_context,\n true,\n )\n .await\n}\n\nasync fn create_training_idempotent(\n State(service): State>,\n Path(training_id): Path,\n headers: HeaderMap,\n body: Option>,\n) -> impl IntoResponse {\n let request = body.map(|Json(r)| r).unwrap_or_else(|| PredictionRequest {\n id: None,\n input: serde_json::json!({}),\n context: Default::default(),\n webhook: None,\n webhook_events_filter: default_webhook_events_filter(),\n });\n\n if let Some(ref req_id) = request.id\n && req_id != &training_id\n {\n return (\n StatusCode::UNPROCESSABLE_ENTITY,\n Json(serde_json::json!({\n \"detail\": [{\n \"loc\": [\"body\", \"id\"],\n \"msg\": \"training ID must match the ID supplied in the URL\",\n \"type\": \"value_error\"\n }]\n })),\n );\n }\n\n // Idempotent: return existing state if already submitted\n if let Some(response) = service.get_prediction_response(&training_id) {\n return (StatusCode::ACCEPTED, Json(response));\n }\n\n let respond_async = should_respond_async(&headers);\n let trace_context = extract_trace_context(&headers);\n create_prediction_with_id(\n service,\n training_id,\n request.input,\n request.context,\n request.webhook,\n request.webhook_events_filter,\n respond_async,\n trace_context,\n true,\n )\n .await\n}\n\nasync fn cancel_training(\n State(service): State>,\n Path(training_id): Path,\n) -> impl IntoResponse {\n cancel_prediction(State(service), Path(training_id)).await\n}\n\n/// Maximum HTTP request body size (100 MiB).\n///\n/// Axum defaults to 2 MiB which is too small for models that accept large\n/// inline inputs (e.g. base64-encoded images). Inputs that exceed the IPC\n/// frame limit are automatically spilled to disk by `build_slot_request`.\nconst MAX_HTTP_BODY_SIZE: usize = 100 * 1024 * 1024;\n\npub fn routes(service: Arc) -> Router {\n Router::new()\n .route(\"/\", get(root))\n .route(\"/health-check\", get(health_check))\n .route(\"/openapi.json\", get(openapi_schema))\n .route(\"/shutdown\", post(shutdown))\n .route(\"/predictions\", post(create_prediction))\n .route(\"/predictions/{id}\", put(create_prediction_idempotent))\n .route(\"/predictions/{id}/cancel\", post(cancel_prediction))\n .route(\"/trainings\", post(create_training))\n .route(\"/trainings/{id}\", put(create_training_idempotent))\n .route(\"/trainings/{id}/cancel\", post(cancel_training))\n .layer(DefaultBodyLimit::max(MAX_HTTP_BODY_SIZE))\n .with_state(service)\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use axum::body::Body;\n use axum::http::{Request, StatusCode};\n use http_body_util::BodyExt;\n use tower::ServiceExt;\n\n async fn response_json(response: axum::response::Response) -> serde_json::Value {\n let body = response.into_body();\n let bytes = body.collect().await.unwrap().to_bytes();\n serde_json::from_slice(&bytes).unwrap()\n }\n\n #[tokio::test]\n async fn health_check_returns_status_and_version() {\n let service = Arc::new(PredictionService::new_no_pool().with_health(Health::Starting));\n let app = routes(service);\n\n let response = app\n .oneshot(Request::get(\"/health-check\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n\n let json = response_json(response).await;\n assert_eq!(json[\"status\"], \"STARTING\");\n assert!(json[\"version\"][\"coglet\"].is_string());\n }\n\n #[tokio::test]\n async fn health_check_unknown_when_no_predictor() {\n let service = Arc::new(PredictionService::new_no_pool());\n let app = routes(service);\n\n let response = app\n .oneshot(Request::get(\"/health-check\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n let json = response_json(response).await;\n assert_eq!(json[\"status\"], \"UNKNOWN\");\n }\n\n #[tokio::test]\n async fn predictions_returns_503_when_not_ready() {\n let service = Arc::new(PredictionService::new_no_pool());\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::post(\"/predictions\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);\n\n let json = response_json(response).await;\n assert_eq!(json[\"status\"], \"failed\");\n assert!(\n json[\"error\"]\n .as_str()\n .unwrap()\n .contains(\"Setup has not finished yet\")\n );\n }\n\n #[tokio::test]\n async fn openapi_returns_503_when_schema_not_available() {\n let service = Arc::new(PredictionService::new_no_pool());\n let app = routes(service);\n\n let response = app\n .oneshot(Request::get(\"/openapi.json\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);\n\n let json = response_json(response).await;\n assert!(json[\"error\"].as_str().unwrap().contains(\"not available\"));\n }\n\n #[tokio::test]\n async fn openapi_returns_schema_when_available() {\n let service = Arc::new(PredictionService::new_no_pool());\n service\n .set_schema(serde_json::json!({\n \"openapi\": \"3.0.2\",\n \"info\": {\"title\": \"Cog\", \"version\": \"0.1.0\"}\n }))\n .await;\n let app = routes(service);\n\n let response = app\n .oneshot(Request::get(\"/openapi.json\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n\n let json = response_json(response).await;\n assert_eq!(json[\"openapi\"], \"3.0.2\");\n assert_eq!(json[\"info\"][\"title\"], \"Cog\");\n }\n\n // --- Tests with MockOrchestrator for full prediction flow ---\n\n use crate::PredictionOutput;\n use crate::bridge::protocol::SlotId;\n use crate::orchestrator::Orchestrator;\n use crate::permit::PermitPool;\n use std::sync::Mutex as StdMutex;\n use std::sync::atomic::{AtomicUsize, Ordering};\n\n /// Mock orchestrator that immediately completes predictions.\n struct MockOrchestrator {\n register_count: AtomicUsize,\n complete_immediately: bool,\n }\n\n impl MockOrchestrator {\n fn new() -> Self {\n Self {\n register_count: AtomicUsize::new(0),\n complete_immediately: true,\n }\n }\n\n /// Create a mock that never completes predictions (for capacity tests).\n fn never_complete() -> Self {\n Self {\n register_count: AtomicUsize::new(0),\n complete_immediately: false,\n }\n }\n }\n\n #[async_trait::async_trait]\n impl Orchestrator for MockOrchestrator {\n async fn register_prediction(\n &self,\n _slot_id: SlotId,\n prediction: Arc>,\n _idle_sender: tokio::sync::oneshot::Sender,\n ) {\n self.register_count.fetch_add(1, Ordering::SeqCst);\n if self.complete_immediately {\n let mut pred = prediction.lock().unwrap();\n pred.set_succeeded(PredictionOutput::Single(serde_json::json!(\"mock output\")));\n }\n }\n\n async fn cancel_by_prediction_id(\n &self,\n _prediction_id: &str,\n ) -> Result<(), crate::orchestrator::OrchestratorError> {\n Ok(())\n }\n\n async fn healthcheck(\n &self,\n ) -> Result\n {\n Ok(crate::orchestrator::HealthcheckResult::healthy())\n }\n\n async fn shutdown(&self) -> Result<(), crate::orchestrator::OrchestratorError> {\n Ok(())\n }\n }\n\n async fn create_test_pool(num_slots: usize) -> Arc {\n use crate::bridge::codec::JsonCodec;\n use crate::bridge::protocol::SlotRequest;\n use futures::StreamExt;\n use tokio::net::UnixStream;\n\n let pool = Arc::new(PermitPool::new(num_slots));\n for _ in 0..num_slots {\n let (a, b) = UnixStream::pair().unwrap();\n let (_read_a, write_a) = a.into_split();\n let (read_b, _write_b) = b.into_split();\n\n // Spawn a task to consume messages from the socket (prevents broken pipe)\n let mut reader =\n tokio_util::codec::FramedRead::new(read_b, JsonCodec::::new());\n tokio::spawn(async move { while reader.next().await.is_some() {} });\n\n let writer =\n tokio_util::codec::FramedWrite::new(write_a, JsonCodec::::new());\n pool.add_permit(SlotId::new(), writer);\n }\n pool\n }\n\n async fn create_ready_service() -> Arc {\n let service = Arc::new(PredictionService::new_no_pool());\n let pool = create_test_pool(2).await;\n let orchestrator = Arc::new(MockOrchestrator::new());\n service.set_orchestrator(pool, orchestrator).await;\n service.set_health(Health::Ready).await;\n service\n }\n\n #[tokio::test]\n async fn health_check_ready_with_orchestrator() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(Request::get(\"/health-check\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n let json = response_json(response).await;\n assert_eq!(json[\"status\"], \"READY\");\n }\n\n #[tokio::test]\n async fn prediction_sync_success() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::post(\"/predictions\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"input\":{\"prompt\":\"hello\"}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n let json = response_json(response).await;\n assert_eq!(json[\"status\"], \"succeeded\");\n assert_eq!(json[\"output\"], \"mock output\");\n assert!(json[\"id\"].is_string());\n }\n\n #[tokio::test]\n async fn prediction_async_returns_accepted() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::post(\"/predictions\")\n .header(\"content-type\", \"application/json\")\n .header(\"prefer\", \"respond-async\")\n .body(Body::from(r#\"{\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::ACCEPTED);\n let json = response_json(response).await;\n assert_eq!(json[\"status\"], \"starting\");\n }\n\n #[tokio::test]\n async fn prediction_with_custom_id() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::post(\"/predictions\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"id\":\"my-pred-123\",\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n let json = response_json(response).await;\n assert_eq!(json[\"id\"], \"my-pred-123\");\n assert_eq!(json[\"status\"], \"succeeded\");\n }\n\n #[tokio::test]\n async fn prediction_idempotent_put() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::put(\"/predictions/idempotent-123\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n let json = response_json(response).await;\n assert_eq!(json[\"id\"], \"idempotent-123\");\n assert_eq!(json[\"status\"], \"succeeded\");\n }\n\n #[tokio::test]\n async fn prediction_idempotent_id_mismatch() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::put(\"/predictions/url-id\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"id\":\"body-id\",\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);\n let json = response_json(response).await;\n assert!(\n json[\"detail\"][0][\"msg\"]\n .as_str()\n .unwrap()\n .contains(\"must match\")\n );\n }\n\n #[tokio::test]\n async fn prediction_at_capacity() {\n let service = Arc::new(PredictionService::new_no_pool());\n let pool = create_test_pool(1).await; // Only 1 slot\n // Use never_complete so the first prediction holds the slot\n let orchestrator = Arc::new(MockOrchestrator::never_complete());\n service.set_orchestrator(pool, orchestrator).await;\n service.set_health(Health::Ready).await;\n\n // Use async mode so first request doesn't block\n let app = routes(Arc::clone(&service));\n let _resp1 = app\n .oneshot(\n Request::post(\"/predictions\")\n .header(\"content-type\", \"application/json\")\n .header(\"prefer\", \"respond-async\")\n .body(Body::from(r#\"{\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n // Small delay to let async task acquire the slot\n tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;\n\n // Second request should get 409 Conflict (at capacity)\n let app2 = routes(service);\n let response = app2\n .oneshot(\n Request::post(\"/predictions\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::CONFLICT);\n let json = response_json(response).await;\n assert!(json[\"error\"].as_str().unwrap().contains(\"capacity\"));\n }\n\n #[tokio::test]\n async fn health_check_busy_when_at_capacity() {\n let service = Arc::new(PredictionService::new_no_pool());\n let pool = create_test_pool(1).await;\n // Use never_complete so the prediction holds the slot\n let orchestrator = Arc::new(MockOrchestrator::never_complete());\n service.set_orchestrator(pool, orchestrator).await;\n service.set_health(Health::Ready).await;\n\n // Use async to hold the slot\n let app = routes(Arc::clone(&service));\n let _resp = app\n .oneshot(\n Request::post(\"/predictions\")\n .header(\"content-type\", \"application/json\")\n .header(\"prefer\", \"respond-async\")\n .body(Body::from(r#\"{\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;\n\n // Health should show BUSY\n let app2 = routes(service);\n let response = app2\n .oneshot(Request::get(\"/health-check\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n let json = response_json(response).await;\n assert_eq!(json[\"status\"], \"BUSY\");\n }\n\n #[tokio::test]\n async fn training_routes_work() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::post(\"/trainings\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n let json = response_json(response).await;\n assert_eq!(json[\"status\"], \"succeeded\");\n }\n\n #[tokio::test]\n async fn training_idempotent_put() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::put(\"/trainings/train-123\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n let json = response_json(response).await;\n assert_eq!(json[\"id\"], \"train-123\");\n assert_eq!(json[\"status\"], \"succeeded\");\n }\n\n #[tokio::test]\n async fn training_idempotent_id_mismatch() {\n let service = create_ready_service().await;\n let app = routes(service);\n\n let response = app\n .oneshot(\n Request::put(\"/trainings/url-id\")\n .header(\"content-type\", \"application/json\")\n .body(Body::from(r#\"{\"id\":\"body-id\",\"input\":{}}\"#))\n .unwrap(),\n )\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);\n let json = response_json(response).await;\n assert!(\n json[\"detail\"][0][\"msg\"]\n .as_str()\n .unwrap()\n .contains(\"must match\")\n );\n }\n\n #[tokio::test]\n async fn shutdown_triggers_service_shutdown() {\n let service = create_ready_service().await;\n let mut rx = service.shutdown_rx();\n let app = routes(service);\n\n assert!(!*rx.borrow());\n\n let response = app\n .oneshot(Request::post(\"/shutdown\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n rx.changed().await.unwrap();\n assert!(*rx.borrow());\n }\n\n #[tokio::test]\n async fn root_returns_discovery_document() {\n let service = Arc::new(PredictionService::new_no_pool());\n let app = routes(service);\n\n let response = app\n .oneshot(Request::get(\"/\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n assert_eq!(\n response.headers().get(\"content-type\").unwrap(),\n \"application/json\"\n );\n\n let json = response_json(response).await;\n // Without a python_sdk version set, falls back to coglet version\n assert_eq!(json[\"cog_version\"], crate::version::COGLET_VERSION);\n assert_eq!(json[\"docs_url\"], \"/docs\");\n assert_eq!(json[\"openapi_url\"], \"/openapi.json\");\n assert_eq!(json[\"shutdown_url\"], \"/shutdown\");\n assert_eq!(json[\"healthcheck_url\"], \"/health-check\");\n assert_eq!(json[\"predictions_url\"], \"/predictions\");\n assert_eq!(\n json[\"predictions_idempotent_url\"],\n \"/predictions/{prediction_id}\"\n );\n assert_eq!(\n json[\"predictions_cancel_url\"],\n \"/predictions/{prediction_id}/cancel\"\n );\n // No training URLs without a TrainingInput schema\n assert!(json.get(\"trainings_url\").is_none());\n assert!(json.get(\"trainings_idempotent_url\").is_none());\n assert!(json.get(\"trainings_cancel_url\").is_none());\n }\n\n #[tokio::test]\n async fn root_includes_training_urls_when_schema_has_training() {\n let service = Arc::new(PredictionService::new_no_pool());\n // Set a schema that includes a TrainingInput component\n service\n .set_schema(serde_json::json!({\n \"openapi\": \"3.0.2\",\n \"info\": {\"title\": \"Cog\", \"version\": \"0.1.0\"},\n \"components\": {\n \"schemas\": {\n \"TrainingInput\": {\n \"type\": \"object\",\n \"properties\": {\n \"data\": {\"type\": \"string\"}\n }\n }\n }\n }\n }))\n .await;\n let app = routes(service);\n\n let response = app\n .oneshot(Request::get(\"/\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n assert_eq!(response.status(), StatusCode::OK);\n\n let json = response_json(response).await;\n // Base fields still present\n assert_eq!(json[\"predictions_url\"], \"/predictions\");\n // Training URLs included\n assert_eq!(json[\"trainings_url\"], \"/trainings\");\n assert_eq!(json[\"trainings_idempotent_url\"], \"/trainings/{training_id}\");\n assert_eq!(\n json[\"trainings_cancel_url\"],\n \"/trainings/{training_id}/cancel\"\n );\n }\n\n #[tokio::test]\n async fn root_cog_version_prefers_python_sdk() {\n let version = VersionInfo::new().with_python_sdk(\"0.14.0\".to_string());\n let service = Arc::new(PredictionService::new_no_pool().with_version(version));\n let app = routes(service);\n\n let response = app\n .oneshot(Request::get(\"/\").body(Body::empty()).unwrap())\n .await\n .unwrap();\n\n let json = response_json(response).await;\n assert_eq!(json[\"cog_version\"], \"0.14.0\");\n }\n}\n" }, { "path": "crates/coglet/src/transport/http/server.rs", "content": "//! HTTP server implementation.\n\nuse std::net::SocketAddr;\nuse std::sync::Arc;\n\nuse tokio::net::TcpListener;\nuse tokio::sync::watch;\nuse tracing::info;\n\nuse crate::service::PredictionService;\n\nuse super::routes::routes;\n\n#[derive(Debug, Clone)]\npub struct ServerConfig {\n pub host: String,\n pub port: u16,\n /// If true, ignore SIGTERM and wait for explicit /shutdown or SIGINT.\n /// Used in Kubernetes to allow graceful draining.\n pub await_explicit_shutdown: bool,\n}\n\nimpl Default for ServerConfig {\n fn default() -> Self {\n Self {\n host: \"0.0.0.0\".to_string(),\n port: 5000,\n await_explicit_shutdown: false,\n }\n }\n}\n\n/// Start the HTTP server with provided service.\npub async fn serve(config: ServerConfig, service: Arc) -> anyhow::Result<()> {\n let shutdown_rx = service.shutdown_rx();\n let app = routes(service.clone());\n\n let addr: SocketAddr = format!(\"{}:{}\", config.host, config.port).parse()?;\n\n let listener = TcpListener::bind(addr).await?;\n let actual_addr = listener.local_addr()?;\n\n info!(\"Starting coglet server on {}\", actual_addr);\n\n axum::serve(listener, app)\n .with_graceful_shutdown(shutdown_signal(config.await_explicit_shutdown, shutdown_rx))\n .await?;\n\n info!(\"Server shutdown complete\");\n\n // Gracefully shutdown the orchestrator worker\n service.shutdown().await;\n\n Ok(())\n}\n\n/// Wait for shutdown signal (SIGTERM, SIGINT, or /shutdown endpoint).\n///\n/// # Panics\n///\n/// Panics if signal handlers cannot be installed. This can only happen if:\n/// - Called from a non-main thread without the runtime being properly configured\n/// - The tokio runtime is not properly initialized\n///\n/// These are unrecoverable configuration errors that should fail fast at startup.\nasync fn shutdown_signal(await_explicit_shutdown: bool, mut shutdown_rx: watch::Receiver) {\n let ctrl_c = async {\n tokio::signal::ctrl_c()\n .await\n .expect(\"failed to install Ctrl+C handler - is tokio runtime configured correctly?\");\n };\n\n #[cfg(unix)]\n let terminate = async {\n if await_explicit_shutdown {\n // Ignore SIGTERM - wait forever (until SIGINT or explicit shutdown)\n tracing::info!(\"await_explicit_shutdown enabled, ignoring SIGTERM\");\n std::future::pending::<()>().await\n } else {\n tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())\n .expect(\n \"failed to install SIGTERM handler - is tokio runtime configured correctly?\",\n )\n .recv()\n .await;\n }\n };\n\n #[cfg(not(unix))]\n let terminate = std::future::pending::<()>();\n\n let explicit_shutdown = async {\n while !*shutdown_rx.borrow() {\n if shutdown_rx.changed().await.is_err() {\n std::future::pending::<()>().await;\n }\n }\n };\n\n tokio::select! {\n _ = ctrl_c => {\n info!(\"Received SIGINT, shutting down...\");\n }\n _ = terminate => {\n info!(\"Received SIGTERM, shutting down...\");\n }\n _ = explicit_shutdown => {\n info!(\"Shutdown requested via /shutdown endpoint...\");\n }\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn server_config_default() {\n let config = ServerConfig::default();\n assert_eq!(config.host, \"0.0.0.0\");\n assert_eq!(config.port, 5000);\n assert!(!config.await_explicit_shutdown);\n }\n}\n" }, { "path": "crates/coglet/src/transport/mod.rs", "content": "//! Transport layer for coglet.\n//!\n//! Currently provides HTTP transport via axum. Future transports\n//! (gRPC, bnet) will be added as separate submodules.\n\npub mod http;\n\npub use http::{ServerConfig, serve};\n" }, { "path": "crates/coglet/src/version.rs", "content": "//! Version information for coglet.\n\n/// Coglet version from Cargo.toml\npub const COGLET_VERSION: &str = env!(\"CARGO_PKG_VERSION\");\n\n/// Version information for the runtime.\n#[derive(Debug, Clone, serde::Serialize)]\npub struct VersionInfo {\n /// Coglet runtime version.\n pub coglet: &'static str,\n /// Git SHA (with optional `-dirty` suffix).\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub git_sha: Option,\n /// Build timestamp (UTC, ISO 8601).\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub build_time: Option,\n /// Python SDK version (if available).\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub python_sdk: Option,\n /// Python version.\n #[serde(skip_serializing_if = \"Option::is_none\")]\n pub python: Option,\n}\n\nimpl Default for VersionInfo {\n fn default() -> Self {\n Self {\n coglet: COGLET_VERSION,\n git_sha: None,\n build_time: None,\n python_sdk: None,\n python: None,\n }\n }\n}\n\nimpl VersionInfo {\n /// Create version info with coglet version only.\n pub fn new() -> Self {\n Self::default()\n }\n\n /// Set git SHA (with optional `-dirty` suffix).\n pub fn with_git_sha(mut self, sha: String) -> Self {\n self.git_sha = Some(sha);\n self\n }\n\n /// Set build timestamp.\n pub fn with_build_time(mut self, time: String) -> Self {\n self.build_time = Some(time);\n self\n }\n\n /// Set Python SDK version.\n pub fn with_python_sdk(mut self, version: String) -> Self {\n self.python_sdk = Some(version);\n self\n }\n\n /// Set Python version.\n pub fn with_python(mut self, version: String) -> Self {\n self.python = Some(version);\n self\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn version_info_has_coglet_version() {\n let info = VersionInfo::new();\n assert_eq!(info.coglet, COGLET_VERSION);\n assert!(info.python_sdk.is_none());\n assert!(info.python.is_none());\n }\n\n #[test]\n fn version_info_builder_pattern() {\n let info = VersionInfo::new()\n .with_git_sha(\"abc1234\".to_string())\n .with_build_time(\"2026-03-12T18:00:00Z\".to_string())\n .with_python_sdk(\"0.9.0\".to_string())\n .with_python(\"3.11.0\".to_string());\n\n assert_eq!(info.git_sha, Some(\"abc1234\".to_string()));\n assert_eq!(info.build_time, Some(\"2026-03-12T18:00:00Z\".to_string()));\n assert_eq!(info.python_sdk, Some(\"0.9.0\".to_string()));\n assert_eq!(info.python, Some(\"3.11.0\".to_string()));\n }\n\n #[test]\n fn version_info_serializes_minimal() {\n // Only coglet when no optional fields set\n let info = VersionInfo {\n coglet: \"0.1.0\",\n git_sha: None,\n build_time: None,\n python_sdk: None,\n python: None,\n };\n insta::assert_json_snapshot!(\"version_minimal\", info);\n }\n\n #[test]\n fn version_info_serializes_full() {\n let info = VersionInfo {\n coglet: \"0.1.0\",\n git_sha: Some(\"abc1234-dirty\".to_string()),\n build_time: Some(\"2026-03-12T18:00:00Z\".to_string()),\n python_sdk: Some(\"0.9.0\".to_string()),\n python: Some(\"3.11.0\".to_string()),\n };\n insta::assert_json_snapshot!(\"version_full\", info);\n }\n}\n" }, { "path": "crates/coglet/src/webhook.rs", "content": "//! Webhook sender for async predictions.\n//!\n//! Implements the cog webhook protocol:\n//! - Throttling (default 500ms between non-terminal updates)\n//! - Terminal webhooks retried with exponential backoff\n//! - WEBHOOK_AUTH_TOKEN bearer authentication\n//! - Events filtering (start, output, logs, completed)\n//!\n//! # Panic Safety\n//!\n//! This module avoids panics:\n//! - `WebhookSender::new()` returns `Result` for HTTP client creation\n//! - Mutex locks use `lock().unwrap_or_else(|e| e.into_inner())` to recover from\n//! poison - worst case is we lose throttle tracking, which is acceptable.\n\nuse std::collections::HashSet;\nuse std::sync::Mutex;\nuse std::time::{Duration, Instant};\n\nuse thiserror::Error;\n\nuse serde::{Deserialize, Serialize};\n\nuse crate::version::COGLET_VERSION;\n\n#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Deserialize, Serialize)]\n#[serde(rename_all = \"lowercase\")]\npub enum WebhookEventType {\n Start,\n Output,\n Logs,\n #[default]\n Completed,\n}\n\nimpl WebhookEventType {\n pub fn is_terminal(&self) -> bool {\n matches!(self, Self::Completed)\n }\n\n pub fn all() -> HashSet {\n [Self::Start, Self::Output, Self::Logs, Self::Completed]\n .into_iter()\n .collect()\n }\n}\n\n/// Error creating a WebhookSender.\n#[derive(Debug, Error)]\npub enum WebhookSenderError {\n #[error(\"failed to create HTTP client: {0}\")]\n HttpClient(#[from] reqwest::Error),\n}\n\n#[derive(Debug, Clone)]\npub struct WebhookConfig {\n pub response_interval: Duration,\n pub events_filter: HashSet,\n pub max_retries: u32,\n pub backoff_base: Duration,\n pub retry_status_codes: Vec,\n}\n\nimpl Default for WebhookConfig {\n fn default() -> Self {\n Self {\n response_interval: Duration::from_millis(\n std::env::var(\"COG_THROTTLE_RESPONSE_INTERVAL\")\n .ok()\n .and_then(|s| s.parse::().ok())\n .map(|s| (s * 1000.0) as u64)\n .unwrap_or(500),\n ),\n events_filter: WebhookEventType::all(),\n max_retries: 12,\n backoff_base: Duration::from_millis(100),\n retry_status_codes: vec![429, 500, 502, 503, 504],\n }\n }\n}\n\n/// W3C Trace Context for distributed tracing.\n#[derive(Debug, Clone, Default)]\npub struct TraceContext {\n pub traceparent: Option,\n pub tracestate: Option,\n}\n\npub struct WebhookSender {\n url: String,\n config: WebhookConfig,\n client: reqwest::Client,\n last_sent: Mutex,\n trace_context: TraceContext,\n}\n\nimpl WebhookSender {\n pub fn new(url: String, config: WebhookConfig) -> Result {\n Self::with_trace_context(url, config, TraceContext::default())\n }\n\n pub fn with_trace_context(\n url: String,\n config: WebhookConfig,\n trace_context: TraceContext,\n ) -> Result {\n let mut headers = reqwest::header::HeaderMap::new();\n\n if let Ok(token) = std::env::var(\"WEBHOOK_AUTH_TOKEN\")\n && let Ok(value) = reqwest::header::HeaderValue::from_str(&format!(\"Bearer {}\", token))\n {\n headers.insert(reqwest::header::AUTHORIZATION, value);\n }\n\n let user_agent = format!(\"coglet/{}\", COGLET_VERSION);\n if let Ok(value) = reqwest::header::HeaderValue::from_str(&user_agent) {\n headers.insert(reqwest::header::USER_AGENT, value);\n }\n\n let client = reqwest::Client::builder()\n .default_headers(headers)\n .timeout(Duration::from_secs(30))\n .build()?;\n\n Ok(Self {\n url,\n config,\n client,\n // Allow immediate first send\n last_sent: Mutex::new(Instant::now() - Duration::from_secs(10)),\n trace_context,\n })\n }\n\n pub fn url(&self) -> &str {\n &self.url\n }\n\n fn should_send(&self, event: WebhookEventType) -> bool {\n if !self.config.events_filter.contains(&event) {\n return false;\n }\n\n if event.is_terminal() {\n return true;\n }\n\n // Output events are never throttled: they are high-value (contain actual\n // prediction results), relatively infrequent (one per output chunk/file),\n // and in the old Python runtime were effectively unthrottled because file\n // uploads were synchronous. Throttling them causes the director to miss\n // intermediate output data.\n if matches!(event, WebhookEventType::Output) {\n return true;\n }\n\n // Recover from poison - losing throttle state is acceptable\n let last = self.last_sent.lock().unwrap_or_else(|e| e.into_inner());\n last.elapsed() >= self.config.response_interval\n }\n\n fn update_last_sent(&self) {\n // Recover from poison - losing throttle state is acceptable\n let mut last = self.last_sent.lock().unwrap_or_else(|e| e.into_inner());\n *last = Instant::now();\n }\n\n fn build_request(&self, payload: &serde_json::Value) -> reqwest::RequestBuilder {\n let mut request = self.client.post(&self.url).json(payload);\n\n if let Some(ref traceparent) = self.trace_context.traceparent {\n request = request.header(\"traceparent\", traceparent);\n }\n if let Some(ref tracestate) = self.trace_context.tracestate {\n request = request.header(\"tracestate\", tracestate);\n }\n\n request\n }\n\n /// Send a non-terminal webhook (fire and forget, no retry).\n pub fn send(&self, event: WebhookEventType, payload: &serde_json::Value) {\n if !self.should_send(event) {\n return;\n }\n\n let request = self.build_request(payload);\n self.update_last_sent();\n\n tokio::spawn(async move {\n if let Err(e) = request.send().await {\n tracing::warn!(error = %e, \"Failed to send webhook (non-terminal)\");\n }\n });\n }\n\n /// Send a terminal webhook with exponential backoff retries.\n pub async fn send_terminal(&self, event: WebhookEventType, payload: &serde_json::Value) {\n if !self.config.events_filter.contains(&event) {\n return;\n }\n\n let mut attempt = 0;\n loop {\n match self.build_request(payload).send().await {\n Ok(response) => {\n let status = response.status().as_u16();\n if response.status().is_success() {\n tracing::debug!(status = %status, \"Terminal webhook sent successfully\");\n return;\n }\n\n if self.config.retry_status_codes.contains(&status) {\n attempt += 1;\n if attempt > self.config.max_retries {\n tracing::error!(\n status = %status,\n attempts = attempt,\n \"Terminal webhook failed after max retries\"\n );\n return;\n }\n\n let backoff = self.config.backoff_base * (1 << attempt.min(10));\n tracing::warn!(\n status = %status,\n attempt = attempt,\n backoff_ms = backoff.as_millis(),\n \"Terminal webhook failed, retrying\"\n );\n tokio::time::sleep(backoff).await;\n continue;\n }\n\n tracing::error!(\n status = %status,\n \"Terminal webhook failed with non-retryable status\"\n );\n return;\n }\n Err(e) => {\n attempt += 1;\n if attempt > self.config.max_retries {\n tracing::error!(\n error = %e,\n attempts = attempt,\n \"Terminal webhook failed after max retries\"\n );\n return;\n }\n\n let backoff = self.config.backoff_base * (1 << attempt.min(10));\n tracing::warn!(\n error = %e,\n attempt = attempt,\n backoff_ms = backoff.as_millis(),\n \"Terminal webhook request error, retrying\"\n );\n tokio::time::sleep(backoff).await;\n }\n }\n }\n }\n\n /// Send a terminal webhook synchronously (for Drop contexts).\n ///\n /// Uses ureq (blocking HTTP) instead of reqwest for non-async contexts.\n pub fn send_terminal_sync(&self, payload: &serde_json::Value) {\n if !self\n .config\n .events_filter\n .contains(&WebhookEventType::Completed)\n {\n return;\n }\n\n let agent = ureq::Agent::config_builder()\n .timeout_global(Some(Duration::from_secs(30)))\n .tls_config(\n ureq::tls::TlsConfig::builder()\n .root_certs(ureq::tls::RootCerts::PlatformVerifier)\n .build(),\n )\n .build()\n .new_agent();\n\n let auth_header = std::env::var(\"WEBHOOK_AUTH_TOKEN\")\n .ok()\n .map(|token| format!(\"Bearer {}\", token));\n\n let user_agent = format!(\"coglet/{}\", COGLET_VERSION);\n\n let mut attempt = 0;\n loop {\n let mut request = agent\n .post(&self.url)\n .header(\"Content-Type\", \"application/json\")\n .header(\"User-Agent\", &user_agent);\n\n if let Some(ref auth) = auth_header {\n request = request.header(\"Authorization\", auth);\n }\n\n if let Some(ref traceparent) = self.trace_context.traceparent {\n request = request.header(\"traceparent\", traceparent);\n }\n if let Some(ref tracestate) = self.trace_context.tracestate {\n request = request.header(\"tracestate\", tracestate);\n }\n\n let result = request.send_json(payload);\n\n match result {\n Ok(response) => {\n let status = response.status().as_u16();\n if (200..300).contains(&status) {\n tracing::debug!(status = %status, \"Terminal webhook (sync) sent successfully\");\n return;\n }\n\n if self.config.retry_status_codes.contains(&status) {\n attempt += 1;\n if attempt > self.config.max_retries {\n tracing::error!(\n status = %status,\n attempts = attempt,\n \"Terminal webhook (sync) failed after max retries\"\n );\n return;\n }\n\n let backoff = self.config.backoff_base * (1 << attempt.min(10));\n tracing::warn!(\n status = %status,\n attempt = attempt,\n backoff_ms = backoff.as_millis(),\n \"Terminal webhook (sync) failed, retrying\"\n );\n std::thread::sleep(backoff);\n continue;\n }\n\n tracing::error!(\n status = %status,\n \"Terminal webhook (sync) failed with non-retryable status\"\n );\n return;\n }\n Err(e) => {\n attempt += 1;\n if attempt > self.config.max_retries {\n tracing::error!(\n error = %e,\n attempts = attempt,\n \"Terminal webhook (sync) failed after max retries\"\n );\n return;\n }\n\n let backoff = self.config.backoff_base * (1 << attempt.min(10));\n tracing::warn!(\n error = %e,\n attempt = attempt,\n backoff_ms = backoff.as_millis(),\n \"Terminal webhook (sync) request error, retrying\"\n );\n std::thread::sleep(backoff);\n }\n }\n }\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use wiremock::matchers::{method, path};\n use wiremock::{Mock, MockServer, ResponseTemplate};\n\n #[test]\n fn config_defaults() {\n let config = WebhookConfig::default();\n assert_eq!(config.response_interval, Duration::from_millis(500));\n assert_eq!(config.max_retries, 12);\n assert!(config.events_filter.contains(&WebhookEventType::Start));\n assert!(config.events_filter.contains(&WebhookEventType::Completed));\n }\n\n #[test]\n fn event_is_terminal() {\n assert!(!WebhookEventType::Start.is_terminal());\n assert!(!WebhookEventType::Output.is_terminal());\n assert!(!WebhookEventType::Logs.is_terminal());\n assert!(WebhookEventType::Completed.is_terminal());\n }\n\n fn test_config() -> WebhookConfig {\n WebhookConfig {\n response_interval: Duration::ZERO,\n max_retries: 2,\n backoff_base: Duration::from_millis(1),\n ..Default::default()\n }\n }\n\n #[tokio::test]\n async fn send_terminal_posts_json() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(200))\n .expect(1)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let sender = WebhookSender::new(url, test_config()).unwrap();\n\n sender\n .send_terminal(\n WebhookEventType::Completed,\n &serde_json::json!({\"id\": \"pred_123\", \"status\": \"succeeded\"}),\n )\n .await;\n }\n\n #[tokio::test]\n async fn send_terminal_retries_on_500() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(500))\n .up_to_n_times(1)\n .mount(&server)\n .await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(200))\n .expect(1)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let sender = WebhookSender::new(url, test_config()).unwrap();\n\n sender\n .send_terminal(\n WebhookEventType::Completed,\n &serde_json::json!({\"status\": \"succeeded\"}),\n )\n .await;\n }\n\n #[tokio::test]\n async fn send_terminal_no_retry_on_400() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(400))\n .expect(1)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let sender = WebhookSender::new(url, test_config()).unwrap();\n\n sender\n .send_terminal(\n WebhookEventType::Completed,\n &serde_json::json!({\"status\": \"succeeded\"}),\n )\n .await;\n }\n\n #[tokio::test]\n async fn send_terminal_respects_filter() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(200))\n .expect(0)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let config = WebhookConfig {\n events_filter: [WebhookEventType::Start].into_iter().collect(),\n ..test_config()\n };\n let sender = WebhookSender::new(url, config).unwrap();\n\n sender\n .send_terminal(\n WebhookEventType::Completed,\n &serde_json::json!({\"status\": \"succeeded\"}),\n )\n .await;\n }\n\n #[tokio::test]\n async fn send_non_terminal_fires_and_forgets() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(200))\n .expect(1)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let sender = WebhookSender::new(url, test_config()).unwrap();\n\n sender.send(\n WebhookEventType::Start,\n &serde_json::json!({\"status\": \"starting\"}),\n );\n\n tokio::time::sleep(Duration::from_millis(50)).await;\n }\n\n #[tokio::test]\n async fn send_non_terminal_logs_throttled() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(200))\n .expect(1)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let config = WebhookConfig {\n response_interval: Duration::from_secs(10),\n ..test_config()\n };\n let sender = WebhookSender::new(url, config).unwrap();\n\n sender.send(\n WebhookEventType::Logs,\n &serde_json::json!({\"logs\": \"line 1\"}),\n );\n // Second send should be throttled\n sender.send(\n WebhookEventType::Logs,\n &serde_json::json!({\"logs\": \"line 2\"}),\n );\n\n tokio::time::sleep(Duration::from_millis(50)).await;\n }\n\n #[tokio::test]\n async fn send_output_not_throttled() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(200))\n .expect(2)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let config = WebhookConfig {\n response_interval: Duration::from_secs(10),\n ..test_config()\n };\n let sender = WebhookSender::new(url, config).unwrap();\n\n // Output events bypass throttling — both should be sent\n sender.send(\n WebhookEventType::Output,\n &serde_json::json!({\"output\": \"1\"}),\n );\n sender.send(\n WebhookEventType::Output,\n &serde_json::json!({\"output\": \"2\"}),\n );\n\n tokio::time::sleep(Duration::from_millis(50)).await;\n }\n\n #[tokio::test]\n async fn send_terminal_sync_posts_json() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(200))\n .expect(1)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let sender = WebhookSender::new(url, test_config()).unwrap();\n\n sender.send_terminal_sync(&serde_json::json!({\"id\": \"pred_123\", \"status\": \"succeeded\"}));\n }\n\n #[tokio::test]\n async fn send_terminal_sync_retries_on_500() {\n let server = MockServer::start().await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(500))\n .up_to_n_times(1)\n .mount(&server)\n .await;\n\n Mock::given(method(\"POST\"))\n .and(path(\"/webhook\"))\n .respond_with(ResponseTemplate::new(200))\n .expect(1)\n .mount(&server)\n .await;\n\n let url = format!(\"{}/webhook\", server.uri());\n let sender = WebhookSender::new(url, test_config()).unwrap();\n\n sender.send_terminal_sync(&serde_json::json!({\"status\": \"succeeded\"}));\n }\n}\n" }, { "path": "crates/coglet/src/worker.rs", "content": "//! Worker subprocess - runs inside the Python subprocess.\n//!\n//! This module provides the child-side of the worker subprocess protocol.\n//! The parent side (spawning, message routing) is in orchestrator.rs.\n//!\n//! Architecture:\n//! - Control channel (stdin/stdout): Cancel, Shutdown signals + Ready, Idle responses\n//! - Slot sockets: Prediction data + streaming logs\n//!\n//! Each slot runs predictions independently.\n\nuse std::collections::HashMap;\nuse std::io;\nuse std::path::PathBuf;\nuse std::sync::Arc;\nuse std::sync::OnceLock;\nuse std::sync::atomic::{AtomicUsize, Ordering};\n\nuse futures::{SinkExt, StreamExt};\nuse tokio::runtime::Handle;\nuse tokio::sync::mpsc;\nuse tokio_util::codec::{FramedRead, FramedWrite};\n\nuse crate::bridge::protocol::truncate_worker_log;\n\n// ============================================================================\n// Dropped log tracking\n// ============================================================================\n\n/// Counter for logs dropped due to channel backpressure during setup.\nstatic DROPPED_SETUP_LOG_COUNT: AtomicUsize = AtomicUsize::new(0);\n\n/// Increment the dropped log counter.\n/// Called by ControlChannelLogSender in coglet-python when try_send fails.\npub fn increment_dropped_log_count() {\n DROPPED_SETUP_LOG_COUNT.fetch_add(1, Ordering::Relaxed);\n}\n\n/// Report and reset dropped log count.\n/// Returns the number of logs dropped since last call.\nfn report_dropped_logs(tx: &mpsc::Sender, interval_millis: u64) {\n let dropped = DROPPED_SETUP_LOG_COUNT.swap(0, Ordering::Relaxed);\n if dropped > 0 {\n let _ = tx.try_send(ControlResponse::DroppedLogs {\n count: dropped,\n interval_millis,\n });\n }\n}\n\n// ============================================================================\n// Fatal worker shutdown\n// ============================================================================\n\nstruct FatalContext {\n tx: mpsc::Sender,\n}\n\nstatic FATAL_CONTEXT: OnceLock = OnceLock::new();\n\nfn init_fatal_context(tx: mpsc::Sender) {\n let _ = FATAL_CONTEXT.set(FatalContext { tx });\n}\n\n/// Install a panic hook that sends a Fatal IPC message and aborts.\n///\n/// Any panic in the worker is an invariant violation. The hook sends a best-effort\n/// `ControlResponse::Fatal` so the parent can poison all slots, then aborts.\n/// This means `.expect()` / `panic!()` at any call site automatically gets\n/// the correct fatal behavior — no special helpers needed.\nfn install_panic_hook() {\n let prev = std::panic::take_hook();\n std::panic::set_hook(Box::new(move |info| {\n // Run the default hook first (prints to stderr).\n prev(info);\n\n let msg = if let Some(s) = info.payload().downcast_ref::<&str>() {\n (*s).to_string()\n } else if let Some(s) = info.payload().downcast_ref::() {\n s.clone()\n } else {\n \"\".to_string()\n };\n\n let reason = match info.location() {\n Some(loc) => format!(\"panic at {}:{}: {}\", loc.file(), loc.line(), msg),\n None => format!(\"panic: {}\", msg),\n };\n\n if let Some(ctx) = FATAL_CONTEXT.get() {\n let _ = ctx.tx.try_send(ControlResponse::Fatal { reason });\n }\n\n // If panic=abort is not set, abort explicitly.\n std::process::abort();\n }));\n}\n\n// ============================================================================\n// Tracing initialization\n// ============================================================================\n\nfn init_worker_tracing(tx: mpsc::Sender) {\n use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};\n\n let filter = if std::env::var(\"RUST_LOG\").is_ok() {\n EnvFilter::from_default_env()\n } else {\n let base_level = match std::env::var(\"COG_LOG_LEVEL\").as_deref() {\n Ok(\"debug\") => \"debug\",\n Ok(\"warn\") | Ok(\"warning\") => \"warn\",\n Ok(\"error\") => \"error\",\n _ => \"info\",\n };\n\n let filter_str = format!(\n \"coglet={level},coglet::setup=info,coglet::user=info,coglet_worker={level},coglet_worker::schema=off,coglet_worker::protocol=off\",\n level = base_level\n );\n\n EnvFilter::new(filter_str)\n };\n\n let worker_layer = WorkerTracingLayer::new(tx);\n\n let subscriber = tracing_subscriber::registry()\n .with(filter)\n .with(worker_layer);\n\n let _ = subscriber.try_init();\n}\n\nuse crate::bridge::codec::JsonCodec;\nuse crate::bridge::protocol::{\n ControlRequest, ControlResponse, FileOutputKind, LogSource, MAX_INLINE_IPC_SIZE, MetricMode,\n SlotId, SlotOutcome, SlotRequest, SlotResponse,\n};\nuse crate::bridge::transport::{ChildTransportInfo, connect_transport};\nuse crate::orchestrator::HealthcheckResult;\nuse crate::worker_tracing_layer::WorkerTracingLayer;\n\ntype SlotWriter =\n Arc>>>;\n\n/// Handle for sending messages on a slot socket.\n///\n/// Used by log writers to stream logs during prediction. Thread-safe via\n/// tokio mpsc channel - logs are queued and written asynchronously.\n#[derive(Clone)]\npub struct SlotSender {\n tx: mpsc::UnboundedSender,\n output_dir: PathBuf,\n file_counter: Arc,\n}\n\nimpl SlotSender {\n pub fn new(tx: mpsc::UnboundedSender, output_dir: PathBuf) -> Self {\n Self {\n tx,\n output_dir,\n file_counter: Arc::new(AtomicUsize::new(0)),\n }\n }\n\n /// Generate a unique filename in the output dir.\n fn next_output_path(&self, extension: &str) -> PathBuf {\n let n = self.file_counter.fetch_add(1, Ordering::Relaxed);\n self.output_dir.join(format!(\"{n}.{extension}\"))\n }\n\n pub fn send_log(&self, source: LogSource, data: &str) -> io::Result<()> {\n if data.is_empty() {\n return Ok(());\n }\n\n let msg = SlotResponse::Log {\n source,\n data: truncate_worker_log(data.to_string()),\n };\n\n self.tx\n .send(msg)\n .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, \"slot channel closed\"))\n }\n\n /// Write raw bytes to a file in the output dir and send as FileOutput.\n ///\n /// Used by FFI workers (Python, Node, etc.) to hand off file data without\n /// needing language-specific file I/O — SlotSender owns the write.\n pub fn write_file_output(\n &self,\n data: &[u8],\n extension: &str,\n mime_type: Option,\n ) -> io::Result<()> {\n let path = self.next_output_path(extension);\n std::fs::write(&path, data)?;\n self.send_file_output(path, mime_type)\n }\n\n /// Send a file-typed output (e.g. Path, File return types).\n ///\n /// The file is already on disk at `path` — we just send the path reference.\n /// `mime_type` is an explicit MIME type; when None the parent guesses from extension.\n pub fn send_file_output(&self, path: PathBuf, mime_type: Option) -> io::Result<()> {\n let filename = path\n .to_str()\n .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, \"non-UTF-8 path\"))?\n .to_string();\n let msg = SlotResponse::FileOutput {\n filename,\n kind: FileOutputKind::FileType,\n mime_type,\n };\n self.tx\n .send(msg)\n .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, \"slot channel closed\"))\n }\n\n /// Send a user metric to the parent process.\n pub fn send_metric(\n &self,\n name: String,\n value: serde_json::Value,\n mode: MetricMode,\n ) -> io::Result<()> {\n let msg = SlotResponse::Metric { name, value, mode };\n self.tx\n .send(msg)\n .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, \"slot channel closed\"))\n }\n\n /// Send prediction output, either inline or spilled to disk if too large.\n pub fn send_output(&self, output: serde_json::Value) -> io::Result<()> {\n let msg = build_output_message(&self.output_dir, output)?;\n self.tx\n .send(msg)\n .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, \"slot channel closed\"))\n }\n}\n\n/// Build an output message, spilling to disk if larger than the IPC frame limit.\nfn build_output_message(\n output_dir: &std::path::Path,\n output: serde_json::Value,\n) -> io::Result {\n let serialized =\n serde_json::to_vec(&output).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;\n\n if serialized.len() > MAX_INLINE_IPC_SIZE {\n let path = output_dir.join(format!(\"spill_{}.json\", uuid::Uuid::new_v4()));\n std::fs::write(&path, &serialized)?;\n let filename = path\n .to_str()\n .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, \"non-UTF-8 path\"))?\n .to_string();\n Ok(SlotResponse::FileOutput {\n filename,\n kind: FileOutputKind::Oversized,\n mime_type: None,\n })\n } else {\n Ok(SlotResponse::Output { output })\n }\n}\n\n/// Setup phase errors.\n///\n/// These errors occur during predictor loading and setup, before predictions\n/// can run. They affect health status (SETUP_FAILED) rather than HTTP status.\n#[derive(Debug, thiserror::Error)]\npub enum SetupError {\n /// Failed to import or instantiate the predictor class.\n #[error(\"failed to load predictor: {message}\")]\n Load { message: String },\n\n /// The setup() method raised an exception.\n #[error(\"setup failed: {message}\")]\n Setup { message: String },\n\n /// Internal error (e.g., GIL acquisition failed).\n #[error(\"internal error: {message}\")]\n Internal { message: String },\n}\n\nimpl SetupError {\n pub fn load(message: impl Into) -> Self {\n Self::Load {\n message: message.into(),\n }\n }\n\n pub fn setup(message: impl Into) -> Self {\n Self::Setup {\n message: message.into(),\n }\n }\n\n pub fn internal(message: impl Into) -> Self {\n Self::Internal {\n message: message.into(),\n }\n }\n}\n\n/// Trait for the prediction handler - abstracts the Python integration.\n#[async_trait::async_trait]\npub trait PredictHandler: Send + Sync + 'static {\n /// Initialize the predictor (load model, run setup).\n async fn setup(&self) -> Result<(), SetupError>;\n\n /// Run a prediction.\n async fn predict(\n &self,\n slot: SlotId,\n id: String,\n input: serde_json::Value,\n slot_sender: Arc,\n context: std::collections::HashMap,\n ) -> PredictResult;\n\n /// Request cancellation of prediction on a slot.\n fn cancel(&self, slot: SlotId);\n\n /// Run user-defined healthcheck. Default: healthy.\n async fn healthcheck(&self) -> HealthcheckResult {\n HealthcheckResult::healthy()\n }\n}\n\n/// Path to the pre-built OpenAPI schema file inside the container.\n/// Written during `cog build` and COPYed into the image.\nconst BUNDLED_SCHEMA_PATH: &str = \".cog/openapi_schema.json\";\n\n/// Load the bundled OpenAPI schema from disk.\n///\n/// Returns `Some(schema)` if the file exists and parses correctly.\n/// Returns `None` if missing or unparseable — the predictor will accept\n/// any input without schema validation.\nfn load_bundled_schema() -> Option {\n let path = std::path::Path::new(BUNDLED_SCHEMA_PATH);\n match std::fs::read_to_string(path) {\n Ok(contents) => match serde_json::from_str(&contents) {\n Ok(schema) => {\n tracing::info!(\"Loaded OpenAPI schema from {}\", BUNDLED_SCHEMA_PATH);\n Some(schema)\n }\n Err(e) => {\n tracing::warn!(\n \"Failed to parse {}: {}. Running without schema — all input types accepted.\",\n BUNDLED_SCHEMA_PATH,\n e,\n );\n None\n }\n },\n Err(_) => {\n tracing::warn!(\n \"No schema file at {}. Running without schema — all input types accepted. \\\n Rebuild with a recent version of cog to generate the schema.\",\n BUNDLED_SCHEMA_PATH,\n );\n None\n }\n }\n}\n\n/// The outcome of a prediction\n#[derive(Debug, Clone, PartialEq)]\npub enum PredictionOutcome {\n /// Prediction completed successfully\n Success {\n output: serde_json::Value,\n predict_time: f64,\n /// True when the predictor returned a list, generator, or iterator.\n is_stream: bool,\n },\n /// Prediction failed with an error\n Failed { error: String, predict_time: f64 },\n /// Prediction was cancelled\n Cancelled { predict_time: f64 },\n}\n\n#[derive(Debug)]\npub struct PredictResult {\n pub outcome: PredictionOutcome,\n}\n\nimpl PredictResult {\n pub fn success(output: serde_json::Value, predict_time: f64, is_stream: bool) -> Self {\n Self {\n outcome: PredictionOutcome::Success {\n output,\n predict_time,\n is_stream,\n },\n }\n }\n\n pub fn failed(error: String, predict_time: f64) -> Self {\n Self {\n outcome: PredictionOutcome::Failed {\n error,\n predict_time,\n },\n }\n }\n\n pub fn cancelled(predict_time: f64) -> Self {\n Self {\n outcome: PredictionOutcome::Cancelled { predict_time },\n }\n }\n}\n\n/// Callback for setup log routing.\n///\n/// Called before setup() with a sender for routing logs to the control channel.\n/// Returns a cleanup function that unregisters the sender.\npub type SetupLogHook =\n Box) -> Box + Send>;\n\npub struct WorkerConfig {\n pub num_slots: usize,\n /// Hook for setup log routing. Called before setup() to register a log sender.\n pub setup_log_hook: Option,\n}\n\nimpl Default for WorkerConfig {\n fn default() -> Self {\n Self {\n num_slots: 1,\n setup_log_hook: None,\n }\n }\n}\n\nstruct SlotCompletion {\n outcome: SlotOutcome,\n}\n\nimpl SlotCompletion {\n fn idle(slot: SlotId) -> Self {\n Self {\n outcome: SlotOutcome::idle(slot),\n }\n }\n\n fn poisoned(slot: SlotId, error: impl Into) -> Self {\n Self {\n outcome: SlotOutcome::poisoned(slot, error),\n }\n }\n}\n\n/// Run the worker event loop.\n///\n/// Connects to slot sockets, runs setup, then processes predictions.\n/// Reads control messages from stdin, prediction requests from slot sockets.\npub async fn run_worker(\n handler: Arc,\n config: WorkerConfig,\n transport_info: ChildTransportInfo,\n) -> io::Result<()> {\n let num_slots = config.num_slots;\n\n let (setup_log_tx, mut setup_log_rx) = mpsc::channel::(5000);\n\n init_worker_tracing(setup_log_tx.clone());\n\n // CRITICAL: Redirect fds BEFORE any FFI initialization to prevent subprocesses\n // from polluting the control channel\n let control_fds =\n crate::fd_redirect::redirect_fds_for_subprocess_isolation(setup_log_tx.clone())?;\n\n // Connect to slot sockets (transport info from Init message)\n tracing::trace!(?transport_info, \"Connecting to slot transport\");\n let mut transport = connect_transport(transport_info).await?;\n tracing::info!(num_slots, \"Connected to slot transport\");\n\n // Control channel via redirected fds (not stdin/stdout)\n let ctrl_stdin = tokio::fs::File::from_std(control_fds.stdin_fd.into());\n let ctrl_stdout = tokio::fs::File::from_std(control_fds.stdout_fd.into());\n\n let mut ctrl_reader = FramedRead::new(ctrl_stdin, JsonCodec::::new());\n let ctrl_writer = Arc::new(tokio::sync::Mutex::new(FramedWrite::new(\n ctrl_stdout,\n JsonCodec::::new(),\n )));\n\n // Generate unique SlotIds for each socket\n let slot_ids: Vec = (0..num_slots).map(|_| SlotId::new()).collect();\n\n init_fatal_context(setup_log_tx.clone());\n install_panic_hook();\n\n let setup_cleanup = config.setup_log_hook.map(|hook| hook(setup_log_tx.clone()));\n\n // Forward logs to control channel (runs for entire worker lifetime)\n // Receives logs from both Python (during setup) and fd_redirect capture threads (always)\n let ctrl_writer_for_logs = Arc::clone(&ctrl_writer);\n let _log_forwarder = tokio::spawn(async move {\n let mut log_count = 0;\n let mut total_bytes = 0;\n while let Some(msg) = setup_log_rx.recv().await {\n if let ControlResponse::Log { ref data, .. } = msg {\n let msg_size = data.len();\n log_count += 1;\n total_bytes += msg_size;\n tracing::trace!(\n log_number = log_count,\n msg_size_bytes = msg_size,\n total_bytes,\n \"Forwarding log\"\n );\n }\n let mut w = ctrl_writer_for_logs.lock().await;\n if let Err(e) = w.send(msg).await {\n tracing::warn!(\n error = %e,\n log_count,\n total_bytes,\n \"Failed to forward log\"\n );\n break;\n }\n }\n tracing::debug!(\n total_logs = log_count,\n total_bytes,\n total_kb = total_bytes / 1024,\n \"Log forwarder exiting\"\n );\n });\n\n // Periodic reporter for dropped logs (runs for entire worker lifetime)\n let dropped_log_tx = setup_log_tx.clone();\n let _dropped_log_reporter = tokio::spawn(async move {\n let mut interval = tokio::time::interval(std::time::Duration::from_millis(5000));\n loop {\n interval.tick().await;\n report_dropped_logs(&dropped_log_tx, 5000);\n }\n });\n\n // Run setup\n tracing::info!(\"Worker starting setup\");\n let setup_start = std::time::Instant::now();\n let setup_result = handler.setup().await;\n let setup_elapsed = setup_start.elapsed();\n tracing::debug!(\n elapsed_ms = setup_elapsed.as_millis() as u64,\n success = setup_result.is_ok(),\n \"Setup handler returned\"\n );\n\n // Unregister Python's setup sender, but keep log_forwarder running\n // The fd_redirect capture threads will continue sending subprocess logs\n if let Some(cleanup) = setup_cleanup {\n tracing::debug!(\"Running cleanup (unregistering Python setup sender)\");\n cleanup();\n }\n // Note: We DON'T drop setup_log_tx or wait for log_forwarder\n // The log_forwarder continues running to forward subprocess output throughout worker lifetime\n\n // Handle setup failure\n if let Err(e) = setup_result {\n tracing::error!(\n error = %e,\n elapsed_ms = setup_elapsed.as_millis() as u64,\n \"Setup failed\"\n );\n let slot = slot_ids.first().copied().unwrap_or_else(SlotId::new);\n let mut w = ctrl_writer.lock().await;\n let _ = w\n .send(ControlResponse::Failed {\n slot,\n error: format!(\"Setup failed: {}\", e),\n })\n .await;\n return Ok(());\n }\n\n // Load the pre-built schema from .cog/openapi_schema.json (written during `cog build`).\n // No runtime generation — if the file doesn't exist, no schema.\n let schema = load_bundled_schema();\n if let Some(ref s) = schema {\n let schema_json = serde_json::to_string(s).unwrap_or_else(|_| \"{}\".to_string());\n let schema_size = schema_json.len();\n tracing::info!(\n schema_size_bytes = schema_size,\n schema_size_kb = schema_size / 1024,\n \"Schema loaded\"\n );\n if schema_size > 1024 * 1024 {\n // Log first 500 chars if schema is >1MB\n tracing::warn!(\n schema_preview = &schema_json[..500.min(schema_json.len())],\n \"Large schema detected\"\n );\n }\n }\n tracing::trace!(num_slots, ?slot_ids, \"Sending Ready to parent\");\n {\n let mut w = ctrl_writer.lock().await;\n w.send(ControlResponse::Ready {\n slots: slot_ids.clone(),\n schema,\n })\n .await?;\n }\n\n // Channel for slot completions\n let (completion_tx, mut completion_rx) = mpsc::channel::(num_slots);\n\n // Track slot state\n let mut slot_busy: HashMap = slot_ids.iter().map(|id| (*id, false)).collect();\n let mut slot_poisoned: HashMap = slot_ids.iter().map(|id| (*id, false)).collect();\n\n // Set up slot socket readers/writers\n let sockets = transport.drain_sockets();\n let mut slot_readers: HashMap<\n SlotId,\n FramedRead>,\n > = HashMap::new();\n let mut slot_writers: HashMap<\n SlotId,\n FramedWrite>,\n > = HashMap::new();\n\n for (slot_id, socket) in slot_ids.iter().zip(sockets) {\n let (read_half, write_half) = socket.into_split();\n slot_readers.insert(*slot_id, FramedRead::new(read_half, JsonCodec::new()));\n slot_writers.insert(*slot_id, FramedWrite::new(write_half, JsonCodec::new()));\n }\n\n // Channel for incoming slot requests\n let (request_tx, mut request_rx) = mpsc::channel::<(SlotId, SlotRequest)>(num_slots);\n\n // Spawn reader task for each slot\n for (slot_id, reader) in slot_readers {\n let tx = request_tx.clone();\n tokio::spawn(async move {\n slot_reader_task(slot_id, reader, tx).await;\n });\n }\n drop(request_tx);\n\n // Wrap writers for sharing with prediction tasks\n let slot_writers: HashMap = slot_writers\n .into_iter()\n .map(|(id, w)| (id, Arc::new(tokio::sync::Mutex::new(w))))\n .collect();\n\n // Main event loop\n loop {\n tokio::select! {\n biased;\n\n ctrl_msg = ctrl_reader.next() => {\n match ctrl_msg {\n Some(Ok(ControlRequest::Init { .. })) => {\n tracing::warn!(\"Received Init in event loop (should be at startup)\");\n }\n Some(Ok(ControlRequest::Cancel { slot })) => {\n tracing::trace!(%slot, \"Cancel requested\");\n handler.cancel(slot);\n }\n Some(Ok(ControlRequest::Shutdown)) => {\n tracing::info!(\"Shutdown requested\");\n let mut w = ctrl_writer.lock().await;\n let _ = w.send(ControlResponse::ShuttingDown).await;\n break;\n }\n Some(Ok(ControlRequest::Healthcheck { id })) => {\n tracing::trace!(%id, \"Healthcheck requested, invoking handler\");\n let result = handler.healthcheck().await;\n tracing::trace!(\n %id,\n status = ?result.status,\n error = ?result.error,\n \"Healthcheck handler returned\"\n );\n let mut w = ctrl_writer.lock().await;\n let _ = w.send(ControlResponse::HealthcheckResult {\n id,\n status: result.status,\n error: result.error,\n }).await;\n }\n Some(Err(e)) => {\n tracing::error!(error = %e, \"Control channel error\");\n break;\n }\n None => {\n tracing::error!(\"Control channel closed (parent died?), exiting\");\n break;\n }\n }\n }\n\n Some(completion) = completion_rx.recv() => {\n let slot = completion.outcome.slot_id();\n slot_busy.insert(slot, false);\n\n if completion.outcome.is_poisoned() {\n slot_poisoned.insert(slot, true);\n }\n {\n let mut w = ctrl_writer.lock().await;\n let _ = w.send(completion.outcome.into_control_response()).await;\n }\n\n if slot_poisoned.values().all(|&p| p) {\n tracing::error!(\"All slots poisoned, exiting\");\n break;\n }\n }\n\n Some((slot_id, request)) = request_rx.recv() => {\n if slot_busy.get(&slot_id).copied().unwrap_or(false) {\n tracing::warn!(%slot_id, \"Request received for busy slot, ignoring\");\n continue;\n }\n if slot_poisoned.get(&slot_id).copied().unwrap_or(false) {\n tracing::warn!(%slot_id, \"Request received for poisoned slot, ignoring\");\n continue;\n }\n\n // Extract the prediction ID before consuming the request, so we\n // can report a failure even if rehydration itself fails.\n let prediction_id = request.prediction_id().to_string();\n\n match request.rehydrate_input() {\n Ok((id, input, output_dir, context)) => {\n tracing::trace!(%slot_id, %id, \"Prediction request received\");\n slot_busy.insert(slot_id, true);\n\n let writer = match slot_writers.get(&slot_id) {\n Some(w) => Arc::clone(w),\n None => {\n tracing::error!(%slot_id, \"No writer for slot\");\n continue;\n }\n };\n\n let handler = Arc::clone(&handler);\n let completion_tx = completion_tx.clone();\n tokio::spawn(async move {\n let completion = run_prediction(\n slot_id,\n id,\n input,\n PathBuf::from(output_dir),\n handler,\n writer,\n context,\n ).await;\n let _ = completion_tx.send(completion).await;\n });\n }\n Err(e) => {\n tracing::error!(%slot_id, %prediction_id, error = %e, \"Failed to rehydrate input\");\n // Send a failure response so the prediction doesn't hang forever.\n if let Some(writer) = slot_writers.get(&slot_id) {\n let mut w = writer.lock().await;\n let fail_msg = SlotResponse::Failed {\n id: prediction_id,\n error: format!(\"Failed to rehydrate input: {}\", e),\n };\n if let Err(send_err) = w.send(fail_msg).await {\n tracing::error!(%slot_id, error = %send_err, \"Failed to send rehydrate error response\");\n }\n }\n let _ = completion_tx.send(SlotCompletion::idle(slot_id)).await;\n }\n }\n }\n }\n }\n\n tracing::info!(\"Worker exiting\");\n Ok(())\n}\n\nasync fn slot_reader_task(\n slot_id: SlotId,\n mut reader: FramedRead>,\n tx: mpsc::Sender<(SlotId, SlotRequest)>,\n) {\n loop {\n match reader.next().await {\n Some(Ok(request)) => {\n if tx.send((slot_id, request)).await.is_err() {\n break;\n }\n }\n Some(Err(e)) => {\n tracing::error!(%slot_id, error = %e, \"Slot reader error\");\n break;\n }\n None => {\n tracing::trace!(%slot_id, \"Slot socket closed\");\n break;\n }\n }\n }\n}\n\nasync fn run_prediction(\n slot_id: SlotId,\n prediction_id: String,\n input: serde_json::Value,\n output_dir: PathBuf,\n handler: Arc,\n writer: SlotWriter,\n context: std::collections::HashMap,\n) -> SlotCompletion {\n tracing::trace!(%slot_id, %prediction_id, \"run_prediction starting\");\n\n // Create channel for log streaming\n let (log_tx, mut log_rx) = mpsc::unbounded_channel::();\n let slot_sender = Arc::new(SlotSender::new(log_tx, output_dir.clone()));\n\n // Forward logs to slot socket\n let writer_for_logs = Arc::clone(&writer);\n let log_forwarder = tokio::spawn(async move {\n while let Some(msg) = log_rx.recv().await {\n let mut w = writer_for_logs.lock().await;\n if let Err(e) = w.send(msg).await {\n tracing::warn!(error = %e, \"Failed to forward log\");\n break;\n }\n }\n tracing::trace!(\"Prediction log forwarder exiting\");\n });\n\n // Run prediction — slot_sender is moved in, dropped when predict returns,\n // which closes the log channel and lets the log forwarder exit.\n //\n // block_in_place tells tokio this thread will block (Python GIL acquisition),\n // allowing the runtime to move other tasks (like log_forwarder) to free\n // threads. Without this, the log forwarder can be work-stolen onto the\n // same thread as the prediction and starved until predict returns, causing\n // all logs to arrive in a single batch at prediction end.\n let result = tokio::task::block_in_place(|| {\n Handle::current().block_on(handler.predict(\n slot_id,\n prediction_id.clone(),\n input,\n slot_sender,\n context,\n ))\n });\n tracing::trace!(%slot_id, %prediction_id, \"handler.predict returned\");\n\n // Wait for log forwarder\n tracing::trace!(%slot_id, %prediction_id, \"Waiting for log forwarder\");\n let _ = log_forwarder.await;\n tracing::trace!(%slot_id, %prediction_id, \"Log forwarder done\");\n\n // Send result on slot socket.\n // Output is always sent separately from Done so that large values get\n // spilled to disk and never exceed the IPC frame limit.\n let mut w = writer.lock().await;\n let response = match result.outcome {\n PredictionOutcome::Success {\n output,\n predict_time,\n is_stream,\n } => {\n // Send output as a separate message (handles spilling for large values).\n // Skip if null or empty array — those mean \"already streamed\" (generators).\n if !output.is_null() && output != serde_json::Value::Array(vec![]) {\n let output_msg = match build_output_message(&output_dir, output) {\n Ok(msg) => msg,\n Err(e) => {\n tracing::error!(error = %e, \"Failed to build output message\");\n return SlotCompletion::poisoned(\n slot_id,\n format!(\"Output spill error: {}\", e),\n );\n }\n };\n if let Err(e) = w.send(output_msg).await {\n tracing::error!(error = %e, \"Failed to send prediction output\");\n return SlotCompletion::poisoned(slot_id, format!(\"Socket write error: {}\", e));\n }\n }\n SlotResponse::Done {\n id: prediction_id.clone(),\n output: None,\n predict_time,\n is_stream,\n }\n }\n PredictionOutcome::Cancelled { .. } => SlotResponse::Cancelled {\n id: prediction_id.clone(),\n },\n PredictionOutcome::Failed { error, .. } => SlotResponse::Failed {\n id: prediction_id.clone(),\n error,\n },\n };\n\n if let Err(e) = w.send(response).await {\n tracing::error!(error = %e, \"Failed to send prediction response\");\n return SlotCompletion::poisoned(slot_id, format!(\"Socket write error: {}\", e));\n }\n\n SlotCompletion::idle(slot_id)\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn predict_result_success() {\n let r = PredictResult::success(serde_json::json!(\"hello\"), 0.5, false);\n assert!(matches!(r.outcome, PredictionOutcome::Success { .. }));\n }\n\n #[test]\n fn predict_result_success_stream() {\n let r = PredictResult::success(serde_json::json!([]), 0.5, true);\n assert!(matches!(\n r.outcome,\n PredictionOutcome::Success {\n is_stream: true,\n ..\n }\n ));\n }\n\n #[test]\n fn predict_result_failed() {\n let r = PredictResult::failed(\"oops\".into(), 0.5);\n assert!(matches!(\n r.outcome,\n PredictionOutcome::Failed { ref error, .. } if error == \"oops\"\n ));\n }\n\n #[test]\n fn predict_result_cancelled() {\n let r = PredictResult::cancelled(0.5);\n assert!(matches!(r.outcome, PredictionOutcome::Cancelled { .. }));\n }\n\n #[test]\n fn worker_config_default() {\n let config = WorkerConfig::default();\n assert_eq!(config.num_slots, 1);\n }\n}\n" }, { "path": "crates/coglet/src/worker_tracing_layer.rs", "content": "//! Custom tracing layer for worker subprocess.\n//!\n//! Ships structured tracing events over IPC to orchestrator, preserving target and level.\n//! Optionally writes to fd 101 for direct debugging (controlled by RUST_WORKER_DIRECT_LOG=1).\n\nuse std::io::Write;\nuse std::sync::{Arc, Mutex};\n\nuse tokio::sync::mpsc;\nuse tracing::{Level, Subscriber};\nuse tracing_subscriber::layer::{Context, Layer};\n\nuse crate::bridge::protocol::{ControlResponse, truncate_worker_log};\n\npub struct WorkerTracingLayer {\n tx: mpsc::Sender,\n direct_log_fd: Option>>,\n}\n\nimpl WorkerTracingLayer {\n pub fn new(tx: mpsc::Sender) -> Self {\n let direct_log_fd = if std::env::var(\"RUST_WORKER_DIRECT_LOG\").as_deref() == Ok(\"1\") {\n // fd 101 is the original stderr preserved during fd_redirect\n let fd = unsafe { std::fs::File::from_raw_fd(101) };\n Some(Arc::new(Mutex::new(fd)))\n } else {\n None\n };\n\n Self { tx, direct_log_fd }\n }\n\n fn level_to_string(level: &Level) -> &'static str {\n match *level {\n Level::TRACE => \"trace\",\n Level::DEBUG => \"debug\",\n Level::INFO => \"info\",\n Level::WARN => \"warn\",\n Level::ERROR => \"error\",\n }\n }\n}\n\nimpl Layer for WorkerTracingLayer\nwhere\n S: Subscriber,\n{\n fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {\n let metadata = event.metadata();\n let target = metadata.target();\n let level = Self::level_to_string(metadata.level());\n\n let mut visitor = MessageVisitor::default();\n event.record(&mut visitor);\n let message = truncate_worker_log(visitor.message);\n\n // Targets excluded from IPC:\n // - coglet::bridge::codec: feedback loop when encoding WorkerLog messages\n // - coglet::worker_local: diagnostics that must stay on the worker process\n let is_local_only = target.starts_with(\"coglet::bridge::codec\")\n || target.starts_with(\"coglet::worker_local\");\n\n if !is_local_only {\n let _ = self.tx.try_send(ControlResponse::WorkerLog {\n target: target.to_string(),\n level: level.to_string(),\n message: message.clone(),\n });\n }\n\n // Write to preserved stderr (fd 101) for:\n // - worker_local targets (always, these are worker-only diagnostics)\n // - all targets when RUST_WORKER_DIRECT_LOG=1 is set\n if let Some(ref fd) = self.direct_log_fd\n && let Ok(mut file) = fd.lock()\n {\n let _ = writeln!(file, \"worker::{} [{}] {}\", target, level, message);\n } else if is_local_only {\n // No direct_log_fd but this is a local-only event — write to fd 101 directly.\n // Safety: fd 101 is the preserved original stderr from fd_redirect.\n #[cfg(unix)]\n {\n use std::os::unix::io::FromRawFd;\n let mut file = unsafe { std::fs::File::from_raw_fd(101) };\n let _ = writeln!(file, \"worker::{} [{}] {}\", target, level, message);\n std::mem::forget(file); // Don't close fd 101\n }\n }\n }\n}\n\n#[derive(Default)]\nstruct MessageVisitor {\n message: String,\n}\n\nimpl tracing::field::Visit for MessageVisitor {\n fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {\n if field.name() == \"message\" {\n self.message = format!(\"{:?}\", value);\n if self.message.starts_with('\"') && self.message.ends_with('\"') {\n self.message = self.message[1..self.message.len() - 1].to_string();\n }\n }\n }\n\n fn record_str(&mut self, field: &tracing::field::Field, value: &str) {\n if field.name() == \"message\" {\n self.message = value.to_string();\n }\n }\n}\n\n#[cfg(unix)]\nuse std::os::unix::io::FromRawFd;\n" }, { "path": "crates/coglet-python/Cargo.toml", "content": "[package]\nname = \"coglet-python\"\nversion.workspace = true\nedition.workspace = true\nlicense.workspace = true\npublish = false # Published to PyPI as 'coglet' wheel, not crates.io\nbuild = \"build.rs\"\n\n[lib]\nname = \"coglet\"\ncrate-type = [\"cdylib\", \"rlib\"]\n\n[dependencies]\nasync-trait = \"0.1.89\"\nbase64 = \"0.22\"\ncoglet_core = { path = \"../coglet\", package = \"coglet\" }\nfutures.workspace = true\npyo3.workspace = true\npyo3-async-runtimes.workspace = true\npyo3-stub-gen.workspace = true\nserde_json.workspace = true\ntokio.workspace = true\ntokio-util = { workspace = true, features = [\"codec\"] }\ntracing.workspace = true\ntracing-subscriber.workspace = true\n\n[target.'cfg(unix)'.dependencies]\nlibc = \"0.2\"\n\n[dev-dependencies]\npyo3 = { workspace = true, features = [\"auto-initialize\"] }\n\n[features]\nextension-module = [\"pyo3/extension-module\"]\n" }, { "path": "crates/coglet-python/README.md", "content": "# coglet-python\n\nPyO3 bindings that bridge the Rust coglet library to Python. This crate implements\nthe `PredictHandler` trait by wrapping Python predictor classes.\n\n## Overview\n\n```\n┌─────────────────────────────────────────────────────────────────────────────┐\n│ coglet-python │\n│ │\n│ ┌──────────────────────────────────────────────────────────────────────┐ │\n│ │ lib.rs │ │\n│ │ Python module: serve(), active(), _run_worker(), _is_cancelable() │ │\n│ └──────────────────────────────────────────────────────────────────────┘ │\n│ │ │\n│ ┌───────────────────────┼───────────────────────┐ │\n│ ▼ ▼ ▼ │\n│ ┌─────────────────────┐ ┌─────────────────────┐ ┌──────────────────┐ │\n│ │ worker_bridge.rs │ │ predictor.rs │ │ log_writer.rs │ │\n│ │ PredictHandler │ │ PythonPredictor │ │ SlotLogWriter │ │\n│ │ impl for Python │ │ load/setup/predict │ │ ContextVar │ │\n│ └─────────────────────┘ └─────────────────────┘ └──────────────────┘ │\n│ │ │ │ │\n│ │ │ │ │\n│ ▼ ▼ ▼ │\n│ ┌─────────────────────┐ ┌─────────────────────┐ ┌──────────────────┐ │\n│ │ input.rs │ │ output.rs │ │ audit.rs │ │\n│ │ Pydantic/ADT │ │ JSON serialization │ │ TeeWriter │ │\n│ │ input processing │ │ make_encodeable │ │ stream protect │ │\n│ └─────────────────────┘ └─────────────────────┘ └──────────────────┘ │\n│ │\n│ ┌──────────────────────────────────────────────────────────────────────┐ │\n│ │ cancel.rs │ │\n│ │ SIGUSR1 handling, CancelableGuard, KeyboardInterrupt injection │ │\n│ └──────────────────────────────────────────────────────────────────────┘ │\n│ │\n└─────────────────────────────────────────────────────────────────────────────┘\n```\n\n## Directory Structure\n\n```\ncoglet-python/\n├── Cargo.toml\n├── coglet.pyi # Type stubs for Python IDE support\n└── src/\n ├── lib.rs # Python module definition, serve/active/_run_worker\n ├── predictor.rs # PythonPredictor: wraps Python Predictor class\n ├── worker_bridge.rs # PythonPredictHandler: implements PredictHandler\n ├── input.rs # Input processing (Pydantic validation, ADT)\n ├── output.rs # Output processing (make_encodeable, upload_files)\n ├── log_writer.rs # SlotLogWriter, ContextVar routing, SetupLogSender\n ├── audit.rs # Audit hook, TeeWriter for stream protection\n └── cancel.rs # Cancellation: SIGUSR1, CancelableGuard\n```\n\n## Critical Concepts\n\n### The `active()` Flag\n\n```python\nimport coglet\n\nif coglet.server.active:\n # Running inside worker subprocess\n # stdout/stderr are captured, print goes to slot routing\nelse:\n # Running in parent or standalone\n # Normal stdout/stderr behavior\n```\n\nSet to `True` at the start of `_run_worker()`. Used by user code and cog internals\nto detect worker context.\n\n### Single Async Event Loop\n\nAsync predictors run on a **single** Python asyncio event loop, created at worker\nstartup. All slots share this loop.\n\n```\nWorker Subprocess\n┌─────────────────────────────────────────────────────────────┐\n│ Tokio Runtime (Rust) │\n│ └─ run_worker event loop │\n│ └─ For each SlotRequest::Predict: │\n│ └─ tokio::spawn prediction task │\n│ └─ Python::attach (acquire GIL) │\n│ └─ asyncio.run_coroutine_threadsafe() │\n│ └─ Predictor.predict() coroutine │\n└─────────────────────────────────────────────────────────────┘\n\nasyncio event loop (Python)\n┌─────────────────────────────────────────────────────────────┐\n│ Single event loop, started once at worker init │\n│ - concurrent.futures.Future per async prediction │\n│ - ContextVar propagates prediction_id to spawned tasks │\n│ - Cancellation via future.cancel() │\n└─────────────────────────────────────────────────────────────┘\n```\n\n**Why single loop?** \n- Python asyncio has one event loop per thread\n- We use `run_coroutine_threadsafe` to submit from Rust/Tokio\n- Multiple slots can have concurrent predictions (up to `max_concurrency`)\n\n### Prediction Execution\n\n**Sync Predictors:**\n```\nSlotRequest::Predict arrives\n │\n ├─▶ Python::attach (acquire GIL)\n ├─▶ set_sync_prediction_id(id) # For log routing\n ├─▶ predictor.predict(input) # Blocking call\n ├─▶ set_sync_prediction_id(None)\n └─▶ Return PredictResult\n```\n\n**Async Predictors:**\n```\nSlotRequest::Predict arrives\n │\n ├─▶ Python::attach (acquire GIL)\n ├─▶ Create wrapped coroutine:\n │ async def _ctx_wrapper(coro, prediction_id, contextvar):\n │ contextvar.set(prediction_id) # Set in this task's context\n │ return await coro\n │\n ├─▶ asyncio.run_coroutine_threadsafe(wrapper, loop)\n ├─▶ py.detach() (release GIL)\n ├─▶ future.result() (block Rust task, Python runs)\n └─▶ Return PredictResult\n```\n\n### STDOUT/STDERR Routing\n\nAll output from user code must be captured and routed through the slot socket.\n\n**Architecture:**\n```\nsys.stdout = SlotLogWriter(stdout)\nsys.stderr = SlotLogWriter(stderr)\n\nSlotLogWriter.write(data)\n │\n ├─▶ Get current prediction_id from:\n │ 1. SYNC_PREDICTION_ID static (for sync predictors)\n │ 2. ContextVar (for async predictors/spawned tasks)\n │\n ├─▶ Look up SlotSender in PREDICTION_REGISTRY\n │\n └─▶ Route:\n Found sender → slot_sender.send_log(source, data)\n No sender → Check setup sender (during setup)\n Neither → Log as orphan to stderr\n```\n\n**Line Buffering:**\nSlotLogWriter buffers writes until a newline. This coalesces Python's `print()`\nwhich does separate writes for content and `\\n`.\n\n### Audit Hook Protection\n\nUser code might replace `sys.stdout`:\n```python\nsys.stdout = open(\"mylog.txt\", \"w\")\n```\n\nWe can't prevent this, but we can intercept it with a Python audit hook.\n\n**Strategy: TeeWriter**\n```\nUser replaces sys.stdout\n │\n ├─▶ Audit hook fires on object.__setattr__(sys, \"stdout\", value)\n │\n ├─▶ Check: is value already SlotLogWriter? → Allow (it's us)\n │\n ├─▶ Check: is value already TeeWriter? → Allow (already wrapped)\n │\n ├─▶ Create TeeWriter(inner=SlotLogWriter, user_stream=value)\n │\n └─▶ Schedule: sys.stdout = tee (via Timer to avoid recursion)\n\nTeeWriter.write(data)\n │\n ├─▶ inner.write(data) # Our SlotLogWriter (routing works)\n └─▶ user_stream.write(data) # User's stream (their code works)\n```\n\n**Result:** Both our log routing AND the user's stream receive the data.\n\n### Cancellation\n\n**Sync Predictors:**\n```\nParent: ControlRequest::Cancel { slot }\n │\n ├─▶ Worker: handler.cancel(slot)\n │ └─▶ Set CANCEL_REQUESTED flag for slot\n │\n ├─▶ Worker: send SIGUSR1 to self\n │\n └─▶ Signal handler: raise KeyboardInterrupt (if in cancelable region)\n\nPrediction code:\n with CancelableGuard(): # Sets CANCELABLE=true\n predictor.predict() # Can be interrupted\n # CANCELABLE=false on exit\n```\n\n**Async Predictors:**\n```\nParent: ControlRequest::Cancel { slot }\n │\n └─▶ Worker: handler.cancel(slot)\n │\n ├─▶ Get future from slot state\n └─▶ future.cancel()\n │\n └─▶ Python raises asyncio.CancelledError\n```\n\n### Setup Log Routing\n\nDuring setup (before any prediction), logs go through the control channel:\n\n```\nworker_bridge.setup()\n │\n ├─▶ register_setup_sender(tx) # Control channel sender\n │\n ├─▶ predictor.load() + predictor.setup()\n │ │\n │ └─▶ print(\"Loading model...\")\n │ │\n │ └─▶ SlotLogWriter.write()\n │ │\n │ ├─▶ No prediction_id (not in prediction)\n │ └─▶ get_setup_sender() → ControlResponse::Log\n │\n └─▶ unregister_setup_sender()\n```\n\n### Behaviors\n\n**Worker Startup:**\n1. `set_active()` - Mark as worker subprocess\n2. `init_tracing()` - Configure logging (stderr, COG_LOG_LEVEL env)\n3. `install_slot_log_writers()` - Replace sys.stdout/stderr\n4. `install_audit_hook()` - Protect streams\n5. `install_signal_handler()` - SIGUSR1 for cancellation\n6. Read Init message from stdin\n7. Connect to slot sockets\n8. `handler.setup()` - Load and initialize predictor\n9. Send Ready message\n10. Enter event loop\n\n**Shutdown:**\n- ControlRequest::Shutdown → Send ShuttingDown, exit\n- stdin closes (parent died) → Exit immediately\n- All slots poisoned → Exit\n\n**Error Handling:**\n- SetupError::Load - Failed to import/instantiate predictor\n- SetupError::Setup - setup() raised exception\n- PredictionError - Prediction failed, slot stays healthy\n- Slot write error → Slot poisoned (no more predictions on that slot)\n" }, { "path": "crates/coglet-python/build.rs", "content": "//! Build script for coglet-python.\n//!\n//! Captures build metadata and converts semver to PEP 440 for Python compatibility.\n\nuse std::process::Command;\n\nfn main() {\n // Convert CARGO_PKG_VERSION (semver) to PEP 440\n let version = env!(\"CARGO_PKG_VERSION\");\n let pep440 = semver_to_pep440(version);\n println!(\"cargo:rustc-env=COGLET_PEP440_VERSION={pep440}\");\n\n // Git SHA (short)\n let git_sha = Command::new(\"git\")\n .args([\"rev-parse\", \"--short\", \"HEAD\"])\n .output()\n .ok()\n .filter(|o| o.status.success())\n .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())\n .unwrap_or_else(|| \"unknown\".to_string());\n println!(\"cargo:rustc-env=COGLET_GIT_SHA={git_sha}\");\n\n // Git dirty flag\n let git_dirty = Command::new(\"git\")\n .args([\"status\", \"--porcelain\"])\n .output()\n .ok()\n .filter(|o| o.status.success())\n .map(|o| {\n if String::from_utf8_lossy(&o.stdout).trim().is_empty() {\n \"false\"\n } else {\n \"true\"\n }\n })\n .unwrap_or(\"unknown\");\n println!(\"cargo:rustc-env=COGLET_GIT_DIRTY={git_dirty}\");\n\n // Build timestamp (UTC, ISO 8601)\n let build_time = Command::new(\"date\")\n .args([\"-u\", \"+%Y-%m-%dT%H:%M:%SZ\"])\n .output()\n .ok()\n .filter(|o| o.status.success())\n .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())\n .unwrap_or_else(|| \"unknown\".to_string());\n println!(\"cargo:rustc-env=COGLET_BUILD_TIME={build_time}\");\n\n // Rustc version\n let rustc_version = Command::new(\"rustc\")\n .args([\"--version\"])\n .output()\n .ok()\n .filter(|o| o.status.success())\n .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())\n .unwrap_or_else(|| \"unknown\".to_string());\n println!(\"cargo:rustc-env=COGLET_RUSTC_VERSION={rustc_version}\");\n\n // Rebuild if git HEAD changes or files are staged\n println!(\"cargo:rerun-if-changed=../../.git/HEAD\");\n println!(\"cargo:rerun-if-changed=../../.git/refs\");\n println!(\"cargo:rerun-if-changed=../../.git/index\");\n}\n\n/// Convert a semver version string to PEP 440 format.\n///\n/// Mapping:\n/// 0.17.0 → 0.17.0\n/// 0.17.0-alpha.2 → 0.17.0a2\n/// 0.17.0-beta.1 → 0.17.0b1\n/// 0.17.0-rc.3 → 0.17.0rc3\n/// 0.17.0-dev.4 → 0.17.0.dev4\nfn semver_to_pep440(version: &str) -> String {\n let Some((base, pre)) = version.split_once('-') else {\n return version.to_string();\n };\n\n if let Some(n) = pre.strip_prefix(\"alpha.\") {\n format!(\"{base}a{n}\")\n } else if let Some(n) = pre.strip_prefix(\"alpha\") {\n if n.is_empty() {\n format!(\"{base}a0\")\n } else {\n format!(\"{base}a{n}\")\n }\n } else if let Some(n) = pre.strip_prefix(\"beta.\") {\n format!(\"{base}b{n}\")\n } else if let Some(n) = pre.strip_prefix(\"beta\") {\n if n.is_empty() {\n format!(\"{base}b0\")\n } else {\n format!(\"{base}b{n}\")\n }\n } else if let Some(n) = pre.strip_prefix(\"rc.\") {\n format!(\"{base}rc{n}\")\n } else if let Some(n) = pre.strip_prefix(\"rc\") {\n if n.is_empty() {\n format!(\"{base}rc0\")\n } else {\n format!(\"{base}rc{n}\")\n }\n } else if let Some(n) = pre.strip_prefix(\"dev.\") {\n format!(\"{base}.dev{n}\")\n } else if let Some(n) = pre.strip_prefix(\"dev\") {\n if n.is_empty() {\n format!(\"{base}.dev0\")\n } else {\n format!(\"{base}.dev{n}\")\n }\n } else {\n // Unknown pre-release format, pass through\n version.to_string()\n }\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n\n #[test]\n fn test_stable_version() {\n assert_eq!(semver_to_pep440(\"0.17.0\"), \"0.17.0\");\n assert_eq!(semver_to_pep440(\"1.0.0\"), \"1.0.0\");\n }\n\n #[test]\n fn test_alpha() {\n assert_eq!(semver_to_pep440(\"0.17.0-alpha.2\"), \"0.17.0a2\");\n assert_eq!(semver_to_pep440(\"0.17.0-alpha.0\"), \"0.17.0a0\");\n assert_eq!(semver_to_pep440(\"0.17.0-alpha\"), \"0.17.0a0\");\n }\n\n #[test]\n fn test_beta() {\n assert_eq!(semver_to_pep440(\"0.17.0-beta.1\"), \"0.17.0b1\");\n assert_eq!(semver_to_pep440(\"0.17.0-beta\"), \"0.17.0b0\");\n }\n\n #[test]\n fn test_rc() {\n assert_eq!(semver_to_pep440(\"0.17.0-rc.3\"), \"0.17.0rc3\");\n assert_eq!(semver_to_pep440(\"0.17.0-rc\"), \"0.17.0rc0\");\n }\n\n #[test]\n fn test_dev() {\n assert_eq!(semver_to_pep440(\"0.17.0-dev.4\"), \"0.17.0.dev4\");\n assert_eq!(semver_to_pep440(\"0.17.0-dev\"), \"0.17.0.dev0\");\n }\n}\n" }, { "path": "crates/coglet-python/coglet/__init__.py", "content": "\"\"\"coglet — high-performance Rust prediction server for Cog ML models.\"\"\"\n\nfrom coglet._impl import CancelationException, __build__, __version__, server\nfrom coglet._impl import _sdk as _sdk\n\n__all__ = [\"__version__\", \"__build__\", \"server\", \"CancelationException\"]\n" }, { "path": "crates/coglet-python/coglet/__init__.pyi", "content": "# This file is automatically generated by stub_gen\n# ruff: noqa: E501, F401\n\nfrom coglet._impl import __build__ as __build__, __version__ as __version__, server as server, CancelationException as CancelationException\nfrom . import _sdk as _sdk\n\n__all__ = ['__build__', '__version__', 'server', 'CancelationException']\n" }, { "path": "crates/coglet-python/coglet/_impl.pyi", "content": "# This file is automatically generated by pyo3_stub_gen\n# ruff: noqa: E501, F401, F403, F405\n\nimport builtins\nimport typing\nfrom . import _sdk\n__all__ = [\n \"BuildInfo\",\n \"CancelationException\",\n \"Server\",\n \"server\",\n]\n\n__build__: BuildInfo\n__version__: builtins.str\nserver: Server\n@typing.final\nclass BuildInfo:\n r\"\"\"\n Frozen build metadata exposed as `coglet.__build__`.\n \"\"\"\n @property\n def version(self) -> builtins.str: ...\n @property\n def git_sha(self) -> builtins.str: ...\n @property\n def build_time(self) -> builtins.str: ...\n @property\n def rustc_version(self) -> builtins.str: ...\n def __repr__(self) -> builtins.str: ...\n\nclass CancelationException(builtins.BaseException):\n r\"\"\"\n Raised when a running prediction or training is cancelled.\n \n Derives from ``BaseException`` (not ``Exception``) so that bare\n ``except Exception`` blocks do not accidentally swallow cancellation.\n This matches the semantics of ``KeyboardInterrupt`` and\n ``asyncio.CancelledError``.\n \"\"\"\n ...\n\n@typing.final\nclass Server:\n r\"\"\"\n The coglet prediction server.\n \n Access via `coglet.server`. Frozen — attributes cannot be set or deleted.\n \n - `coglet.server.active` — `True` when running inside a worker subprocess\n - `coglet.server.serve(...)` — start the HTTP prediction server (blocking)\n \"\"\"\n @property\n def active(self) -> builtins.bool:\n r\"\"\"\n `True` when running inside a coglet worker subprocess.\n \"\"\"\n def serve(self, predictor_ref: typing.Optional[builtins.str] = None, host: builtins.str = '0.0.0.0', port: builtins.int = 5000, await_explicit_shutdown: builtins.bool = False, is_train: builtins.bool = False, output_temp_dir_base: builtins.str = '/tmp/coglet/output', upload_url: typing.Optional[builtins.str] = None) -> None:\n r\"\"\"\n Start the HTTP prediction server. Blocks until shutdown.\n \"\"\"\n def _run_worker(self) -> None:\n r\"\"\"\n Worker subprocess entry point. Called by the orchestrator.\n \n Sets the active flag, installs log writers and audit hooks,\n then enters the worker event loop.\n \"\"\"\n def __repr__(self) -> builtins.str: ...\n\n" }, { "path": "crates/coglet-python/coglet/_sdk/__init__.pyi", "content": "# This file is automatically generated by pyo3_stub_gen\n# ruff: noqa: E501, F401, F403, F405\n\nimport builtins\nimport typing\n\n__all__ = [\n \"MetricRecorder\",\n \"Scope\",\n \"current_scope\",\n]\n\n@typing.final\nclass MetricRecorder:\n r\"\"\"\n Metric recorder with type invariant enforcement.\n\n Accessed via `scope.metrics`. Supports:\n - `scope.metrics.record(key, value, mode=\"replace\")` — full API\n - `scope.metrics.delete(key)` — delete (required before type change)\n - `scope.metrics[key] = value` — dict-style set (replace mode)\n - `del scope.metrics[key]` — dict-style delete\n \"\"\"\n def record(\n self,\n key: builtins.str,\n value: typing.Any,\n mode: typing.Optional[builtins.str] = None,\n ) -> None:\n r\"\"\"\n Record a metric value.\n\n Args:\n key: Metric name. Dot-separated keys (e.g. \"timing.preprocess\") create\n nested objects in the response.\n value: Must be bool, int, float, str, list, or dict. Once a key is set\n with a type, it cannot be changed without calling delete() first.\n mode: Accumulation mode — \"replace\" (default), \"incr\" (increment numeric),\n or \"append\" (push to array).\n \"\"\"\n def delete(self, key: builtins.str) -> None:\n r\"\"\"\n Delete a metric key. Required before changing a metric's type.\n \"\"\"\n def __setitem__(self, key: builtins.str, value: typing.Any) -> None:\n r\"\"\"\n Dict-style set: `scope.metrics[\"key\"] = value`\n \"\"\"\n def __delitem__(self, key: builtins.str) -> None:\n r\"\"\"\n Dict-style delete: `del scope.metrics[\"key\"]`\n \"\"\"\n def __repr__(self) -> builtins.str: ...\n\n@typing.final\nclass Scope:\n r\"\"\"\n Prediction scope, obtained via `current_scope()`.\n\n Provides access to `scope.metrics` for recording metrics, and\n `scope.record_metric()` as a convenience shorthand.\n \"\"\"\n @property\n def metrics(self) -> MetricRecorder:\n r\"\"\"\n The metric recorder for this prediction.\n \"\"\"\n @property\n def context(self) -> dict[builtins.str, builtins.str]:\n r\"\"\"\n Per-prediction context passed in the request body.\n\n Returns a `dict[str, str]` (empty dict if no context was provided).\n \"\"\"\n def record_metric(\n self,\n key: builtins.str,\n value: typing.Any,\n mode: typing.Optional[builtins.str] = None,\n ) -> None:\n r\"\"\"\n Convenience: record a metric value.\n\n Equivalent to `scope.metrics.record(key, value, mode)`.\n \"\"\"\n def __repr__(self) -> builtins.str: ...\n\n@typing.final\nclass _SlotLogWriter:\n r\"\"\"\n A Python file-like object that routes writes via the prediction_id ContextVar.\n\n This is installed as sys.stdout/stderr once at worker startup.\n Each write looks up the current prediction_id from the ContextVar and routes\n to the appropriate SlotSender.\n\n If no prediction_id is set, or the prediction has completed (orphan task),\n writes go to tracing (logged as orphan).\n\n Uses line buffering: accumulates writes until a newline is received, then\n emits complete lines. This coalesces Python's print() which does separate\n writes for content and newline.\n \"\"\"\n @property\n def closed(self) -> builtins.bool:\n r\"\"\"\n Whether writes should be ignored (used after errors).\n \"\"\"\n @property\n def encoding(self) -> typing.Optional[builtins.str]:\n r\"\"\"\n Encoding property - needed for compatibility.\n \"\"\"\n @property\n def newlines(self) -> typing.Optional[builtins.str]:\n r\"\"\"\n Newlines property - needed for compatibility.\n \"\"\"\n @property\n def buffer(self) -> typing.Any:\n r\"\"\"\n Buffer property - some code checks for this.\n \"\"\"\n def write(self, data: builtins.str) -> builtins.int:\n r\"\"\"\n Write data, routing to the appropriate destination.\n\n Uses line buffering: accumulates data until a newline is received, then\n emits complete lines. This coalesces Python's print() which does separate\n writes for content and the trailing newline.\n\n Priority for routing:\n 1. If inside a prediction (ContextVar set), route to slot sender\n 2. If setup sender registered, route to control channel\n 3. Fall back to stderr (for orphan tasks or unexpected cases)\n \"\"\"\n def emit_data(self, data: builtins.str) -> None:\n r\"\"\"\n Emit data to the appropriate destination.\n \"\"\"\n def flush(self) -> None:\n r\"\"\"\n Flush the stream.\n\n Emits any buffered content that hasn't been terminated with a newline.\n \"\"\"\n def readable(self) -> builtins.bool:\n r\"\"\"\n Return whether the stream is readable.\n \"\"\"\n def writable(self) -> builtins.bool:\n r\"\"\"\n Return whether the stream is writable.\n \"\"\"\n def seekable(self) -> builtins.bool:\n r\"\"\"\n Return whether the stream is seekable.\n \"\"\"\n def isatty(self) -> builtins.bool:\n r\"\"\"\n Return whether the stream is a TTY.\n \"\"\"\n def fileno(self) -> builtins.int:\n r\"\"\"\n Return the file number.\n \"\"\"\n def close(self) -> None:\n r\"\"\"\n Close the stream.\n \"\"\"\n def __enter__(self) -> _SlotLogWriter:\n r\"\"\"\n Context manager enter.\n \"\"\"\n def __exit__(\n self,\n _exc_type: typing.Optional[typing.Any],\n _exc_val: typing.Optional[typing.Any],\n _exc_tb: typing.Optional[typing.Any],\n ) -> builtins.bool:\n r\"\"\"\n Context manager exit.\n \"\"\"\n\n@typing.final\nclass _TeeWriter:\n r\"\"\"\n Tee writer that sends writes to both our slot routing and user's stream.\n\n - inner: Our _SlotLogWriter for slot-based log routing\n - user_stream: The stream user code tried to install\n \"\"\"\n @property\n def inner(self) -> typing.Any:\n r\"\"\"\n Our _SlotLogWriter (does ContextVar-based routing)\n \"\"\"\n @property\n def user_stream(self) -> typing.Any:\n r\"\"\"\n User's replacement stream\n \"\"\"\n @property\n def name(self) -> builtins.str:\n r\"\"\"\n Stream name (stdout or stderr)\n \"\"\"\n @property\n def closed(self) -> builtins.bool:\n r\"\"\"\n Closed flag\n \"\"\"\n @property\n def encoding(self) -> typing.Optional[builtins.str]: ...\n @property\n def newlines(self) -> typing.Optional[builtins.str]: ...\n def __new__(\n cls, inner: typing.Any, user_stream: typing.Any, name: builtins.str\n ) -> _TeeWriter: ...\n def write(self, data: builtins.str) -> builtins.int:\n r\"\"\"\n Write to both streams.\n \"\"\"\n def flush(self) -> None:\n r\"\"\"\n Flush both streams.\n \"\"\"\n def readable(self) -> builtins.bool: ...\n def writable(self) -> builtins.bool: ...\n def seekable(self) -> builtins.bool: ...\n def isatty(self) -> builtins.bool: ...\n def fileno(self) -> builtins.int: ...\n def close(self) -> None: ...\n def __enter__(self) -> _TeeWriter: ...\n def __exit__(\n self,\n _exc_type: typing.Optional[typing.Any],\n _exc_val: typing.Optional[typing.Any],\n _exc_tb: typing.Optional[typing.Any],\n ) -> builtins.bool: ...\n\ndef current_scope() -> Scope:\n r\"\"\"\n Python-callable: get the current Scope.\n\n Returns the active scope if inside a prediction, or a no-op scope otherwise.\n \"\"\"\n" }, { "path": "crates/coglet-python/coglet/py.typed", "content": "" }, { "path": "crates/coglet-python/pyproject.toml", "content": "[build-system]\nrequires = [\"maturin>=1.0,<2.0\"]\nbuild-backend = \"maturin\"\n\n[project]\nname = \"coglet\"\ndescription = \"High-performance Rust prediction server for Cog ML models\"\nreadme = \"README.md\"\nlicense = {text = \"Apache-2.0\"}\nrequires-python = \">=3.10\"\nauthors = [\n {name = \"Replicate\", email = \"team@replicate.com\"},\n]\nkeywords = [\"machine-learning\", \"inference\", \"cog\", \"prediction\"]\nclassifiers = [\n \"Development Status :: 4 - Beta\",\n \"Intended Audience :: Developers\",\n \"License :: OSI Approved :: Apache Software License\",\n \"Operating System :: MacOS\",\n \"Operating System :: POSIX :: Linux\",\n \"Programming Language :: Rust\",\n \"Programming Language :: Python :: Implementation :: CPython\",\n \"Programming Language :: Python :: 3.10\",\n \"Programming Language :: Python :: 3.11\",\n \"Programming Language :: Python :: 3.12\",\n \"Programming Language :: Python :: 3.13\",\n \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n]\ndynamic = [\"version\"]\ndependencies = []\n\n[project.urls]\nHomepage = \"https://cog.run\"\nDocumentation = \"https://cog.run/docs\"\nRepository = \"https://github.com/replicate/cog\"\nIssues = \"https://github.com/replicate/cog/issues\"\n\n[project.optional-dependencies]\ntest = [\n \"pytest>=8.0\",\n \"requests>=2.31\",\n]\n\n[tool.maturin]\nfeatures = [\"pyo3/extension-module\"]\n# Mixed layout: coglet/__init__.py is hand-managed, .so is named _impl.\nmodule-name = \"coglet._impl\"\n# Tell pyo3-stub-gen where the Python source root is (also used by maturin)\npython-source = \".\"\n# Use manylinux2014 (glibc 2.17) for compatibility with Python 3.10+ base images\ncompatibility = \"manylinux2014\"\n\n[tool.pytest.ini_options]\ntestpaths = [\"tests\"]\n" }, { "path": "crates/coglet-python/src/audit.rs", "content": "//! Audit hooks to protect Rust-injected runtime objects.\n//!\n//! Uses sys.addaudithook to intercept operations that could interfere with\n//! our runtime machinery. The hook cannot be removed once added.\n//!\n//! ## Protection: sys.stdout/stderr (Tee pattern)\n//!\n//! If user code replaces stdout/stderr, we wrap their replacement in a _TeeWriter\n//! that sends data to BOTH our slot routing AND their stream. User's code works\n//! as they expect, but we still get our logs.\n//!\n//! If they replace again, we unwrap the inner _SlotLogWriter from the current\n//! _TeeWriter and re-tee with the new stream. No nested _TeeWriters.\n\nuse std::sync::atomic::{AtomicBool, Ordering};\nuse std::sync::{Mutex, OnceLock};\n\nuse pyo3::prelude::*;\nuse pyo3_stub_gen::derive::*;\n\n/// Whether the audit hook has been installed.\nstatic HOOK_INSTALLED: AtomicBool = AtomicBool::new(false);\n\n/// Re-entrancy guard for the audit hook.\n/// Prevents infinite recursion when the hook itself sets sys.stdout/stderr.\nstatic IN_HOOK: AtomicBool = AtomicBool::new(false);\n\n/// Serializes stream replacement so concurrent threads don't race\n/// on the read-current → create-tee → set-new sequence.\nstatic STREAM_LOCK: Mutex<()> = Mutex::new(());\n\n/// Reference to sys module for identity comparison in hook.\nstatic SYS_MODULE: OnceLock> = OnceLock::new();\n\n/// Reference to our _SlotLogWriter class for isinstance checks.\nstatic SLOT_LOG_WRITER_TYPE: OnceLock> = OnceLock::new();\n\n/// Install the audit hook. Called once at worker startup.\n///\n/// The hook intercepts object.__setattr__ on sys for stdout/stderr.\npub fn install_audit_hook(py: Python<'_>) -> PyResult<()> {\n if HOOK_INSTALLED.swap(true, Ordering::SeqCst) {\n return Ok(());\n }\n\n // Store sys module reference for identity comparison\n let sys = py.import(\"sys\")?;\n let _ = SYS_MODULE.set(sys.as_any().clone().unbind());\n\n // Store our _SlotLogWriter type for isinstance checks\n if let Ok(coglet) = py.import(\"coglet\")\n && let Ok(writer_type) = coglet.getattr(\"_SlotLogWriter\")\n {\n let _ = SLOT_LOG_WRITER_TYPE.set(writer_type.unbind());\n }\n\n // Register the Rust audit hook callable\n let hook = wrap_pyfunction!(_coglet_audit_hook, py)?;\n sys.call_method1(\"addaudithook\", (hook,))?;\n\n tracing::debug!(\"Installed audit hook for runtime protection\");\n Ok(())\n}\n\n/// Audit hook implemented in Rust.\n///\n/// Intercepts `object.__setattr__` events on `sys` for stdout/stderr.\n/// Uses an AtomicBool re-entrancy guard instead of deferred threading.Timer.\n#[pyfunction]\nfn _coglet_audit_hook(py: Python<'_>, event: &str, args: &Bound<'_, PyAny>) -> PyResult<()> {\n if event != \"object.__setattr__\" {\n return Ok(());\n }\n\n // Re-entrancy guard: skip if we're already inside the hook\n // (because we're setting sys.stdout/stderr ourselves).\n if IN_HOOK.load(Ordering::SeqCst) {\n return Ok(());\n }\n\n // args is (obj, name, value)\n let obj = args.get_item(0)?;\n let name: String = args.get_item(1)?.extract()?;\n\n if name != \"stdout\" && name != \"stderr\" {\n return Ok(());\n }\n\n // Check if obj is the sys module (identity comparison)\n let Some(sys_ref) = SYS_MODULE.get() else {\n return Ok(());\n };\n if !obj.is(sys_ref.bind(py)) {\n return Ok(());\n }\n\n let value = args.get_item(2)?;\n handle_stream_replacement(py, &name, &value)?;\n\n Ok(())\n}\n\n/// Handle user code replacing sys.stdout or sys.stderr.\n///\n/// If the new value is already our _SlotLogWriter, this is our own setup — skip.\n/// Otherwise, find our _SlotLogWriter from the current stream (direct or inside\n/// a _TeeWriter), and wrap the user's new stream in a fresh _TeeWriter.\nfn handle_stream_replacement(py: Python<'_>, name: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {\n // If value is our _SlotLogWriter, this is us installing — skip\n if is_slot_log_writer(py, value) {\n return Ok(());\n }\n\n // Serialize the read-current → create-tee → set-new sequence.\n // Without this, two threads replacing stdout simultaneously could race\n // and one tee gets silently dropped.\n // The lock protects no data (just `()`), so poisoned is safe to recover.\n let _lock = STREAM_LOCK.lock().unwrap_or_else(|poisoned| {\n tracing::warn!(\n target: \"coglet::worker_local\",\n \"stream lock was poisoned (a thread panicked during stream replacement) — \\\n recovering, but log routing may be inconsistent\"\n );\n poisoned.into_inner()\n });\n\n // Get current writer from sys\n let sys = py.import(\"sys\")?;\n let current = sys.getattr(name)?;\n\n // Find our _SlotLogWriter — either it IS current, or it's inside a _TeeWriter\n let slot_writer = if is_slot_log_writer(py, ¤t) {\n Some(current.clone().unbind())\n } else if is_tee_writer(¤t) {\n get_inner_writer(py, ¤t).ok()\n } else {\n None\n };\n\n let Some(slot_writer) = slot_writer else {\n // No _SlotLogWriter installed — nothing to protect\n return Ok(());\n };\n\n // Create new _TeeWriter wrapping our _SlotLogWriter and user's stream\n let tee = _TeeWriter::new(slot_writer, value.clone().unbind(), name.to_string());\n let tee_obj = tee.into_pyobject(py)?;\n\n // Set under re-entrancy guard to prevent hook from re-triggering\n IN_HOOK.store(true, Ordering::SeqCst);\n let result = sys.setattr(name, tee_obj);\n IN_HOOK.store(false, Ordering::SeqCst);\n\n result\n}\n\n// ============================================================================\n// Type checks — pub(crate) only, not exported to Python\n// ============================================================================\n\n/// Check if a value is a _SlotLogWriter.\npub(crate) fn is_slot_log_writer(py: Python<'_>, value: &Bound<'_, PyAny>) -> bool {\n if let Some(writer_type) = SLOT_LOG_WRITER_TYPE.get()\n && let Ok(true) = value.is_instance(writer_type.bind(py))\n {\n return true;\n }\n\n // Fallback: check by class name (handles cross-module edge cases)\n if let Ok(type_name) = value.get_type().name() {\n return type_name == \"_SlotLogWriter\";\n }\n\n false\n}\n\n/// Check if a value is a _TeeWriter.\npub(crate) fn is_tee_writer(value: &Bound<'_, PyAny>) -> bool {\n if value.is_instance_of::<_TeeWriter>() {\n return true;\n }\n\n if let Ok(type_name) = value.get_type().name() {\n return type_name == \"_TeeWriter\";\n }\n\n false\n}\n\n/// Get the inner _SlotLogWriter from a _TeeWriter.\npub(crate) fn get_inner_writer(py: Python<'_>, tee: &Bound<'_, PyAny>) -> PyResult> {\n if let Ok(tee_writer) = tee.extract::>() {\n return Ok(tee_writer.inner.clone_ref(py));\n }\n\n if let Ok(inner) = tee.getattr(\"inner\") {\n return Ok(inner.unbind());\n }\n\n Err(pyo3::exceptions::PyTypeError::new_err(\n \"Expected _TeeWriter with inner attribute\",\n ))\n}\n\n// ============================================================================\n// _TeeWriter — private pyclass\n// ============================================================================\n\n/// Tee writer that sends writes to both our slot routing and user's stream.\n///\n/// - inner: Our _SlotLogWriter for slot-based log routing\n/// - user_stream: The stream user code tried to install\n#[gen_stub_pyclass]\n#[pyclass(name = \"_TeeWriter\", module = \"coglet._sdk\")]\npub struct _TeeWriter {\n /// Our _SlotLogWriter (does ContextVar-based routing)\n #[pyo3(get)]\n inner: Py,\n /// User's replacement stream\n #[pyo3(get)]\n user_stream: Py,\n /// Stream name (stdout or stderr)\n #[pyo3(get)]\n name: String,\n /// Closed flag\n #[pyo3(get)]\n closed: bool,\n}\n\n#[gen_stub_pymethods]\n#[pymethods]\nimpl _TeeWriter {\n #[new]\n fn new(inner: Py, user_stream: Py, name: String) -> Self {\n Self {\n inner,\n user_stream,\n name,\n closed: false,\n }\n }\n\n /// Write to both streams.\n fn write(&self, py: Python<'_>, data: &str) -> PyResult {\n if self.closed || data.is_empty() {\n return Ok(data.len());\n }\n\n if let Err(e) = self.inner.call_method1(py, \"write\", (data,)) {\n tracing::warn!(error = %e, \"_TeeWriter: failed to write to inner\");\n }\n\n if let Err(e) = self.user_stream.call_method1(py, \"write\", (data,)) {\n tracing::warn!(error = %e, \"_TeeWriter: failed to write to user stream\");\n }\n\n Ok(data.len())\n }\n\n /// Flush both streams.\n fn flush(&self, py: Python<'_>) -> PyResult<()> {\n let _ = self.inner.call_method0(py, \"flush\");\n let _ = self.user_stream.call_method0(py, \"flush\");\n Ok(())\n }\n\n fn readable(&self) -> bool {\n false\n }\n\n fn writable(&self) -> bool {\n !self.closed\n }\n\n fn seekable(&self) -> bool {\n false\n }\n\n fn isatty(&self, py: Python<'_>) -> PyResult {\n let result = self.user_stream.call_method0(py, \"isatty\")?;\n result.extract(py)\n }\n\n fn fileno(&self, py: Python<'_>) -> PyResult {\n let result = self.user_stream.call_method0(py, \"fileno\")?;\n result.extract(py)\n }\n\n fn close(&mut self) {\n self.closed = true;\n }\n\n fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {\n slf\n }\n\n fn __exit__(\n &mut self,\n _exc_type: Option<&Bound<'_, PyAny>>,\n _exc_val: Option<&Bound<'_, PyAny>>,\n _exc_tb: Option<&Bound<'_, PyAny>>,\n ) -> bool {\n false\n }\n\n #[getter]\n fn encoding(&self, py: Python<'_>) -> PyResult> {\n match self.user_stream.getattr(py, \"encoding\") {\n Ok(enc) => enc.extract(py),\n Err(_) => Ok(Some(\"utf-8\".to_string())),\n }\n }\n\n #[getter]\n fn newlines(&self) -> Option {\n None\n }\n}\n" }, { "path": "crates/coglet-python/src/bin/stub_gen.rs", "content": "//! Generate Python stub files for coglet.\n//!\n//! Run with: cargo run --bin stub_gen\n//!\n//! Custom generate logic: pyo3-stub-gen places classes from the native\n//! `coglet._impl` module into the `coglet` parent package stub. We redirect\n//! that output to `coglet/_impl.pyi` so native module types are preserved in\n//! the right place, then generate `coglet/__init__.pyi` ourselves to\n//! re-export the public API — matching the hand-maintained `__init__.py`.\n\nuse pyo3_stub_gen::Result;\nuse std::fs;\nuse std::io::Write;\n\n/// Public items re-exported from `coglet._impl` in `coglet/__init__.pyi`.\n/// Uses the `X as X` pattern to mark explicit re-exports (PEP 484).\nconst PUBLIC_REEXPORTS: &[&str] = &[\"__build__\", \"__version__\", \"server\", \"CancelationException\"];\n\n/// Private submodules re-exported with `from . import X as X`.\n///\n/// These use a relative import (not `from coglet._impl`) because `_sdk` is a\n/// subpackage that type checkers resolve via the filesystem, not an attribute\n/// of the native extension module. Not included in `__all__`.\nconst PRIVATE_REEXPORTS: &[&str] = &[\"_sdk\"];\n\nfn main() -> Result<()> {\n let stub = coglet::stub_info()?;\n\n for (name, module) in &stub.modules {\n let normalized = name.replace('-', \"_\");\n\n let dest = if normalized == \"coglet\" {\n // Native module classes land here — redirect to _impl.pyi\n stub.python_root.join(\"coglet\").join(\"_impl.pyi\")\n } else {\n // Submodules like \"coglet._sdk\" → coglet/_sdk/__init__.pyi\n let path = normalized.replace('.', \"/\");\n stub.python_root.join(&path).join(\"__init__.pyi\")\n };\n\n let dir = dest.parent().expect(\"cannot get parent directory\");\n if !dir.exists() {\n fs::create_dir_all(dir)?;\n }\n\n let mut f = fs::File::create(&dest)?;\n write!(f, \"{module}\")?;\n eprintln!(\"Generated stub: {}\", dest.display());\n }\n\n // Generate coglet/__init__.pyi — re-exports from _impl\n let init_pyi = stub.python_root.join(\"coglet\").join(\"__init__.pyi\");\n let mut f = fs::File::create(&init_pyi)?;\n\n writeln!(f, \"# This file is automatically generated by stub_gen\")?;\n writeln!(f, \"# ruff: noqa: E501, F401\")?;\n writeln!(f)?;\n\n // `from coglet._impl import X as X, Y as Y, ...`\n let reexports: Vec = PUBLIC_REEXPORTS\n .iter()\n .map(|name| format!(\"{name} as {name}\"))\n .collect();\n writeln!(f, \"from coglet._impl import {}\", reexports.join(\", \"))?;\n\n // `from . import _sdk as _sdk` — relative import so ty resolves the\n // subpackage via coglet/_sdk/__init__.pyi, not through _impl.\n let private: Vec = PRIVATE_REEXPORTS\n .iter()\n .map(|name| format!(\"{name} as {name}\"))\n .collect();\n writeln!(f, \"from . import {}\", private.join(\", \"))?;\n writeln!(f)?;\n\n // __all__ only includes public items (no underscore-prefixed names)\n let all_items: Vec = PUBLIC_REEXPORTS\n .iter()\n .map(|name| format!(\"'{name}'\"))\n .collect();\n writeln!(f, \"__all__ = [{}]\", all_items.join(\", \"))?;\n\n eprintln!(\"Generated stub: {}\", init_pyi.display());\n\n Ok(())\n}\n" }, { "path": "crates/coglet-python/src/cancel.rs", "content": "//! Cancellation support for predictions.\n//!\n//! Sync predictors use `PyThreadState_SetAsyncExc` to inject a\n//! `CancelationException` (a `BaseException` subclass) into the Python\n//! thread running `predict()`.\n//!\n//! Async predictors use asyncio task cancellation:\n//! - Store task reference when prediction starts\n//! - Call task.cancel() when cancel requested\n//! - Python raises asyncio.CancelledError\n//!\n//! `CancelationException` deliberately derives from `BaseException` (not\n//! `Exception`) so that bare `except Exception` blocks in user code cannot\n//! swallow it — matching the semantics of `KeyboardInterrupt` and\n//! `asyncio.CancelledError`.\n\nuse pyo3::prelude::*;\n\n// Static exception type with automatic stub generation.\n// Derives from BaseException so `except Exception` does not catch it.\npyo3_stub_gen::create_exception!(\n coglet,\n CancelationException,\n pyo3::exceptions::PyBaseException,\n \"Raised when a running prediction or training is cancelled.\\n\\\n \\n\\\n Derives from ``BaseException`` (not ``Exception``) so that bare\\n\\\n ``except Exception`` blocks do not accidentally swallow cancellation.\\n\\\n This matches the semantics of ``KeyboardInterrupt`` and\\n\\\n ``asyncio.CancelledError``.\"\n);\n\n/// Inject CancelationException into a specific Python thread.\n///\n/// Uses CPython's `PyThreadState_SetAsyncExc` to raise the exception at the\n/// next bytecode boundary. This works on any thread (not just the main thread),\n/// unlike SIGUSR1-based cancellation.\n///\n/// Requires the GIL — `Python::attach` acquires it, blocking briefly if the\n/// prediction thread currently holds it (CPython releases it every ~5ms).\npub fn cancel_sync_thread(py_thread_id: std::ffi::c_long) {\n Python::attach(|py| {\n let exc = py.get_type::().as_ptr();\n\n // SAFETY: We hold the GIL. exc is a valid Python type pointer\n // obtained from the interpreter's type registry.\n let result = unsafe { pyo3::ffi::PyThreadState_SetAsyncExc(py_thread_id, exc) };\n\n match result {\n 0 => {\n tracing::warn!(\n py_thread_id,\n \"PyThreadState_SetAsyncExc: thread not found (prediction may have completed)\"\n );\n }\n 1 => {\n tracing::debug!(\n py_thread_id,\n \"Injected CancelationException into Python thread\"\n );\n }\n _ => {\n // CPython docs: if > 1, call again with NULL to reset\n tracing::error!(\n py_thread_id,\n count = result,\n \"PyThreadState_SetAsyncExc modified multiple thread states, resetting\"\n );\n unsafe {\n pyo3::ffi::PyThreadState_SetAsyncExc(py_thread_id, std::ptr::null_mut());\n }\n }\n }\n });\n}\n\n/// Get the current Python thread identifier (for later use with `cancel_sync_thread`).\n///\n/// Uses `threading.get_ident()` which returns the same value as\n/// `PyThreadState_SetAsyncExc` expects for the thread id argument.\n/// Can be called from any thread (acquires the GIL briefly).\npub fn current_py_thread_id() -> std::ffi::c_long {\n Python::attach(|py| {\n let threading = py.import(\"threading\").expect(\"failed to import threading\");\n threading\n .call_method0(\"get_ident\")\n .expect(\"threading.get_ident() failed\")\n .extract::()\n .expect(\"thread ident is not an integer\")\n })\n}\n" }, { "path": "crates/coglet-python/src/input.rs", "content": "//! Input processing for cog predictors.\n//!\n//! This module handles file downloads for cog predictor inputs.\n//! Input validation is performed at the HTTP edge using the OpenAPI schema;\n//! the worker only needs to download URLPath inputs and pass them through.\n\nuse std::collections::HashSet;\n\nuse pyo3::prelude::*;\nuse pyo3::types::PyDict;\n\n/// Type alias for Python object.\ntype PyObject = Py;\n\n/// RAII wrapper for prepared input that cleans up temp files on drop.\n///\n/// When URLPath inputs are downloaded, they create temp files. This struct\n/// ensures those files are cleaned up when the prediction completes (success,\n/// failure, or cancellation).\npub struct PreparedInput {\n /// The prepared input dict (ready for predict(**kwargs))\n dict: Py,\n /// Paths to cleanup on drop (downloaded temp files)\n cleanup_paths: Vec,\n}\n\nimpl PreparedInput {\n /// Create a new PreparedInput with the given dict and paths to cleanup.\n pub fn new(dict: Py, cleanup_paths: Vec) -> Self {\n Self {\n dict,\n cleanup_paths,\n }\n }\n\n /// Get the input dict bound to the given Python context.\n pub fn dict<'py>(&self, py: Python<'py>) -> Bound<'py, PyDict> {\n self.dict.bind(py).clone()\n }\n}\n\nimpl Drop for PreparedInput {\n fn drop(&mut self) {\n if self.cleanup_paths.is_empty() {\n return;\n }\n\n Python::attach(|py| {\n for path in &self.cleanup_paths {\n let path_bound = path.bind(py);\n let kwargs = PyDict::new(py);\n if kwargs.set_item(\"missing_ok\", true).is_ok()\n && let Err(e) = path_bound.call_method(\"unlink\", (), Some(&kwargs))\n {\n tracing::warn!(error = %e, \"Failed to cleanup temp file\");\n }\n }\n });\n }\n}\n\n// Safety: PyObject is Send in PyO3 0.23+, we only access through Python::attach\nunsafe impl Send for PreparedInput {}\n\n/// Prepare input for prediction.\n///\n/// Coerces URL strings to the appropriate cog types based on the function's\n/// type annotations: `File`-annotated params get `File.validate()` (IO-like),\n/// `Path`-annotated params get `Path.validate()` (filesystem path + download).\n/// Returns a PreparedInput that cleans up temp files on drop.\n///\n/// Input validation is handled at the HTTP edge via the OpenAPI schema —\n/// this function only handles URL->Path/File coercion and file downloads.\n///\n/// `func` is the Python predict/train callable used to introspect type annotations.\npub fn prepare_input(\n py: Python<'_>,\n input: &Bound<'_, PyDict>,\n func: &Bound<'_, PyAny>,\n) -> PyResult {\n let file_fields = detect_file_fields(py, func)?;\n coerce_url_strings(py, input, &file_fields)?;\n let cleanup_paths = download_url_paths_into_dict(py, input)?;\n Ok(PreparedInput::new(input.clone().unbind(), cleanup_paths))\n}\n\n/// Inspect a Python function's type annotations to find parameters typed as\n/// `cog.File` (or `list[File]`, `Optional[File]`, `File | None`,\n/// `Optional[list[File]]`, etc.). Returns a set of field names that should use\n/// `File.validate()` instead of `Path.validate()`.\nfn detect_file_fields(py: Python<'_>, func: &Bound<'_, PyAny>) -> PyResult> {\n let mut file_fields = HashSet::new();\n\n let cog_file_class = py.import(\"cog.types\")?.getattr(\"File\")?;\n\n // typing.get_type_hints resolves string annotations and handles forward refs\n let typing = py.import(\"typing\")?;\n let get_type_hints = typing.getattr(\"get_type_hints\")?;\n let get_origin = typing.getattr(\"get_origin\")?;\n let get_args = typing.getattr(\"get_args\")?;\n let builtins_list = py.eval(c\"list\", None, None)?;\n let union_type = typing.getattr(\"Union\")?;\n\n let hints = match get_type_hints.call1((func,)) {\n Ok(h) => h,\n Err(_) => return Ok(file_fields), // If we can't get hints, don't coerce as File\n };\n\n // Helper closure: returns true if `ty` is `File` or `list[File]`.\n let is_file_like = |ty: &Bound<'_, PyAny>| -> PyResult {\n if ty.is(&cog_file_class) {\n return Ok(true);\n }\n let inner_origin = get_origin.call1((ty,))?;\n if !inner_origin.is_none() && inner_origin.is(&builtins_list) {\n let inner_args = get_args.call1((ty,))?;\n if let Ok(t) = inner_args.cast::()\n && !t.is_empty()\n && t.get_item(0)?.is(&cog_file_class)\n {\n return Ok(true);\n }\n }\n Ok(false)\n };\n\n let hints_dict = hints.cast::()?;\n for (name, annotation) in hints_dict.iter() {\n let name_str: String = match name.extract() {\n Ok(s) => s,\n Err(_) => continue,\n };\n if name_str == \"return\" {\n continue;\n }\n\n // Direct File annotation: `param: File`\n // Also covers `list[File]` via is_file_like.\n if is_file_like(&annotation)? {\n file_fields.insert(name_str);\n continue;\n }\n\n // Union annotation: Optional[File], File | None, Optional[list[File]], etc.\n // typing.get_origin(Optional[X]) -> typing.Union\n // typing.get_args(Optional[X]) -> (X, NoneType)\n let origin = get_origin.call1((&annotation,))?;\n if !origin.is_none() && origin.is(&union_type) {\n let args = get_args.call1((&annotation,))?;\n if let Ok(args_tuple) = args.cast::() {\n for arg in args_tuple.iter() {\n // Skip NoneType\n if arg.is_none() || arg.is(py.None().into_bound(py)) {\n continue;\n }\n // Check if this variant is NoneType by comparing to type(None)\n let nonetype = py.eval(c\"type(None)\", None, None)?;\n if arg.is(&nonetype) {\n continue;\n }\n if is_file_like(&arg)? {\n file_fields.insert(name_str.clone());\n break;\n }\n }\n }\n }\n }\n\n if !file_fields.is_empty() {\n tracing::debug!(\"Detected File-typed fields: {:?}\", file_fields);\n }\n\n Ok(file_fields)\n}\n\n/// Coerce URL string values in the input dict to the appropriate cog types.\n///\n/// After `json.loads()`, all values are plain Python types. URL strings\n/// (http://, https://, data:) that represent file inputs need to be converted:\n/// - `File`-typed fields -> `File.validate()` -> returns IO-like `URLFile`\n/// - `Path`-typed fields -> `Path.validate()` -> returns `URLPath` (downloaded later)\n///\n/// This replaces the type coercion that `_adt.py`'s `PrimitiveType.normalize()`\n/// previously performed.\nfn coerce_url_strings(\n py: Python<'_>,\n payload: &Bound<'_, PyDict>,\n file_fields: &HashSet,\n) -> PyResult<()> {\n let cog_types = py.import(\"cog.types\")?;\n let path_validate = cog_types.getattr(\"Path\")?.getattr(\"validate\")?;\n let file_validate = cog_types.getattr(\"File\")?.getattr(\"validate\")?;\n\n for (key, value) in payload.iter() {\n let key_str: String = key.extract().unwrap_or_default();\n let use_file = file_fields.contains(&key_str);\n let validate = if use_file {\n &file_validate\n } else {\n &path_validate\n };\n\n // Single string value -- check if it's a URL\n if let Ok(s) = value.extract::() {\n if s.starts_with(\"http://\") || s.starts_with(\"https://\") || s.starts_with(\"data:\") {\n let coerced = validate.call1((&value,))?;\n payload.set_item(&key, coerced)?;\n }\n }\n // List of strings -- check if any are URLs\n else if let Ok(list) = value.extract::>() {\n let mut any_coerced = false;\n let new_items = pyo3::types::PyList::empty(py);\n for item in list.iter() {\n if let Ok(s) = item.extract::()\n && (s.starts_with(\"http://\")\n || s.starts_with(\"https://\")\n || s.starts_with(\"data:\"))\n {\n let coerced = validate.call1((&item,))?;\n new_items.append(coerced)?;\n any_coerced = true;\n continue;\n }\n new_items.append(item)?;\n }\n if any_coerced {\n payload.set_item(&key, new_items)?;\n }\n }\n }\n Ok(())\n}\n\n/// Download URLPath inputs in parallel and replace them in the payload dict.\n///\n/// This replicates the behavior from cog's worker.py:\n/// - Find all URLPath instances in the payload dict\n/// - Download them in parallel using ThreadPoolExecutor\n/// - Replace URLPath values with local Path in the dict\n///\n/// Returns the downloaded Path objects for cleanup on drop.\nfn download_url_paths_into_dict(\n py: Python<'_>,\n payload: &Bound<'_, PyDict>,\n) -> PyResult> {\n let cog_types = py.import(\"cog.types\")?;\n let url_path_class = cog_types.getattr(\"URLPath\")?;\n\n // Collect URLPath fields that need downloading\n // Structure: (key, value, is_list)\n let mut url_path_keys: Vec<(String, bool)> = Vec::new();\n\n for (key, value) in payload.iter() {\n let key_str: String = key.extract()?;\n\n if value.is_instance(&url_path_class)? {\n url_path_keys.push((key_str, false));\n }\n // Check for lists of URLPath\n else if let Ok(list) = value.extract::>()\n && !list.is_empty()\n {\n let all_url_paths = list\n .iter()\n .all(|item| item.is_instance(&url_path_class).unwrap_or(false));\n if all_url_paths {\n url_path_keys.push((key_str, true));\n }\n }\n }\n\n if url_path_keys.is_empty() {\n return Ok(Vec::new());\n }\n\n tracing::debug!(\"Downloading {} URLPath input(s)\", url_path_keys.len());\n\n // Use ThreadPoolExecutor to download in parallel (like worker.py)\n let concurrent_futures = py.import(\"concurrent.futures\")?;\n let executor_class = concurrent_futures.getattr(\"ThreadPoolExecutor\")?;\n let executor = executor_class.call1((8,))?; // max_workers=8\n\n // Structure to track futures: (key, future_or_futures, is_list)\n let mut futs: std::collections::HashMap>, bool)> =\n std::collections::HashMap::new();\n let mut all_futures: Vec> = Vec::new();\n\n for (key, is_list) in &url_path_keys {\n let value = payload.get_item(key)?.ok_or_else(|| {\n pyo3::exceptions::PyKeyError::new_err(format!(\n \"Input key '{}' disappeared during processing\",\n key\n ))\n })?;\n\n if *is_list {\n let list = value.extract::>()?;\n let mut futures_for_key = Vec::new();\n for item in list.iter() {\n let convert_method = item.getattr(\"convert\")?;\n let future = executor.call_method1(\"submit\", (convert_method,))?;\n futures_for_key.push(future.clone());\n all_futures.push(future);\n }\n futs.insert(key.clone(), (futures_for_key, true));\n } else {\n let convert_method = value.getattr(\"convert\")?;\n let future = executor.call_method1(\"submit\", (convert_method,))?;\n all_futures.push(future.clone());\n futs.insert(key.clone(), (vec![future], false));\n }\n }\n\n // Wait for all futures\n let future_list = pyo3::types::PyList::new(py, &all_futures)?;\n let wait_fn = concurrent_futures.getattr(\"wait\")?;\n let wait_result = wait_fn.call1((&future_list,))?;\n let done = wait_result.get_item(0)?;\n let not_done = wait_result.get_item(1)?;\n\n // Check for failures\n let not_done_len: usize = not_done.len()?;\n if not_done_len > 0 {\n // Cancel remaining and find the exception\n for item in not_done.try_iter()? {\n let fut = item?;\n let _ = fut.call_method0(\"cancel\");\n }\n // Find and raise the exception\n for item in done.try_iter()? {\n let fut = item?;\n fut.call_method0(\"result\")?; // raises if future finished with exception\n }\n return Err(PyErr::new::(\n \"Download failed\",\n ));\n }\n\n // All downloads complete - replace URLPath with local Path in payload\n // Collect the Path objects for cleanup\n let mut cleanup_paths: Vec = Vec::new();\n\n for (key, (futures, is_list)) in futs {\n if is_list {\n let mut results = Vec::new();\n for fut in futures {\n let result = fut.call_method0(\"result\")?;\n cleanup_paths.push(result.clone().unbind());\n results.push(result);\n }\n let result_list = pyo3::types::PyList::new(py, &results)?;\n payload.set_item(&key, result_list)?;\n } else {\n let result = futures[0].call_method0(\"result\")?;\n cleanup_paths.push(result.clone().unbind());\n payload.set_item(&key, result)?;\n }\n }\n\n // Shutdown executor\n executor.call_method0(\"shutdown\")?;\n\n tracing::debug!(\n \"URLPath downloads complete, {} paths to cleanup\",\n cleanup_paths.len()\n );\n Ok(cleanup_paths)\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n /// Helper: define a Python function with the given parameter annotations and\n /// return the set of field names that `detect_file_fields` identifies as File-typed.\n fn file_fields_for(py_func_src: &str) -> HashSet {\n pyo3::Python::initialize();\n Python::attach(|py| {\n // Ensure cog.types.File is importable\n py.run(c\"import cog.types\", None, None)\n .expect(\"cog.types must be importable for tests\");\n\n let locals = PyDict::new(py);\n py.run(\n &std::ffi::CString::new(py_func_src).unwrap(),\n None,\n Some(&locals),\n )\n .expect(\"failed to define test function\");\n\n let func = locals.get_item(\"func\").unwrap().unwrap();\n detect_file_fields(py, &func).expect(\"detect_file_fields failed\")\n })\n }\n\n #[test]\n #[ignore] // Requires cog Python package in PYTHONPATH\n fn detect_direct_file() {\n let fields = file_fields_for(\"from cog import File\\ndef func(a: File, b: str): ...\");\n assert!(fields.contains(\"a\"), \"direct File annotation not detected\");\n assert!(!fields.contains(\"b\"), \"str incorrectly flagged as File\");\n }\n\n #[test]\n #[ignore] // Requires cog Python package in PYTHONPATH\n fn detect_list_file() {\n let fields = file_fields_for(\"from cog import File\\ndef func(a: list[File]): ...\");\n assert!(fields.contains(\"a\"), \"list[File] annotation not detected\");\n }\n\n #[test]\n #[ignore] // Requires cog Python package in PYTHONPATH\n fn detect_optional_file() {\n let fields = file_fields_for(\n \"from typing import Optional\\nfrom cog import File\\ndef func(a: Optional[File]): ...\",\n );\n assert!(\n fields.contains(\"a\"),\n \"Optional[File] annotation not detected\"\n );\n }\n\n #[test]\n #[ignore] // Requires cog Python package in PYTHONPATH\n fn detect_file_union_none() {\n let fields = file_fields_for(\n \"from typing import Union\\nfrom cog import File\\ndef func(a: Union[File, None]): ...\",\n );\n assert!(\n fields.contains(\"a\"),\n \"File | None / Union[File, None] annotation not detected\"\n );\n }\n\n #[test]\n #[ignore] // Requires cog Python package in PYTHONPATH\n fn detect_optional_list_file() {\n let fields = file_fields_for(\n \"from typing import Optional\\nfrom cog import File\\ndef func(a: Optional[list[File]]): ...\",\n );\n assert!(\n fields.contains(\"a\"),\n \"Optional[list[File]] annotation not detected\"\n );\n }\n\n #[test]\n #[ignore] // Requires cog Python package in PYTHONPATH\n fn non_file_types_not_detected() {\n let fields = file_fields_for(\n \"from pathlib import Path\\nfrom typing import Optional\\ndef func(a: str, b: int, c: Optional[str], d: Path): ...\",\n );\n assert!(\n fields.is_empty(),\n \"non-File types incorrectly detected: {:?}\",\n fields\n );\n }\n}\n" }, { "path": "crates/coglet-python/src/lib.rs", "content": "//! coglet-python: PyO3 bindings for coglet.\n\nmod audit;\nmod cancel;\nmod input;\nmod log_writer;\nmod metric_scope;\nmod output;\nmod predictor;\nmod worker_bridge;\n\nuse std::sync::Arc;\nuse std::sync::atomic::{AtomicBool, Ordering};\n\nuse pyo3::prelude::*;\nuse pyo3_stub_gen::derive::*;\nuse tracing::{debug, error, info, warn};\nuse tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt};\n\n// Define stub info gatherer for generating .pyi files\npyo3_stub_gen::define_stub_info_gatherer!(stub_info);\n\n// Module-level attributes (pyo3-stub-gen can't see m.add() calls).\n// Uses \"coglet\" because that's the module key in StubInfo for the native module.\npyo3_stub_gen::module_variable!(\"coglet\", \"__version__\", &str);\npyo3_stub_gen::module_variable!(\"coglet\", \"__build__\", BuildInfo);\npyo3_stub_gen::module_variable!(\"coglet\", \"server\", CogletServer);\n\nuse coglet_core::{\n Health, PredictionService, SetupResult, VersionInfo,\n transport::{ServerConfig, serve as http_serve},\n};\n\n/// Global flag: true when running inside a worker subprocess.\nstatic ACTIVE: AtomicBool = AtomicBool::new(false);\n\n/// Frozen build metadata exposed as `coglet.__build__`.\n#[gen_stub_pyclass]\n#[pyclass(name = \"BuildInfo\", module = \"coglet\", frozen)]\npub struct BuildInfo {\n #[pyo3(get)]\n version: String,\n #[pyo3(get)]\n git_sha: String,\n #[pyo3(get)]\n dirty: bool,\n #[pyo3(get)]\n build_time: String,\n #[pyo3(get)]\n rustc_version: String,\n}\n\n#[gen_stub_pymethods]\n#[pymethods]\nimpl BuildInfo {\n fn __repr__(&self) -> String {\n format!(\n \"BuildInfo(version='{}', git_sha='{}', dirty={}, build_time='{}', rustc_version='{}')\",\n self.version,\n self.git_sha,\n if self.dirty { \"True\" } else { \"False\" },\n self.build_time,\n self.rustc_version\n )\n }\n}\n\nimpl BuildInfo {\n fn new() -> Self {\n Self {\n version: env!(\"COGLET_PEP440_VERSION\").to_string(),\n git_sha: env!(\"COGLET_GIT_SHA\").to_string(),\n dirty: env!(\"COGLET_GIT_DIRTY\") == \"true\",\n build_time: env!(\"COGLET_BUILD_TIME\").to_string(),\n rustc_version: env!(\"COGLET_RUSTC_VERSION\").to_string(),\n }\n }\n\n /// Git SHA with optional `-dirty` suffix.\n fn sha_display(&self) -> String {\n if self.dirty {\n format!(\"{}-dirty\", self.git_sha)\n } else {\n self.git_sha.clone()\n }\n }\n}\n\nfn set_active() {\n ACTIVE.store(true, Ordering::SeqCst);\n}\n\n/// Initialize tracing with COG_LOG_LEVEL and LOG_FORMAT support.\n/// Returns optional receiver for draining setup logs.\nfn init_tracing(\n _to_stderr: bool,\n setup_log_tx: Option>,\n) -> Option> {\n let filter = if std::env::var(\"RUST_LOG\").is_ok() {\n EnvFilter::from_default_env()\n } else {\n let base_level = match std::env::var(\"COG_LOG_LEVEL\").as_deref() {\n Ok(\"debug\") => \"debug\",\n Ok(\"warn\") | Ok(\"warning\") => \"warn\",\n Ok(\"error\") => \"error\",\n _ => \"info\",\n };\n\n let filter_str = format!(\n \"coglet={level},coglet::setup=info,coglet::user=info,coglet_worker={level},coglet_worker::schema=off,coglet_worker::protocol=off\",\n level = base_level\n );\n\n EnvFilter::new(filter_str)\n };\n\n let use_json = std::env::var(\"LOG_FORMAT\").as_deref() != Ok(\"console\");\n\n if let Some(tx) = setup_log_tx {\n let accumulator = coglet_core::SetupLogAccumulator::new(tx);\n\n if use_json {\n let subscriber = tracing_subscriber::registry()\n .with(filter)\n .with(accumulator)\n .with(fmt::layer().json().with_writer(std::io::stderr));\n let _ = subscriber.try_init();\n } else {\n let subscriber = tracing_subscriber::registry()\n .with(filter)\n .with(accumulator)\n .with(fmt::layer().with_writer(std::io::stderr));\n let _ = subscriber.try_init();\n }\n None\n } else {\n if use_json {\n let subscriber = tracing_subscriber::registry()\n .with(filter)\n .with(fmt::layer().json().with_writer(std::io::stderr));\n let _ = subscriber.try_init();\n } else {\n let subscriber = tracing_subscriber::registry()\n .with(filter)\n .with(fmt::layer().with_writer(std::io::stderr));\n let _ = subscriber.try_init();\n }\n None\n }\n}\n\nfn detect_version(py: Python<'_>, build: &BuildInfo) -> VersionInfo {\n let mut version = VersionInfo::new()\n .with_git_sha(build.sha_display())\n .with_build_time(build.build_time.clone());\n\n if let Ok(sys) = py.import(\"sys\")\n && let Ok(py_version) = sys.getattr(\"version\")\n && let Ok(v) = py_version.extract::()\n {\n let short_version = v.split_whitespace().next().unwrap_or(&v);\n version = version.with_python(short_version.to_string());\n }\n\n if let Ok(cog) = py.import(\"cog\")\n && let Ok(cog_version) = cog.getattr(\"__version__\")\n && let Ok(v) = cog_version.extract::()\n {\n version = version.with_python_sdk(v);\n }\n\n version\n}\n\nfn read_max_concurrency() -> usize {\n match std::env::var(\"COG_MAX_CONCURRENCY\") {\n Ok(val) => val.parse::().unwrap_or(1),\n Err(_) => 1,\n }\n}\n\nfn read_setup_timeout() -> Option {\n match std::env::var(\"COG_SETUP_TIMEOUT\") {\n Ok(val) => match val.parse::() {\n Ok(0) => {\n warn!(\"COG_SETUP_TIMEOUT=0 would cause immediate timeout, ignoring\");\n None\n }\n Ok(secs) => Some(std::time::Duration::from_secs(secs)),\n Err(e) => {\n warn!(\n value = %val,\n error = %e,\n \"Invalid COG_SETUP_TIMEOUT value, ignoring (no timeout will be applied)\"\n );\n None\n }\n },\n Err(_) => None,\n }\n}\n\n// =============================================================================\n// coglet.server — frozen Server object with serve() and active property\n// =============================================================================\n\n/// The coglet prediction server.\n///\n/// Access via `coglet.server`. Frozen — attributes cannot be set or deleted.\n///\n/// - `coglet.server.active` — `True` when running inside a worker subprocess\n/// - `coglet.server.serve(...)` — start the HTTP prediction server (blocking)\n#[gen_stub_pyclass]\n#[pyclass(name = \"Server\", module = \"coglet\", frozen)]\npub struct CogletServer {}\n\n#[gen_stub_pymethods]\n#[pymethods]\nimpl CogletServer {\n /// `True` when running inside a coglet worker subprocess.\n #[getter]\n fn active(&self) -> bool {\n ACTIVE.load(Ordering::SeqCst)\n }\n\n /// Start the HTTP prediction server. Blocks until shutdown.\n #[allow(clippy::too_many_arguments)]\n #[pyo3(signature = (predictor_ref=None, host=\"0.0.0.0\".to_string(), port=5000, await_explicit_shutdown=false, is_train=false, output_temp_dir_base=\"/tmp/coglet/output\".to_string(), upload_url=None))]\n fn serve(\n &self,\n py: Python<'_>,\n predictor_ref: Option,\n host: String,\n port: u16,\n await_explicit_shutdown: bool,\n is_train: bool,\n output_temp_dir_base: String,\n upload_url: Option,\n ) -> PyResult<()> {\n serve_impl(\n py,\n predictor_ref,\n host,\n port,\n await_explicit_shutdown,\n is_train,\n output_temp_dir_base,\n upload_url,\n )\n }\n\n /// Worker subprocess entry point. Called by the orchestrator.\n ///\n /// Sets the active flag, installs log writers and audit hooks,\n /// then enters the worker event loop.\n #[pyo3(name = \"_run_worker\", signature = ())]\n fn run_worker(&self, py: Python<'_>) -> PyResult<()> {\n set_active();\n\n // Install SlotLogWriters for ContextVar-based log routing\n log_writer::install_slot_log_writers(py)?;\n\n // Install audit hook to protect stdout/stderr from user replacement\n if let Err(e) = audit::install_audit_hook(py) {\n warn!(error = %e, \"Failed to install audit hook, stdout/stderr protection disabled\");\n }\n\n info!(target: \"coglet::worker\", \"Worker subprocess starting, waiting for Init message\");\n\n py.detach(|| {\n let rt = tokio::runtime::Runtime::new()\n .map_err(|e| PyErr::new::(e.to_string()))?;\n\n rt.block_on(async {\n run_worker_with_init()\n .await\n .map_err(|e| PyErr::new::(e.to_string()))\n })\n })\n }\n\n fn __repr__(&self) -> &'static str {\n \"coglet.server\"\n }\n}\n\n#[allow(clippy::too_many_arguments)]\nfn serve_impl(\n py: Python<'_>,\n predictor_ref: Option,\n host: String,\n port: u16,\n await_explicit_shutdown: bool,\n is_train: bool,\n _output_temp_dir_base: String,\n upload_url: Option,\n) -> PyResult<()> {\n let (setup_log_tx, setup_log_rx) = tokio::sync::mpsc::unbounded_channel();\n init_tracing(false, Some(setup_log_tx));\n\n let build = BuildInfo::new();\n info!(\n \"coglet {} ({}, built {}{})\",\n env!(\"CARGO_PKG_VERSION\"),\n build.sha_display(),\n build.build_time,\n if cfg!(debug_assertions) {\n \", debug\"\n } else {\n \"\"\n },\n );\n\n let config = ServerConfig {\n host,\n port,\n await_explicit_shutdown,\n };\n\n // Install Python SIGTERM handler if await_explicit_shutdown\n if await_explicit_shutdown {\n let signal_module = py.import(\"signal\")?;\n let sigterm = signal_module.getattr(\"SIGTERM\")?;\n let sig_ign = signal_module.getattr(\"SIG_IGN\")?;\n signal_module.call_method1(\"signal\", (sigterm, sig_ign))?;\n info!(\"await_explicit_shutdown: installed SIGTERM ignore handler\");\n }\n\n let version = detect_version(py, &build);\n info!(\n \"python sdk {}\",\n version.python_sdk.as_deref().unwrap_or(\"unknown\")\n );\n info!(\"python {}\", version.python.as_deref().unwrap_or(\"unknown\"));\n\n let Some(pred_ref) = predictor_ref else {\n info!(\"No predictor specified, serving health endpoints only\");\n let service = Arc::new(\n PredictionService::new_no_pool()\n .with_health(Health::Unknown)\n .with_version(version),\n );\n return py.detach(|| {\n let rt = tokio::runtime::Runtime::new()\n .map_err(|e| PyErr::new::(e.to_string()))?;\n rt.block_on(async {\n http_serve(config, service)\n .await\n .map_err(|e| PyErr::new::(e.to_string()))\n })\n });\n };\n\n info!(predictor_ref = %pred_ref, is_train, \"Using subprocess isolation\");\n serve_subprocess(\n py,\n pred_ref,\n config,\n version,\n is_train,\n setup_log_rx,\n upload_url,\n )\n}\n\nfn serve_subprocess(\n py: Python<'_>,\n pred_ref: String,\n config: ServerConfig,\n version: VersionInfo,\n is_train: bool,\n mut setup_log_rx: tokio::sync::mpsc::UnboundedReceiver,\n upload_url: Option,\n) -> PyResult<()> {\n let max_concurrency = read_max_concurrency();\n info!(\n max_concurrency,\n \"Configuring subprocess worker via orchestrator\"\n );\n\n let setup_timeout = read_setup_timeout();\n debug!(\n setup_timeout_secs = setup_timeout.map(|d| d.as_secs()),\n is_train, \"Orchestrator configuration\"\n );\n let orch_config = coglet_core::orchestrator::OrchestratorConfig::new(pred_ref)\n .with_num_slots(max_concurrency)\n .with_train(is_train)\n .with_upload_url(upload_url)\n .with_setup_timeout(setup_timeout);\n\n let service = Arc::new(\n PredictionService::new_no_pool()\n .with_health(Health::Starting)\n .with_version(version),\n );\n\n let service_clone = Arc::clone(&service);\n py.detach(|| {\n let rt = tokio::runtime::Runtime::new()\n .map_err(|e| PyErr::new::(e.to_string()))?;\n\n rt.block_on(async {\n let setup_result = SetupResult::starting();\n service_clone.set_setup_result(setup_result.clone()).await;\n\n let setup_service = Arc::clone(&service_clone);\n tokio::spawn(async move {\n info!(\"Spawning worker subprocess\");\n let spawn_start = std::time::Instant::now();\n match coglet_core::orchestrator::spawn_worker(orch_config, &mut setup_log_rx).await\n {\n Ok(ready) => {\n let spawn_elapsed = spawn_start.elapsed();\n debug!(\n elapsed_ms = spawn_elapsed.as_millis() as u64,\n \"Worker ready, configuring service\"\n );\n\n let num_slots = ready.handle.slot_ids().len();\n debug!(num_slots, \"Setting up orchestrator on service\");\n\n setup_service\n .set_orchestrator(ready.pool, Arc::new(ready.handle))\n .await;\n debug!(\"Transitioning health to Ready\");\n setup_service.set_health(Health::Ready).await;\n\n if let Some(s) = ready.schema {\n debug!(\"Setting OpenAPI schema on service\");\n setup_service.set_schema(s).await;\n } else {\n debug!(\"No OpenAPI schema provided by worker\");\n }\n\n let mode = if is_train { \"train\" } else { \"predict\" };\n info!(num_slots, mode, \"Server ready\");\n\n // Drain final logs (includes \"Server ready\" above)\n let final_logs = coglet_core::drain_accumulated_logs(&mut setup_log_rx);\n debug!(\n initial_logs_len = ready.setup_logs.len(),\n final_logs_len = final_logs.len(),\n \"Drained setup logs\"\n );\n drop(setup_log_rx);\n\n // Combine initial + final logs\n let complete_logs = ready.setup_logs + &final_logs;\n setup_service\n .set_setup_result(setup_result.succeeded(complete_logs))\n .await;\n\n info!(\"Setup complete, now accepting requests\");\n }\n Err(e) => {\n let spawn_elapsed = spawn_start.elapsed();\n error!(\n error = %e,\n elapsed_ms = spawn_elapsed.as_millis() as u64,\n \"Worker initialization failed\"\n );\n debug!(\"Transitioning health to SetupFailed\");\n setup_service.set_health(Health::SetupFailed).await;\n setup_service\n .set_setup_result(setup_result.failed(e.to_string()))\n .await;\n }\n }\n });\n\n http_serve(config, service_clone)\n .await\n .map_err(|e| PyErr::new::(e.to_string()))\n })\n })\n}\n\nasync fn run_worker_with_init() -> Result<(), String> {\n use coglet_core::bridge::codec::JsonCodec;\n use coglet_core::bridge::protocol::ControlRequest;\n use futures::StreamExt;\n use tokio::io::stdin;\n use tokio_util::codec::FramedRead;\n\n let mut ctrl_reader = FramedRead::new(stdin(), JsonCodec::::new());\n\n let init_msg = ctrl_reader\n .next()\n .await\n .ok_or_else(|| \"stdin closed before Init received\".to_string())?\n .map_err(|e| format!(\"Failed to read Init: {}\", e))?;\n\n let (predictor_ref, num_slots, transport_info, is_train, _is_async) = match init_msg {\n ControlRequest::Init {\n predictor_ref,\n num_slots,\n transport_info,\n is_train,\n is_async,\n } => (predictor_ref, num_slots, transport_info, is_train, is_async),\n other => {\n return Err(format!(\"Expected Init message, got: {:?}\", other));\n }\n };\n\n info!(predictor_ref = %predictor_ref, num_slots, is_train, \"Init received, connecting to transport\");\n\n let handler = Arc::new(if is_train {\n worker_bridge::PythonPredictHandler::new_train(predictor_ref)\n .map_err(|e| format!(\"Failed to create handler: {}\", e))?\n } else {\n worker_bridge::PythonPredictHandler::new(predictor_ref)\n .map_err(|e| format!(\"Failed to create handler: {}\", e))?\n });\n\n // Setup log hook: registers a global sender for control channel logs\n // This lives for the entire worker lifetime (setup + subprocess output)\n let setup_log_hook: coglet_core::SetupLogHook = Box::new(|tx| {\n let sender = Arc::new(log_writer::ControlChannelLogSender::new(tx));\n log_writer::register_control_channel_sender(sender);\n // Cleanup is a no-op: sender stays registered for worker lifetime\n Box::new(|| {})\n });\n\n let config = coglet_core::WorkerConfig {\n num_slots,\n setup_log_hook: Some(setup_log_hook),\n };\n\n coglet_core::run_worker(handler, config, transport_info)\n .await\n .map_err(|e| format!(\"Worker error: {}\", e))\n}\n\n// =============================================================================\n// Module init\n// =============================================================================\n\n#[pymodule]\n#[pyo3(name = \"_impl\")]\nfn coglet(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {\n // Control what `from ._impl import *` exports into coglet/__init__.py\n m.add(\"__all__\", vec![\"__version__\", \"__build__\", \"server\"])?;\n\n // Static metadata\n m.add(\"__version__\", env!(\"COGLET_PEP440_VERSION\"))?;\n m.add(\"__build__\", BuildInfo::new())?;\n\n // Frozen server object\n m.add(\"server\", CogletServer {})?;\n\n // CancelationException — a BaseException subclass for prediction cancellation.\n // Re-exported through coglet → cog.exceptions → cog.CancelationException.\n m.add(\n \"CancelationException\",\n py.get_type::(),\n )?;\n\n // _sdk submodule — internal Python runtime integration classes\n let sdk = PyModule::new(py, \"_sdk\")?;\n sdk.setattr(\n \"__doc__\",\n \"Internal SDK runtime integration for coglet.\\n\\\n \\n\\\n This submodule contains Rust-backed classes that integrate coglet with\\n\\\n the Python runtime (I/O routing, audit hooks, log streaming). These are\\n\\\n implementation details used by the cog SDK — not part of the public API.\",\n )?;\n sdk.add_class::()?;\n sdk.add_class::()?;\n sdk.add_class::()?;\n sdk.add_class::()?;\n sdk.add_function(wrap_pyfunction!(metric_scope::py_current_scope, &sdk)?)?;\n m.add_submodule(&sdk)?;\n\n Ok(())\n}\n" }, { "path": "crates/coglet-python/src/log_writer.rs", "content": "//! Log routing via prediction_id ContextVar.\n//!\n//! Architecture:\n//! - Rust owns a ContextVar `_coglet_prediction_id` that holds the current prediction ID\n//! - Rust maintains a registry mapping prediction_id -> SlotSender\n//! - SlotLogWriter reads the ContextVar to route logs to the correct sender\n//!\n//! This design supports:\n//! - Async predictions with proper per-task isolation (ContextVar is task-local)\n//! - Orphan task detection (prediction completed but task still running)\n//! - Slot reuse safety (new prediction = new ID, old tasks can't pollute)\n//! - Setup logs routed through control channel before predictions start\n//!\n//! The ContextVar is private (`_coglet_` prefix). Users who need the prediction ID\n//! should use the public API (e.g., `cog.current_prediction_id()`) which we'll\n//! inject onto the cog namespace later.\n\nuse std::collections::HashMap;\nuse std::sync::{Arc, Mutex, OnceLock};\n\nuse pyo3::prelude::*;\nuse pyo3_stub_gen::derive::*;\n\nuse coglet_core::bridge::protocol::{ControlResponse, LogSource};\nuse coglet_core::worker::SlotSender;\nuse tokio::sync::mpsc::Sender;\n\n// ============================================================================\n// Rust-owned ContextVar for prediction routing\n// ============================================================================\n\n/// The Rust-owned ContextVar instance. Created once, lives forever.\n/// Named `cog_prediction_id` - documented as internal, don't modify.\nstatic PREDICTION_CONTEXTVAR: OnceLock> = OnceLock::new();\n\n/// Registry mapping prediction_id -> SlotSender.\n/// When a prediction starts, we register the sender.\n/// When SlotLogWriter.write() is called, we look up the sender here.\nstatic PREDICTION_REGISTRY: OnceLock>>> = OnceLock::new();\n\n/// Current sync prediction ID.\n/// For sync predictions (single slot, blocking), there's exactly one active prediction.\n/// ContextVars don't work across separate Python::attach calls, so we use this.\n/// Protected by mutex since it's accessed from Python callbacks.\nstatic SYNC_PREDICTION_ID: OnceLock>> = OnceLock::new();\n\nfn get_sync_prediction_id_slot() -> &'static Mutex> {\n SYNC_PREDICTION_ID.get_or_init(|| Mutex::new(None))\n}\n\n/// Control channel log sender - used when outside prediction context.\n/// Set by worker before setup(), lives for entire worker lifetime.\nstatic CONTROL_CHANNEL_LOG_SENDER: OnceLock>>> =\n OnceLock::new();\n\nfn get_registry() -> &'static Mutex>> {\n PREDICTION_REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))\n}\n\nfn get_control_channel_sender_slot() -> &'static Mutex>> {\n CONTROL_CHANNEL_LOG_SENDER.get_or_init(|| Mutex::new(None))\n}\n\n// ============================================================================\n// ControlChannelLogSender - sends logs via control channel\n// ============================================================================\n\n/// Sender for logs that go through the control channel.\n/// Used for Python logs during setup and subprocess output throughout worker lifetime.\npub struct ControlChannelLogSender {\n tx: Sender,\n}\n\nimpl ControlChannelLogSender {\n /// Create a new control channel log sender.\n pub fn new(tx: Sender) -> Self {\n Self { tx }\n }\n\n /// Try to send a log message.\n ///\n /// Uses try_send() to avoid blocking (called from Python code on tokio runtime threads).\n /// If the channel is full, the log is dropped and counted for periodic reporting.\n pub fn try_send_log(&self, source: LogSource, data: &str) {\n if self\n .tx\n .try_send(ControlResponse::Log {\n source,\n data: data.to_string(),\n })\n .is_err()\n {\n coglet_core::worker::increment_dropped_log_count();\n }\n }\n}\n\n// NOTE: All mutex locks in the worker use .expect().\n//\n// If a mutex is poisoned (another thread panicked while holding it), the worker\n// is in an unrecoverable state. We cannot safely continue because:\n// - Log routing shares channels with prediction updates\n// - Prediction→slot mappings could be inconsistent\n// - Continuing risks cross-prediction data bleed\n//\n// The panic hook installed by coglet_core::worker sends a Fatal IPC message\n// to the parent (which poisons all slots) and aborts the process.\n\n/// Register the control channel log sender.\n/// Called by worker before setup().\npub fn register_control_channel_sender(sender: Arc) {\n let mut slot = get_control_channel_sender_slot()\n .lock()\n .expect(\"control_channel_sender mutex poisoned\");\n *slot = Some(sender);\n}\n\n/// Unregister the control channel log sender.\n/// Called by worker when shutting down (not after setup).\n#[allow(dead_code)]\npub fn unregister_control_channel_sender() {\n let mut slot = get_control_channel_sender_slot()\n .lock()\n .expect(\"control_channel_sender mutex poisoned\");\n *slot = None;\n}\n\n/// Get the control channel log sender if registered.\nfn get_control_channel_sender() -> Option> {\n let slot = get_control_channel_sender_slot()\n .lock()\n .expect(\"control_channel_sender mutex poisoned\");\n slot.clone()\n}\n\n/// Get or create the Rust-owned ContextVar.\n///\n/// This returns the same ContextVar instance used by SlotLogWriter for log routing.\n/// Public so predictor.rs can pass it to async coroutine wrappers.\npub fn get_prediction_contextvar(py: Python<'_>) -> PyResult<&'static Py> {\n if let Some(cv) = PREDICTION_CONTEXTVAR.get() {\n return Ok(cv);\n }\n\n let contextvars = py.import(\"contextvars\")?;\n let cv = contextvars.call_method1(\"ContextVar\", (\"_coglet_prediction_id\",))?;\n\n // Try to store it. Race is fine - if another thread won, use their value.\n match PREDICTION_CONTEXTVAR.set(cv.unbind()) {\n Ok(()) => {}\n Err(_already_set) => {\n // Another thread initialized it first - that's fine\n }\n }\n\n // This should always succeed now - either we set it or another thread did.\n PREDICTION_CONTEXTVAR.get().ok_or_else(|| {\n pyo3::exceptions::PyRuntimeError::new_err(\n \"Failed to initialize prediction context variable\",\n )\n })\n}\n\n/// Register a SlotSender for a prediction ID.\n/// Called when starting a prediction.\npub fn register_prediction(prediction_id: String, sender: Arc) {\n let mut registry = get_registry()\n .lock()\n .expect(\"prediction_registry mutex poisoned\");\n tracing::trace!(%prediction_id, \"Registering prediction sender\");\n registry.insert(prediction_id, sender);\n}\n\n/// Unregister a prediction ID.\n/// Called when prediction completes.\npub fn unregister_prediction(prediction_id: &str) {\n let mut registry = get_registry()\n .lock()\n .expect(\"prediction_registry mutex poisoned\");\n registry.remove(prediction_id);\n\n // Clear sync prediction ID if it matches\n let mut slot = get_sync_prediction_id_slot()\n .lock()\n .expect(\"sync_prediction_id mutex poisoned\");\n if slot.as_deref() == Some(prediction_id) {\n *slot = None;\n }\n}\n\n/// Get the SlotSender for a prediction ID.\nfn get_prediction_sender(prediction_id: &str) -> Option> {\n let registry = get_registry()\n .lock()\n .expect(\"prediction_registry mutex poisoned\");\n registry.get(prediction_id).cloned()\n}\n\n/// Set the current prediction ID in the ContextVar (for async).\n/// Returns a token that can be used to reset (for explicit cleanup).\npub fn set_current_prediction(py: Python<'_>, prediction_id: &str) -> PyResult> {\n // Set ContextVar for async predictions\n let cv = get_prediction_contextvar(py)?;\n let token = cv.call_method1(py, \"set\", (prediction_id,))?;\n Ok(token)\n}\n\n/// Set the current sync prediction ID (for sync predictions only).\n/// Call this before running a sync prediction, clear after.\npub fn set_sync_prediction_id(prediction_id: Option<&str>) {\n let mut slot = get_sync_prediction_id_slot()\n .lock()\n .expect(\"sync_prediction_id mutex poisoned\");\n *slot = prediction_id.map(|s| s.to_string());\n}\n\n/// Get the current prediction ID from sync static or ContextVar.\n/// Returns None if not set (outside prediction context).\nfn get_current_prediction_id(py: Python<'_>) -> PyResult> {\n // First check sync prediction static (works for sync predictions)\n {\n let slot = get_sync_prediction_id_slot()\n .lock()\n .expect(\"sync_prediction_id mutex poisoned\");\n if let Some(ref prediction_id) = *slot {\n tracing::trace!(%prediction_id, \"Sync prediction ID found\");\n return Ok(Some(prediction_id.clone()));\n }\n }\n\n // Fall back to ContextVar (works for async predictions)\n let cv = get_prediction_contextvar(py)?;\n\n // Try to get the value - returns the value or raises LookupError\n match cv.call_method0(py, \"get\") {\n Ok(val) => {\n let prediction_id: String = val.extract(py)?;\n tracing::trace!(%prediction_id, \"ContextVar lookup succeeded\");\n Ok(Some(prediction_id))\n }\n Err(e) if e.is_instance_of::(py) => {\n // ContextVar not set - outside prediction context\n Ok(None)\n }\n Err(e) => Err(e),\n }\n}\n\n// ============================================================================\n// SlotLogWriter - routes via ContextVar lookup\n// ============================================================================\n\n/// A Python file-like object that routes writes via the prediction_id ContextVar.\n///\n/// This is installed as sys.stdout/stderr once at worker startup.\n/// Each write looks up the current prediction_id from the ContextVar and routes\n/// to the appropriate SlotSender.\n///\n/// If no prediction_id is set, or the prediction has completed (orphan task),\n/// writes go to tracing (logged as orphan).\n///\n/// Uses line buffering: accumulates writes until a newline is received, then\n/// emits complete lines. This coalesces Python's print() which does separate\n/// writes for content and newline.\n#[gen_stub_pyclass]\n#[pyclass(name = \"_SlotLogWriter\", module = \"coglet._sdk\")]\npub struct SlotLogWriter {\n /// Which stream this captures (stdout or stderr).\n source: LogSource,\n /// Original stream (used for delegation of methods like isatty, fileno).\n original: Py,\n /// Whether writes should be ignored (used after errors).\n #[pyo3(get)]\n closed: bool,\n /// Line buffer for coalescing writes into complete lines.\n line_buffer: Mutex,\n}\n\n#[gen_stub_pymethods]\n#[pymethods]\nimpl SlotLogWriter {\n /// Write data, routing to the appropriate destination.\n ///\n /// Uses line buffering: accumulates data until a newline is received, then\n /// emits complete lines. This coalesces Python's print() which does separate\n /// writes for content and the trailing newline.\n ///\n /// Priority for routing:\n /// 1. If inside a prediction (ContextVar set), route to slot sender\n /// 2. If setup sender registered, route to control channel \n /// 3. Fall back to stderr (for orphan tasks or unexpected cases)\n fn write(&self, py: Python<'_>, data: &str) -> PyResult {\n if self.closed || data.is_empty() {\n return Ok(data.len());\n }\n\n let len = data.len();\n\n // Append to line buffer and extract complete lines\n let complete = {\n let mut buffer = self.line_buffer.lock().expect(\"line_buffer mutex poisoned\");\n buffer.push_str(data);\n\n // Check if we have complete lines to emit\n if let Some(last_newline) = buffer.rfind('\\n') {\n // Extract complete lines (including the newline)\n let complete = buffer[..=last_newline].to_string();\n // Keep remainder in buffer\n let remainder = buffer[last_newline + 1..].to_string();\n *buffer = remainder;\n Some(complete)\n } else {\n None\n }\n };\n\n // Emit complete lines (outside lock)\n if let Some(complete) = complete {\n self.emit_data(py, &complete)?;\n }\n\n Ok(len)\n }\n\n /// Emit data to the appropriate destination.\n fn emit_data(&self, py: Python<'_>, data: &str) -> PyResult<()> {\n if data.is_empty() {\n return Ok(());\n }\n\n // Try to get current prediction from ContextVar\n match get_current_prediction_id(py)? {\n Some(prediction_id) => {\n // Have prediction ID - check if still active\n if let Some(sender) = get_prediction_sender(&prediction_id) {\n // Active prediction - route to slot\n tracing::trace!(\n prediction_id = %prediction_id,\n source = ?self.source,\n bytes = data.len(),\n \"Log routed to slot\"\n );\n sender\n .send_log(self.source, data)\n .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?;\n } else {\n // Orphan task - prediction completed but task still running\n tracing::trace!(\n prediction_id = %prediction_id,\n source = ?self.source,\n \"Orphan log (prediction completed)\"\n );\n self.write_outside_prediction(py, data)?;\n }\n }\n None => {\n // Outside prediction context\n // Try setup sender (for setup logs), then fallback to stderr\n self.write_outside_prediction(py, data)?;\n }\n }\n\n Ok(())\n }\n\n /// Flush the stream.\n ///\n /// Emits any buffered content that hasn't been terminated with a newline.\n fn flush(&self, py: Python<'_>) -> PyResult<()> {\n // Emit any buffered content\n let buffered = {\n let mut buffer = self.line_buffer.lock().expect(\"line_buffer mutex poisoned\");\n std::mem::take(&mut *buffer)\n };\n if !buffered.is_empty() {\n self.emit_data(py, &buffered)?;\n }\n\n // Flush the original stream\n self.original.call_method0(py, \"flush\")?;\n Ok(())\n }\n\n /// Return whether the stream is readable.\n fn readable(&self) -> bool {\n false\n }\n\n /// Return whether the stream is writable.\n fn writable(&self) -> bool {\n !self.closed\n }\n\n /// Return whether the stream is seekable.\n fn seekable(&self) -> bool {\n false\n }\n\n /// Return whether the stream is a TTY.\n fn isatty(&self, py: Python<'_>) -> PyResult {\n // Delegate to original\n let result = self.original.call_method0(py, \"isatty\")?;\n result.extract(py)\n }\n\n /// Return the file number.\n fn fileno(&self, py: Python<'_>) -> PyResult {\n // Delegate to original - needed for some libraries\n let result = self.original.call_method0(py, \"fileno\")?;\n result.extract(py)\n }\n\n /// Close the stream.\n fn close(&mut self) {\n self.closed = true;\n }\n\n /// Context manager enter.\n fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {\n slf\n }\n\n /// Context manager exit.\n fn __exit__(\n &mut self,\n _exc_type: Option<&Bound<'_, PyAny>>,\n _exc_val: Option<&Bound<'_, PyAny>>,\n _exc_tb: Option<&Bound<'_, PyAny>>,\n ) -> bool {\n false // Don't suppress exceptions\n }\n\n /// Encoding property - needed for compatibility.\n #[getter]\n fn encoding(&self, py: Python<'_>) -> PyResult> {\n match self.original.getattr(py, \"encoding\") {\n Ok(enc) => enc.extract(py),\n Err(_) => Ok(Some(\"utf-8\".to_string())),\n }\n }\n\n /// Newlines property - needed for compatibility.\n #[getter]\n fn newlines(&self) -> Option {\n None\n }\n\n /// Buffer property - some code checks for this.\n #[getter]\n fn buffer(&self, py: Python<'_>) -> PyResult> {\n // Return original's buffer if it has one, otherwise return self\n match self.original.getattr(py, \"buffer\") {\n Ok(buf) => Ok(buf),\n Err(_) => Ok(self.original.clone_ref(py)),\n }\n }\n}\n\nimpl SlotLogWriter {\n /// Create a new stdout writer.\n pub fn new_stdout(original: Py) -> Self {\n Self {\n source: LogSource::Stdout,\n original,\n closed: false,\n line_buffer: Mutex::new(String::new()),\n }\n }\n\n /// Create a new stderr writer.\n pub fn new_stderr(original: Py) -> Self {\n Self {\n source: LogSource::Stderr,\n original,\n closed: false,\n line_buffer: Mutex::new(String::new()),\n }\n }\n\n /// Write when outside prediction context.\n ///\n /// During setup: routes to control channel (for health-check).\n /// Otherwise: emits via tracing to stderr locally (not shipped).\n fn write_outside_prediction(&self, _py: Python<'_>, data: &str) -> PyResult<()> {\n // Try control channel sender (registered for worker lifetime)\n if let Some(sender) = get_control_channel_sender() {\n sender.try_send_log(self.source, data);\n tracing::trace!(\n source = ?self.source,\n bytes = data.len(),\n \"Log routed via control channel\"\n );\n return Ok(());\n }\n // Outside setup/prediction context - orphan log\n // This happens with orphan tasks or edge cases\n for line in data.lines() {\n tracing::info!(target: \"coglet::user\", \"{}\", line);\n }\n Ok(())\n }\n}\n\n// ============================================================================\n// Installation - called once at worker startup\n// ============================================================================\n\n/// Install SlotLogWriters as sys.stdout/stderr.\n/// Called once at worker startup. The writers persist for the lifetime of the process.\n/// Returns true if installation succeeded.\npub fn install_slot_log_writers(py: Python<'_>) -> PyResult {\n let sys = py.import(\"sys\")?;\n\n // Get originals\n let original_stdout = sys.getattr(\"stdout\")?.unbind();\n let original_stderr = sys.getattr(\"stderr\")?.unbind();\n\n // Create writers\n let stdout_writer = SlotLogWriter::new_stdout(original_stdout);\n let stderr_writer = SlotLogWriter::new_stderr(original_stderr);\n\n // Install\n sys.setattr(\"stdout\", stdout_writer.into_pyobject(py)?)?;\n sys.setattr(\"stderr\", stderr_writer.into_pyobject(py)?)?;\n\n // Initialize the ContextVar\n get_prediction_contextvar(py)?;\n\n tracing::debug!(\"Installed SlotLogWriters with prediction_id routing\");\n Ok(true)\n}\n\n// ============================================================================\n// PredictionLogGuard - RAII guard for prediction context\n// ============================================================================\n\n/// RAII guard that sets the current prediction in the ContextVar.\n///\n/// On creation, registers the SlotSender and sets the ContextVar.\n/// On drop, unregisters the prediction (but ContextVar reset is automatic for async).\npub struct PredictionLogGuard {\n prediction_id: String,\n #[allow(dead_code)]\n token: Py,\n}\n\nimpl PredictionLogGuard {\n /// Enter prediction context.\n ///\n /// Registers the sender and sets the ContextVar.\n pub fn enter(py: Python<'_>, prediction_id: String, sender: Arc) -> PyResult {\n // Register sender in global registry\n register_prediction(prediction_id.clone(), sender);\n\n // Set ContextVar\n let token = set_current_prediction(py, &prediction_id)?;\n\n tracing::trace!(%prediction_id, \"Entered prediction log context\");\n Ok(Self {\n prediction_id,\n token,\n })\n }\n\n /// Get the prediction ID.\n #[allow(dead_code)]\n pub fn prediction_id(&self) -> &str {\n &self.prediction_id\n }\n}\n\nimpl Drop for PredictionLogGuard {\n fn drop(&mut self) {\n // Unregister prediction - this makes orphan tasks fall back to stderr\n unregister_prediction(&self.prediction_id);\n\n // Note: We don't reset the ContextVar here because:\n // 1. For sync: the context resets naturally when the function returns\n // 2. For async: each task has its own ContextVar copy, no reset needed\n // The token is kept just in case we need explicit reset in the future.\n }\n}\n\n// ============================================================================\n// Tests\n// ============================================================================\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use coglet_core::bridge::protocol::SlotResponse;\n use tokio::sync::mpsc;\n\n #[test]\n fn registry_operations() {\n let prediction_id = \"pred_123\".to_string();\n let (tx, _rx) = mpsc::unbounded_channel();\n let sender = Arc::new(SlotSender::new(tx, std::env::temp_dir()));\n\n // Register\n register_prediction(prediction_id.clone(), sender.clone());\n assert!(get_prediction_sender(&prediction_id).is_some());\n\n // Unregister\n unregister_prediction(&prediction_id);\n assert!(get_prediction_sender(&prediction_id).is_none());\n }\n\n #[test]\n fn slot_sender_sends_log() {\n let (tx, mut rx) = mpsc::unbounded_channel();\n let sender = SlotSender::new(tx, std::env::temp_dir());\n\n sender.send_log(LogSource::Stdout, \"hello\").unwrap();\n\n let msg = rx.try_recv().unwrap();\n match msg {\n SlotResponse::Log { source, data } => {\n assert_eq!(source, LogSource::Stdout);\n assert_eq!(data, \"hello\");\n }\n _ => panic!(\"expected Log message\"),\n }\n }\n\n #[test]\n fn slot_sender_ignores_empty() {\n let (tx, mut rx) = mpsc::unbounded_channel();\n let sender = SlotSender::new(tx, std::env::temp_dir());\n\n sender.send_log(LogSource::Stderr, \"\").unwrap();\n\n // No message should be sent\n assert!(rx.try_recv().is_err());\n }\n\n #[test]\n fn slot_sender_detects_closed_channel() {\n let (tx, rx) = mpsc::unbounded_channel::();\n drop(rx); // Close receiver\n\n let sender = SlotSender::new(tx, std::env::temp_dir());\n let result = sender.send_log(LogSource::Stdout, \"hello\");\n\n assert!(result.is_err());\n }\n}\n" }, { "path": "crates/coglet-python/src/metric_scope.rs", "content": "//! Metric scope: type-safe metric recording with ContextVar routing.\n//!\n//! Two PyO3 classes:\n//! - `Scope` — the per-prediction context, obtained via `current_scope()`\n//! - `MetricRecorder` — the `scope.metrics` sub-object with type invariant\n//! enforcement, dict-style access, and accumulation modes\n//!\n//! All validation happens in Rust (PyO3, in-process). IPC sends the validated\n//! metric to the coglet server via SlotSender.\n\nuse std::collections::HashMap;\nuse std::sync::{Arc, Mutex, OnceLock};\n\nuse pyo3::exceptions::PyTypeError;\nuse pyo3::prelude::*;\nuse pyo3::types::PyDict;\nuse pyo3_stub_gen::derive::*;\n\nuse coglet_core::bridge::protocol::MetricMode;\nuse coglet_core::worker::SlotSender;\n\n// ============================================================================\n// Value type tracking for type invariant\n// ============================================================================\n\n/// Coarse type tag for enforcing the type invariant.\n/// Once a key is set with a type, it cannot be changed without deleting first.\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\nenum MetricValueType {\n Bool,\n Int,\n Float,\n Str,\n List,\n Dict,\n}\n\nimpl MetricValueType {\n /// Classify a Python object into a type tag.\n fn from_py(obj: &Bound<'_, PyAny>) -> PyResult {\n // Order matters: bool before int (bool is a subclass of int in Python)\n if obj.is_instance_of::() {\n Ok(Self::Bool)\n } else if obj.is_instance_of::() {\n Ok(Self::Int)\n } else if obj.is_instance_of::() {\n Ok(Self::Float)\n } else if obj.is_instance_of::() {\n Ok(Self::Str)\n } else if obj.is_instance_of::() {\n Ok(Self::List)\n } else if obj.is_instance_of::() {\n Ok(Self::Dict)\n } else {\n let type_name = obj.get_type().name()?.to_string();\n Err(PyTypeError::new_err(format!(\n \"Unsupported metric value type: {}. Expected bool, int, float, str, list, or dict.\",\n type_name\n )))\n }\n }\n\n fn as_str(&self) -> &'static str {\n match self {\n Self::Bool => \"bool\",\n Self::Int => \"int\",\n Self::Float => \"float\",\n Self::Str => \"str\",\n Self::List => \"list\",\n Self::Dict => \"dict\",\n }\n }\n}\n\n// ============================================================================\n// MetricRecorder — scope.metrics sub-object\n// ============================================================================\n\n/// Metric recorder with type invariant enforcement.\n///\n/// Accessed via `scope.metrics`. Supports:\n/// - `scope.metrics.record(key, value, mode=\"replace\")` — full API\n/// - `scope.metrics.delete(key)` — delete (required before type change)\n/// - `scope.metrics[key] = value` — dict-style set (replace mode)\n/// - `del scope.metrics[key]` — dict-style delete\n#[gen_stub_pyclass]\n#[pyclass(name = \"MetricRecorder\", module = \"coglet._sdk\")]\npub struct MetricRecorder {\n inner: Mutex>,\n}\n\nstruct RecorderInner {\n /// Type tag per metric key — enforces type invariant.\n types: HashMap,\n /// IPC sender to the coglet server.\n sender: Arc,\n}\n\nimpl MetricRecorder {\n pub fn new(sender: Arc) -> Self {\n Self {\n inner: Mutex::new(Some(RecorderInner {\n types: HashMap::new(),\n sender,\n })),\n }\n }\n\n pub fn noop() -> Self {\n Self {\n inner: Mutex::new(None),\n }\n }\n}\n\n#[gen_stub_pymethods]\n#[pymethods]\nimpl MetricRecorder {\n /// Record a metric value.\n ///\n /// Args:\n /// key: Metric name. Dot-separated keys (e.g. \"timing.preprocess\") create\n /// nested objects in the response.\n /// value: Must be bool, int, float, str, list, or dict. Once a key is set\n /// with a type, it cannot be changed without calling delete() first.\n /// mode: Accumulation mode — \"replace\" (default), \"incr\" (increment numeric),\n /// or \"append\" (push to array).\n #[pyo3(signature = (key, value, mode=None))]\n fn record(\n &self,\n py: Python<'_>,\n key: &str,\n value: &Bound<'_, PyAny>,\n mode: Option<&str>,\n ) -> PyResult<()> {\n let mode = parse_mode(mode)?;\n\n let mut guard = self.inner.lock().expect(\"metric_recorder mutex poisoned\");\n let Some(inner) = guard.as_mut() else {\n return Ok(()); // no-op outside prediction\n };\n\n record_impl(py, inner, key, value, mode)\n }\n\n /// Delete a metric key. Required before changing a metric's type.\n fn delete(&self, key: &str) -> PyResult<()> {\n let mut guard = self.inner.lock().expect(\"metric_recorder mutex poisoned\");\n let Some(inner) = guard.as_mut() else {\n return Ok(());\n };\n\n delete_impl(inner, key)\n }\n\n /// Dict-style set: `scope.metrics[\"key\"] = value`\n fn __setitem__(&self, py: Python<'_>, key: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {\n if value.is_none() {\n return self.delete(key);\n }\n\n let mut guard = self.inner.lock().expect(\"metric_recorder mutex poisoned\");\n let Some(inner) = guard.as_mut() else {\n return Ok(());\n };\n\n record_impl(py, inner, key, value, MetricMode::Replace)\n }\n\n /// Dict-style delete: `del scope.metrics[\"key\"]`\n fn __delitem__(&self, key: &str) -> PyResult<()> {\n self.delete(key)\n }\n\n fn __repr__(&self) -> String {\n let guard = self.inner.lock().expect(\"metric_recorder mutex poisoned\");\n match guard.as_ref() {\n Some(inner) => format!(\"MetricRecorder(keys={})\", inner.types.len()),\n None => \"MetricRecorder(inactive)\".to_string(),\n }\n }\n}\n\n// ============================================================================\n// Scope — the per-prediction context\n// ============================================================================\n\n/// Prediction scope, obtained via `current_scope()`.\n///\n/// Provides access to `scope.metrics` for recording metrics,\n/// `scope.record_metric()` as a convenience shorthand, and\n/// `scope.context` for per-prediction context passed in the request.\n#[gen_stub_pyclass]\n#[pyclass(name = \"Scope\", module = \"coglet._sdk\")]\npub struct Scope {\n metrics_recorder: Py,\n /// Per-prediction context from the request body (`dict[str, str]`).\n context: Py,\n}\n\nimpl Scope {\n pub fn new(\n py: Python<'_>,\n sender: Arc,\n context: HashMap,\n ) -> PyResult {\n let recorder = Py::new(py, MetricRecorder::new(sender))?;\n let dict = PyDict::new(py);\n for (k, v) in &context {\n dict.set_item(k, v)?;\n }\n Ok(Self {\n metrics_recorder: recorder,\n context: dict.unbind(),\n })\n }\n\n pub fn noop(py: Python<'_>) -> PyResult {\n let recorder = Py::new(py, MetricRecorder::noop())?;\n let dict = PyDict::new(py);\n Ok(Self {\n metrics_recorder: recorder,\n context: dict.unbind(),\n })\n }\n}\n\n#[gen_stub_pymethods]\n#[pymethods]\nimpl Scope {\n /// The metric recorder for this prediction.\n #[getter]\n fn metrics(&self, py: Python<'_>) -> Py {\n self.metrics_recorder.clone_ref(py)\n }\n\n /// Per-prediction context passed in the request body.\n ///\n /// Returns a `dict[str, str]` (empty dict if no context was provided).\n #[getter]\n fn context(&self, py: Python<'_>) -> Py {\n self.context.clone_ref(py)\n }\n\n /// Convenience: record a metric value.\n ///\n /// Equivalent to `scope.metrics.record(key, value, mode)`.\n #[pyo3(signature = (key, value, mode=None))]\n fn record_metric(\n &self,\n py: Python<'_>,\n key: &str,\n value: &Bound<'_, PyAny>,\n mode: Option<&str>,\n ) -> PyResult<()> {\n self.metrics_recorder\n .borrow(py)\n .record(py, key, value, mode)\n }\n\n fn __repr__(&self, py: Python<'_>) -> String {\n let recorder = self.metrics_recorder.borrow(py);\n format!(\"Scope({})\", recorder.__repr__())\n }\n}\n\n// ============================================================================\n// Shared implementation\n// ============================================================================\n\nfn parse_mode(mode: Option<&str>) -> PyResult {\n match mode {\n None | Some(\"replace\") => Ok(MetricMode::Replace),\n Some(\"incr\") | Some(\"increment\") => Ok(MetricMode::Increment),\n Some(\"append\") => Ok(MetricMode::Append),\n Some(other) => Err(PyTypeError::new_err(format!(\n \"Invalid metric mode: '{}'. Expected 'replace', 'incr', or 'append'.\",\n other\n ))),\n }\n}\n\nfn record_impl(\n _py: Python<'_>,\n inner: &mut RecorderInner,\n key: &str,\n value: &Bound<'_, PyAny>,\n mode: MetricMode,\n) -> PyResult<()> {\n let value_type = MetricValueType::from_py(value)?;\n\n // Type invariant check\n if let Some(existing_type) = inner.types.get(key)\n && *existing_type != value_type\n {\n return Err(PyTypeError::new_err(format!(\n \"Metric '{}' has type {}, cannot set to {} without deleting first\",\n key,\n existing_type.as_str(),\n value_type.as_str(),\n )));\n }\n\n // Mode-specific validation\n if mode == MetricMode::Increment\n && !matches!(value_type, MetricValueType::Int | MetricValueType::Float)\n {\n return Err(PyTypeError::new_err(format!(\n \"Increment mode requires int or float, got {}\",\n value_type.as_str()\n )));\n }\n\n let json_value = py_to_json(value)?;\n\n inner.types.insert(key.to_string(), value_type);\n\n inner\n .sender\n .send_metric(key.to_string(), json_value, mode)\n .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!(\"Failed to send metric: {}\", e)))\n}\n\nfn delete_impl(inner: &mut RecorderInner, key: &str) -> PyResult<()> {\n inner.types.remove(key);\n inner\n .sender\n .send_metric(\n key.to_string(),\n serde_json::Value::Null,\n MetricMode::Replace,\n )\n .map_err(|e| {\n pyo3::exceptions::PyIOError::new_err(format!(\"Failed to send metric delete: {}\", e))\n })\n}\n\n// ============================================================================\n// ContextVar-based routing (same pattern as log_writer.rs)\n// ============================================================================\n\n/// Global ContextVar for the current Scope.\nstatic SCOPE_CONTEXTVAR: OnceLock> = OnceLock::new();\n\n/// Current sync scope (for sync predictions where ContextVar doesn't work across attach calls).\nstatic SYNC_SCOPE: OnceLock>>> = OnceLock::new();\n\nfn get_sync_scope_slot() -> &'static Mutex>> {\n SYNC_SCOPE.get_or_init(|| Mutex::new(None))\n}\n\nfn get_scope_contextvar(py: Python<'_>) -> PyResult<&'static Py> {\n if let Some(cv) = SCOPE_CONTEXTVAR.get() {\n return Ok(cv);\n }\n\n let contextvars = py.import(\"contextvars\")?;\n let cv = contextvars.call_method1(\"ContextVar\", (\"_coglet_metric_scope\",))?;\n\n match SCOPE_CONTEXTVAR.set(cv.unbind()) {\n Ok(()) => {}\n Err(_already_set) => {}\n }\n\n SCOPE_CONTEXTVAR.get().ok_or_else(|| {\n pyo3::exceptions::PyRuntimeError::new_err(\"Failed to initialize scope ContextVar\")\n })\n}\n\n/// Set the current scope in the ContextVar (for async predictions).\npub fn set_current_scope(py: Python<'_>, scope: &Py) -> PyResult> {\n let cv = get_scope_contextvar(py)?;\n let token = cv.call_method1(py, \"set\", (scope,))?;\n Ok(token)\n}\n\n/// Set the current sync scope (for sync predictions).\npub fn set_sync_scope(py: Python<'_>, scope: Option<&Py>) {\n let mut slot = get_sync_scope_slot()\n .lock()\n .expect(\"sync_scope mutex poisoned\");\n *slot = scope.map(|s| s.clone_ref(py));\n}\n\n/// Clear the sync scope.\npub fn clear_sync_scope() {\n let mut slot = get_sync_scope_slot()\n .lock()\n .expect(\"sync_scope mutex poisoned\");\n *slot = None;\n}\n\n/// Python-callable: get the current Scope.\n///\n/// Returns the active scope if inside a prediction, or a no-op scope otherwise.\n#[gen_stub_pyfunction(module = \"coglet._sdk\")]\n#[pyfunction]\n#[pyo3(name = \"current_scope\")]\npub fn py_current_scope(py: Python<'_>) -> PyResult> {\n // Try sync scope first\n {\n let slot = get_sync_scope_slot()\n .lock()\n .expect(\"sync_scope mutex poisoned\");\n if let Some(ref scope) = *slot {\n return Ok(scope.clone_ref(py));\n }\n }\n\n // Try ContextVar (async predictions)\n if let Some(cv) = SCOPE_CONTEXTVAR.get() {\n match cv.call_method0(py, \"get\") {\n Ok(val) => {\n let scope: Py = val.extract(py)?;\n return Ok(scope);\n }\n Err(e) if e.is_instance_of::(py) => {\n // Not set — fall through to no-op\n }\n Err(e) => return Err(e),\n }\n }\n\n // Outside prediction context — return no-op scope\n Py::new(py, Scope::noop(py)?)\n}\n\n// ============================================================================\n// RAII guard for prediction scope lifecycle\n// ============================================================================\n\n/// RAII guard that manages the Scope for a prediction.\n///\n/// On creation, creates a Scope with a MetricRecorder and sets it in\n/// ContextVar + sync scope. On drop, clears the scope and releases the\n/// Arc so the log-forwarder channel can close.\npub struct ScopeGuard {\n scope: Py,\n #[allow(dead_code)]\n token: Py,\n}\n\nimpl ScopeGuard {\n /// Enter scope for a prediction.\n pub fn enter(\n py: Python<'_>,\n sender: Arc,\n context: HashMap,\n ) -> PyResult {\n let scope = Py::new(py, Scope::new(py, sender, context)?)?;\n\n let token = set_current_scope(py, &scope)?;\n set_sync_scope(py, Some(&scope));\n\n Ok(Self { scope, token })\n }\n}\n\nimpl Drop for ScopeGuard {\n fn drop(&mut self) {\n clear_sync_scope();\n\n // Acquire the GIL to release the Arc held by the MetricRecorder.\n // Without this, the Py destructor may not run immediately (PyO3\n // defers ref-count decrements when the GIL is not held), keeping the\n // SlotSender channel alive and blocking the log-forwarder shutdown.\n Python::attach(|py| {\n let scope = self.scope.borrow(py);\n let recorder = scope.metrics_recorder.borrow(py);\n let mut guard = recorder\n .inner\n .lock()\n .expect(\"metric_recorder mutex poisoned\");\n // Drop the RecorderInner (and its Arc)\n *guard = None;\n });\n }\n}\n\n// ============================================================================\n// Python → JSON conversion\n// ============================================================================\n\nfn py_to_json(obj: &Bound<'_, PyAny>) -> PyResult {\n if obj.is_none() {\n Ok(serde_json::Value::Null)\n } else if obj.is_instance_of::() {\n Ok(serde_json::Value::Bool(obj.extract::()?))\n } else if obj.is_instance_of::() {\n if let Ok(v) = obj.extract::() {\n Ok(serde_json::json!(v))\n } else {\n Ok(serde_json::json!(obj.extract::()?))\n }\n } else if obj.is_instance_of::() {\n Ok(serde_json::json!(obj.extract::()?))\n } else if obj.is_instance_of::() {\n Ok(serde_json::Value::String(obj.extract::()?))\n } else if obj.is_instance_of::() {\n let list = obj.cast::()?;\n let items: Vec = list\n .iter()\n .map(|item| py_to_json(&item))\n .collect::>()?;\n Ok(serde_json::Value::Array(items))\n } else if obj.is_instance_of::() {\n let dict = obj.cast::()?;\n let mut map = serde_json::Map::new();\n for (k, v) in dict.iter() {\n let key: String = k.extract()?;\n map.insert(key, py_to_json(&v)?);\n }\n Ok(serde_json::Value::Object(map))\n } else {\n let type_name = obj.get_type().name()?.to_string();\n Err(PyTypeError::new_err(format!(\n \"Cannot convert {} to JSON metric value\",\n type_name\n )))\n }\n}\n" }, { "path": "crates/coglet-python/src/output.rs", "content": "//! Output processing for prediction results.\n//!\n//! Converts Python prediction output to JSON-serializable format:\n//! - Pydantic models -> dict (via model_dump() / dict())\n//! - Dataclasses -> dict (via dataclasses.asdict())\n//! - Enums -> .value\n//! - datetime -> .isoformat()\n//! - PathLike -> base64 data URL\n//! - IOBase -> base64 data URL\n//! - numpy int/float/ndarray -> Python int/float/list\n//! - dict/list/set/tuple/generator -> recursive descent\n//!\n//! This replaces the Python modules cog.json and cog.files.\n\nuse pyo3::prelude::*;\nuse pyo3::types::{PyDict, PyFrozenSet, PyList, PySet, PyString, PyTuple};\n\n/// Process prediction output for JSON serialization.\n///\n/// Calls make_encodeable() to normalize, then encode_files() to convert any\n/// remaining Path/IOBase objects to base64 data URLs.\npub fn process_output<'py>(\n py: Python<'py>,\n output: &Bound<'py, PyAny>,\n) -> PyResult> {\n let encodeable = make_encodeable(py, output)?;\n encode_files(py, &encodeable)\n}\n\n/// Process a single output item (for generator outputs).\npub fn process_output_item<'py>(\n py: Python<'py>,\n item: &Bound<'py, PyAny>,\n) -> PyResult> {\n process_output(py, item)\n}\n\n/// Normalize a Python object into a JSON-friendly form.\n///\n/// Handles Pydantic models, dataclasses, enums, datetime, numpy types,\n/// and collections. PathLike objects are passed through (handled later\n/// by encode_files).\nfn make_encodeable<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult> {\n // Pydantic v2: model_dump()\n if let Ok(method) = obj.getattr(\"model_dump\")\n && method.is_callable()\n {\n let dumped = method.call0()?;\n return make_encodeable(py, &dumped);\n }\n\n // Pydantic v1: dict()\n // Skip plain dicts -- they also have .dict but we handle those below\n if !obj.is_instance_of::()\n && let Ok(method) = obj.getattr(\"dict\")\n && method.is_callable()\n {\n let dumped = method.call0()?;\n return make_encodeable(py, &dumped);\n }\n\n // dataclass instances (not the class itself)\n let dataclasses = py.import(\"dataclasses\")?;\n let is_dataclass = dataclasses.getattr(\"is_dataclass\")?;\n if is_dataclass.call1((obj,))?.is_truthy()?\n && !obj.is_instance(py.get_type::().as_any())?\n {\n let asdict = dataclasses.getattr(\"asdict\")?;\n let d = asdict.call1((obj,))?;\n return make_encodeable(py, &d);\n }\n\n // dict\n if let Ok(dict) = obj.cast_exact::() {\n let new_dict = PyDict::new(py);\n for (key, value) in dict.iter() {\n new_dict.set_item(&key, make_encodeable(py, &value)?)?;\n }\n return Ok(new_dict.into_any());\n }\n\n // list, set, frozenset, tuple, generator\n if obj.is_instance_of::()\n || obj.is_instance_of::()\n || obj.is_instance_of::()\n || obj.is_instance_of::()\n || is_generator(py, obj)?\n {\n let iter = obj.try_iter()?;\n let items: Vec> = iter\n .map(|item| make_encodeable(py, &item?))\n .collect::>()?;\n let list = PyList::new(py, &items)?;\n return Ok(list.into_any());\n }\n\n // Enum -> .value\n let enum_mod = py.import(\"enum\")?;\n let enum_cls = enum_mod.getattr(\"Enum\")?;\n if obj.is_instance(&enum_cls)? {\n return obj.getattr(\"value\");\n }\n\n // datetime -> .isoformat()\n let datetime_mod = py.import(\"datetime\")?;\n let datetime_cls = datetime_mod.getattr(\"datetime\")?;\n if obj.is_instance(&datetime_cls)? {\n return obj.call_method0(\"isoformat\");\n }\n\n // os.PathLike -> pathlib.Path (will be encoded to base64 later by encode_files)\n let os_mod = py.import(\"os\")?;\n let pathlike_cls = os_mod.getattr(\"PathLike\")?;\n if obj.is_instance(&pathlike_cls)? {\n let pathlib = py.import(\"pathlib\")?;\n let path_cls = pathlib.getattr(\"Path\")?;\n return path_cls.call1((obj,));\n }\n\n // numpy types (optional)\n if let Ok(np) = py.import(\"numpy\")\n && !obj.is_instance(py.get_type::().as_any())?\n {\n let np_integer = np.getattr(\"integer\")?;\n if obj.is_instance(&np_integer)? {\n let builtins = py.import(\"builtins\")?;\n return builtins.getattr(\"int\")?.call1((obj,));\n }\n let np_floating = np.getattr(\"floating\")?;\n if obj.is_instance(&np_floating)? {\n let builtins = py.import(\"builtins\")?;\n return builtins.getattr(\"float\")?.call1((obj,));\n }\n let np_ndarray = np.getattr(\"ndarray\")?;\n if obj.is_instance(&np_ndarray)? {\n return obj.call_method0(\"tolist\");\n }\n }\n\n // Primitive / unknown -- pass through\n Ok(obj.clone())\n}\n\n/// Recursively walk the output and encode any Path/IOBase objects to base64 data URLs.\nfn encode_files<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult> {\n // str -- return as-is (don't recurse into characters)\n if obj.is_instance_of::() {\n return Ok(obj.clone());\n }\n\n // dict\n if let Ok(dict) = obj.cast_exact::() {\n let new_dict = PyDict::new(py);\n for (key, value) in dict.iter() {\n new_dict.set_item(&key, encode_files(py, &value)?)?;\n }\n return Ok(new_dict.into_any());\n }\n\n // list\n if let Ok(list) = obj.cast_exact::() {\n let items: Vec> = list\n .iter()\n .map(|item| encode_files(py, &item))\n .collect::>()?;\n let new_list = PyList::new(py, &items)?;\n return Ok(new_list.into_any());\n }\n\n // os.PathLike -> open and base64 encode\n let os_mod = py.import(\"os\")?;\n let pathlike_cls = os_mod.getattr(\"PathLike\")?;\n if obj.is_instance(&pathlike_cls)? {\n let builtins = py.import(\"builtins\")?;\n let fh = builtins.getattr(\"open\")?.call1((obj, \"rb\"))?;\n let result = file_to_base64(py, &fh);\n fh.call_method0(\"close\")?;\n return result;\n }\n\n // io.IOBase -> base64 encode\n let io_mod = py.import(\"io\")?;\n let iobase_cls = io_mod.getattr(\"IOBase\")?;\n if obj.is_instance(&iobase_cls)? {\n return file_to_base64(py, obj);\n }\n\n // Primitive -- pass through\n Ok(obj.clone())\n}\n\n/// Encode a file handle to a base64 data URL.\n///\n/// Seeks to start if seekable, reads all bytes, guesses MIME type from\n/// the file name, and returns \"data:{mime};base64,{encoded}\".\nfn file_to_base64<'py>(py: Python<'py>, fh: &Bound<'py, PyAny>) -> PyResult> {\n // Seek to start if possible\n if let Ok(seekable) = fh.call_method0(\"seekable\")\n && seekable.is_truthy()?\n {\n fh.call_method1(\"seek\", (0,))?;\n }\n\n // Read content\n let content = fh.call_method0(\"read\")?;\n let bytes: Vec = if content.is_instance_of::() {\n let s: String = content.extract()?;\n s.into_bytes()\n } else {\n content.extract()?\n };\n\n // Guess MIME type from filename\n let mime_type = if let Ok(name) = fh.getattr(\"name\")\n && !name.is_none()\n {\n let name_str: String = name.extract()?;\n let mimetypes = py.import(\"mimetypes\")?;\n let guess = mimetypes.call_method1(\"guess_type\", (&name_str,))?;\n let first = guess.get_item(0)?;\n if first.is_none() {\n \"application/octet-stream\".to_string()\n } else {\n first.extract()?\n }\n } else {\n \"application/octet-stream\".to_string()\n };\n\n // Base64 encode\n use base64::Engine as _;\n let encoded = base64::engine::general_purpose::STANDARD.encode(&bytes);\n let data_url = format!(\"data:{mime_type};base64,{encoded}\");\n\n Ok(PyString::new(py, &data_url).into_any())\n}\n\n/// Check if a Python object is a generator instance.\nfn is_generator<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult {\n let types_mod = py.import(\"types\")?;\n let gen_type = types_mod.getattr(\"GeneratorType\")?;\n obj.is_instance(&gen_type)\n}\n\n#[cfg(test)]\nmod tests {\n use super::*;\n use base64::Engine as _;\n use pyo3::types::PyDict;\n\n /// Helper: evaluate a Python expression and run make_encodeable on it.\n fn encodeable(py_expr: &str) -> String {\n pyo3::Python::initialize();\n Python::attach(|py| {\n let locals = PyDict::new(py);\n // Run any setup + expression, storing result in `obj`\n let code = format!(\"import json\\n{}\\nresult = obj\", py_expr);\n py.run(&std::ffi::CString::new(code).unwrap(), None, Some(&locals))\n .expect(\"failed to evaluate test expression\");\n\n let obj = locals.get_item(\"result\").unwrap().unwrap();\n let encoded = make_encodeable(py, &obj).expect(\"make_encodeable failed\");\n\n // Convert to JSON string for easy assertion\n let json_mod = py.import(\"json\").unwrap();\n let json_str = json_mod\n .call_method1(\"dumps\", (&encoded,))\n .expect(\"json.dumps failed\");\n json_str.extract::().unwrap()\n })\n }\n\n /// Helper: evaluate a Python expression and run process_output on it.\n fn processed(py_expr: &str) -> String {\n pyo3::Python::initialize();\n Python::attach(|py| {\n let locals = PyDict::new(py);\n let code = format!(\"import json\\n{}\\nresult = obj\", py_expr);\n py.run(&std::ffi::CString::new(code).unwrap(), None, Some(&locals))\n .expect(\"failed to evaluate test expression\");\n\n let obj = locals.get_item(\"result\").unwrap().unwrap();\n let output = process_output(py, &obj).expect(\"process_output failed\");\n\n let json_mod = py.import(\"json\").unwrap();\n let json_str = json_mod\n .call_method1(\"dumps\", (&output,))\n .expect(\"json.dumps failed\");\n json_str.extract::().unwrap()\n })\n }\n\n // ── make_encodeable: primitives ──────────────────────────────────\n\n #[test]\n fn encodeable_string() {\n assert_eq!(encodeable(\"obj = 'hello'\"), r#\"\"hello\"\"#);\n }\n\n #[test]\n fn encodeable_int() {\n assert_eq!(encodeable(\"obj = 42\"), \"42\");\n }\n\n #[test]\n fn encodeable_float() {\n assert_eq!(encodeable(\"obj = 3.14\"), \"3.14\");\n }\n\n #[test]\n fn encodeable_bool() {\n assert_eq!(encodeable(\"obj = True\"), \"true\");\n }\n\n #[test]\n fn encodeable_none() {\n assert_eq!(encodeable(\"obj = None\"), \"null\");\n }\n\n // ── make_encodeable: collections ─────────────────────────────────\n\n #[test]\n fn encodeable_list() {\n assert_eq!(encodeable(\"obj = [1, 2, 3]\"), \"[1, 2, 3]\");\n }\n\n #[test]\n fn encodeable_dict() {\n assert_eq!(\n encodeable(r#\"obj = {\"a\": 1, \"b\": 2}\"#),\n r#\"{\"a\": 1, \"b\": 2}\"#\n );\n }\n\n #[test]\n fn encodeable_tuple_to_list() {\n assert_eq!(encodeable(\"obj = (1, 2, 3)\"), \"[1, 2, 3]\");\n }\n\n #[test]\n fn encodeable_set_to_list() {\n // Set with single element to avoid ordering issues\n assert_eq!(encodeable(\"obj = {42}\"), \"[42]\");\n }\n\n #[test]\n fn encodeable_frozenset_to_list() {\n assert_eq!(encodeable(\"obj = frozenset([99])\"), \"[99]\");\n }\n\n #[test]\n fn encodeable_nested_dict() {\n assert_eq!(\n encodeable(r#\"obj = {\"outer\": {\"inner\": [1, 2]}}\"#),\n r#\"{\"outer\": {\"inner\": [1, 2]}}\"#\n );\n }\n\n // ── make_encodeable: enum ────────────────────────────────────────\n\n #[test]\n fn encodeable_enum() {\n assert_eq!(\n encodeable(\"import enum\\nclass Color(enum.Enum):\\n RED = 'red'\\nobj = Color.RED\"),\n r#\"\"red\"\"#\n );\n }\n\n #[test]\n fn encodeable_int_enum() {\n assert_eq!(\n encodeable(\n \"import enum\\nclass Priority(enum.IntEnum):\\n HIGH = 1\\nobj = Priority.HIGH\"\n ),\n \"1\"\n );\n }\n\n // ── make_encodeable: datetime ────────────────────────────────────\n\n #[test]\n fn encodeable_datetime() {\n let result =\n encodeable(\"from datetime import datetime\\nobj = datetime(2025, 1, 15, 10, 30, 0)\");\n assert_eq!(result, r#\"\"2025-01-15T10:30:00\"\"#);\n }\n\n // ── make_encodeable: dataclass ───────────────────────────────────\n\n #[test]\n fn encodeable_dataclass() {\n assert_eq!(\n encodeable(\n \"from dataclasses import dataclass\\n\\\n @dataclass\\n\\\n class Point:\\n\\\n \\tx: int\\n\\\n \\ty: int\\n\\\n obj = Point(x=1, y=2)\"\n ),\n r#\"{\"x\": 1, \"y\": 2}\"#\n );\n }\n\n #[test]\n fn encodeable_nested_dataclass() {\n assert_eq!(\n encodeable(\n \"from dataclasses import dataclass, asdict\\n\\\n @dataclass\\n\\\n class Inner:\\n\\\n \\tval: str\\n\\\n # Build nested via dict so class scoping isn't an issue\\n\\\n obj = {'inner': asdict(Inner(val='hello')), 'name': 'test'}\"\n ),\n r#\"{\"inner\": {\"val\": \"hello\"}, \"name\": \"test\"}\"#\n );\n }\n\n // ── make_encodeable: generator ───────────────────────────────────\n\n #[test]\n fn encodeable_generator() {\n assert_eq!(encodeable(\"obj = (x * 2 for x in range(3))\"), \"[0, 2, 4]\");\n }\n\n // ── make_encodeable: enum value in collection ────────────────────\n\n #[test]\n fn encodeable_enum_in_list() {\n assert_eq!(\n encodeable(\n \"import enum\\n\\\n class Status(enum.Enum):\\n\\\n \\tOK = 'ok'\\n\\\n \\tERR = 'err'\\n\\\n obj = [Status.OK, Status.ERR]\"\n ),\n r#\"[\"ok\", \"err\"]\"#\n );\n }\n\n // ── encode_files / file_to_base64 ────────────────────────────────\n\n #[test]\n fn encode_pathlike_to_base64() {\n let result = processed(\n \"import tempfile, pathlib\\n\\\n f = tempfile.NamedTemporaryFile(suffix='.txt', delete=False)\\n\\\n f.write(b'hello world')\\n\\\n f.close()\\n\\\n obj = pathlib.Path(f.name)\",\n );\n assert!(\n result.starts_with(r#\"\"data:text/plain;base64,\"#),\n \"expected data URL, got: {result}\"\n );\n // Verify the base64 content decodes correctly\n let b64_part = result.trim_matches('\"').split(\",\").nth(1).unwrap();\n let decoded = base64::engine::general_purpose::STANDARD\n .decode(b64_part)\n .unwrap();\n assert_eq!(decoded, b\"hello world\");\n }\n\n #[test]\n fn encode_iobase_to_base64() {\n let result = processed(\n \"import io\\n\\\n obj = io.BytesIO(b'test bytes')\",\n );\n assert!(\n result.starts_with(r#\"\"data:application/octet-stream;base64,\"#),\n \"expected data URL, got: {result}\"\n );\n let b64_part = result.trim_matches('\"').split(\",\").nth(1).unwrap();\n let decoded = base64::engine::general_purpose::STANDARD\n .decode(b64_part)\n .unwrap();\n assert_eq!(decoded, b\"test bytes\");\n }\n\n #[test]\n fn encode_iobase_seeks_to_start() {\n let result = processed(\n \"import io\\n\\\n buf = io.BytesIO(b'rewind me')\\n\\\n buf.read() # advance to end\\n\\\n obj = buf\",\n );\n let b64_part = result.trim_matches('\"').split(\",\").nth(1).unwrap();\n let decoded = base64::engine::general_purpose::STANDARD\n .decode(b64_part)\n .unwrap();\n assert_eq!(decoded, b\"rewind me\", \"should seek to start before reading\");\n }\n\n #[test]\n fn encode_file_in_dict() {\n let result = processed(\n \"import io\\n\\\n obj = {'output': io.BytesIO(b'nested')}\",\n );\n // Parse the JSON to verify structure\n let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();\n let url = parsed[\"output\"].as_str().unwrap();\n assert!(\n url.starts_with(\"data:application/octet-stream;base64,\"),\n \"expected data URL in dict value\"\n );\n }\n\n #[test]\n fn encode_file_in_list() {\n let result = processed(\n \"import io\\n\\\n obj = [io.BytesIO(b'item1'), io.BytesIO(b'item2')]\",\n );\n let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();\n assert!(parsed.as_array().unwrap().len() == 2);\n for item in parsed.as_array().unwrap() {\n assert!(item.as_str().unwrap().starts_with(\"data:\"));\n }\n }\n\n #[test]\n fn encode_string_passthrough() {\n // Strings should NOT be recursed into\n assert_eq!(processed(\"obj = 'just a string'\"), r#\"\"just a string\"\"#);\n }\n\n #[test]\n fn encode_mime_type_guessing() {\n let result = processed(\n \"import tempfile, pathlib\\n\\\n f = tempfile.NamedTemporaryFile(suffix='.png', delete=False)\\n\\\n f.write(b'\\\\x89PNG')\\n\\\n f.close()\\n\\\n obj = pathlib.Path(f.name)\",\n );\n assert!(\n result.contains(\"image/png\"),\n \"expected image/png MIME type, got: {result}\"\n );\n }\n\n // ── process_output: end-to-end ───────────────────────────────────\n\n #[test]\n fn process_output_primitives_passthrough() {\n assert_eq!(processed(\"obj = 'hello'\"), r#\"\"hello\"\"#);\n assert_eq!(processed(\"obj = 42\"), \"42\");\n assert_eq!(processed(\"obj = None\"), \"null\");\n }\n\n #[test]\n fn process_output_dataclass_with_file() {\n let result = processed(\n \"from dataclasses import dataclass\\n\\\n import pathlib, tempfile\\n\\\n @dataclass\\n\\\n class Output:\\n\\\n \\ttext: str\\n\\\n \\tdata: object\\n\\\n f = tempfile.NamedTemporaryFile(suffix='.bin', delete=False)\\n\\\n f.write(b'binary')\\n\\\n f.close()\\n\\\n obj = Output(text='result', data=pathlib.Path(f.name))\",\n );\n let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();\n assert_eq!(parsed[\"text\"], \"result\");\n assert!(parsed[\"data\"].as_str().unwrap().starts_with(\"data:\"));\n }\n}\n" }, { "path": "crates/coglet-python/src/predictor.rs", "content": "//! Python predictor loading and invocation.\n\nuse std::sync::{Arc, OnceLock};\n\nuse pyo3::prelude::*;\nuse pyo3::types::PyDict;\n\nuse coglet_core::worker::SlotSender;\nuse coglet_core::{PredictionError, PredictionOutput, PredictionResult};\n\nuse crate::cancel;\nuse crate::input::{self, PreparedInput};\nuse crate::output;\n\n// =============================================================================\n// Async helper functions — defined as Python strings, initialized once.\n//\n// These must be Python `async def` functions to participate in asyncio's event\n// loop. They cannot be expressed as pure Rust because they use Python's async\n// iteration protocol and ContextVar.set() before awaiting a coroutine.\n// =============================================================================\n\n/// Collects an async generator into a list. Initialized once, reused per-call.\nstatic COLLECT_ASYNC_GEN: OnceLock> = OnceLock::new();\n\n/// Sets a ContextVar then awaits a coroutine. Initialized once, reused per-call.\nstatic CTX_WRAPPER: OnceLock> = OnceLock::new();\n\n/// Get or initialize the `_collect_async_gen` Python helper.\nfn get_collect_async_gen(py: Python<'_>) -> Result, PredictionError> {\n if let Some(f) = COLLECT_ASYNC_GEN.get() {\n return Ok(f.clone_ref(py));\n }\n\n let code = c\"\\\nasync def _collect_async_gen(agen):\n results = []\n async for item in agen:\n results.append(item)\n return results\n\";\n let globals = PyDict::new(py);\n py.run(code, Some(&globals), None).map_err(|e| {\n PredictionError::Failed(format!(\"Failed to define _collect_async_gen: {e}\"))\n })?;\n let f = globals\n .get_item(\"_collect_async_gen\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to get _collect_async_gen: {e}\")))?\n .ok_or_else(|| PredictionError::Failed(\"_collect_async_gen not found\".to_string()))?\n .unbind();\n let _ = COLLECT_ASYNC_GEN.set(f.clone_ref(py));\n Ok(f)\n}\n\n/// Get or initialize the `_ctx_wrapper` Python helper.\nfn get_ctx_wrapper(py: Python<'_>) -> Result, PredictionError> {\n if let Some(f) = CTX_WRAPPER.get() {\n return Ok(f.clone_ref(py));\n }\n\n let code = c\"\\\nasync def _ctx_wrapper(coro, prediction_id, contextvar):\n contextvar.set(prediction_id)\n return await coro\n\";\n let globals = PyDict::new(py);\n py.run(code, Some(&globals), None)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to define _ctx_wrapper: {e}\")))?;\n let f = globals\n .get_item(\"_ctx_wrapper\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to get _ctx_wrapper: {e}\")))?\n .ok_or_else(|| PredictionError::Failed(\"_ctx_wrapper not found\".to_string()))?\n .unbind();\n let _ = CTX_WRAPPER.set(f.clone_ref(py));\n Ok(f)\n}\n\n/// Check if a PyErr is a CancelationException or asyncio.CancelledError.\nfn is_cancelation_exception(py: Python<'_>, err: &PyErr) -> bool {\n // Check for our static CancelationException type\n if err.is_instance_of::(py) {\n return true;\n }\n\n // Check for asyncio.CancelledError\n if let Ok(asyncio) = py.import(\"asyncio\")\n && let Ok(cancelled_error) = asyncio.getattr(\"CancelledError\")\n && err.is_instance(py, &cancelled_error)\n {\n return true;\n }\n\n false\n}\n\n/// Format a Python validation error.\n///\n/// Cog validation errors are already formatted as \"field: message\".\nfn format_validation_error(py: Python<'_>, err: &PyErr) -> String {\n err.value(py).to_string()\n}\n\n/// Send a single output item over IPC, routing file outputs to disk.\n///\n/// For Path outputs (os.PathLike): sends the existing file path via send_file_output.\n/// For IOBase outputs: reads bytes, writes to output_dir via write_file_output.\n/// For everything else: processes through make_encodeable + upload_files, then send_output.\nfn send_output_item(\n py: Python<'_>,\n item: &Bound<'_, PyAny>,\n json_module: &Bound<'_, PyAny>,\n slot_sender: &SlotSender,\n) -> Result<(), PredictionError> {\n let os = py\n .import(\"os\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to import os: {}\", e)))?;\n let io_mod = py\n .import(\"io\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to import io: {}\", e)))?;\n let pathlike = os\n .getattr(\"PathLike\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to get os.PathLike: {}\", e)))?;\n let iobase = io_mod\n .getattr(\"IOBase\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to get io.IOBase: {}\", e)))?;\n\n if item.is_instance(&pathlike).unwrap_or(false) {\n // Path output — file already on disk, send path reference\n let path_str: String = item\n .call_method0(\"__fspath__\")\n .and_then(|p| p.extract())\n .map_err(|e| PredictionError::Failed(format!(\"Failed to get fspath: {}\", e)))?;\n slot_sender\n .send_file_output(std::path::PathBuf::from(path_str), None)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to send file output: {}\", e)))?;\n return Ok(());\n }\n\n if item.is_instance(&iobase).unwrap_or(false) {\n // IOBase output — read bytes, write to disk via SlotSender\n // Seek to start if seekable\n if item\n .call_method0(\"seekable\")\n .and_then(|r| r.extract::())\n .unwrap_or(false)\n {\n let _ = item.call_method1(\"seek\", (0,));\n }\n let data: Vec = item\n .call_method0(\"read\")\n .and_then(|d| d.extract())\n .map_err(|e| PredictionError::Failed(format!(\"Failed to read IOBase: {}\", e)))?;\n\n // Try to guess extension from filename\n let ext = item\n .getattr(\"name\")\n .and_then(|n| n.extract::())\n .ok()\n .and_then(|name| {\n std::path::Path::new(&name)\n .extension()\n .and_then(|e| e.to_str())\n .map(|s| s.to_string())\n })\n .unwrap_or_else(|| \"bin\".to_string());\n\n slot_sender\n .write_file_output(&data, &ext, None)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to write file output: {}\", e)))?;\n return Ok(());\n }\n\n // Non-file output - process normally\n let processed = output::process_output_item(py, item)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to process output item: {}\", e)))?;\n\n let item_str: String = json_module\n .call_method1(\"dumps\", (&processed,))\n .map_err(|e| PredictionError::Failed(format!(\"Failed to serialize output item: {}\", e)))?\n .extract()\n .map_err(|e| PredictionError::Failed(format!(\"Failed to extract output string: {}\", e)))?;\n\n let item_json: serde_json::Value = serde_json::from_str(&item_str)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to parse output JSON: {}\", e)))?;\n\n slot_sender\n .send_output(item_json)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to send output: {}\", e)))?;\n\n Ok(())\n}\n\n/// Type alias for Python object (Py).\ntype PyObject = Py;\n\n/// How a predict() method executes\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum PredictKind {\n /// Synchronous function: def predict(self, **input) -> Output\n Sync,\n /// Async coroutine: async def predict(self, **input) -> Output\n Async,\n /// Async generator: async def predict(self, **input) -> AsyncIterator[Output]\n AsyncGen,\n}\n\n/// Whether and how train() exists\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum TrainKind {\n /// No train() method\n None,\n /// Synchronous: def train(self, **input) -> Output\n Sync,\n /// Async: async def train(self, **input) -> Output\n Async,\n}\n\n/// The predictor's structure and invocation target\n#[derive(Debug, Clone, PartialEq, Eq)]\npub enum PredictorKind {\n /// Class instance with predict() method, optionally train()\n Class {\n predict: PredictKind,\n train: TrainKind,\n },\n /// Standalone function (e.g., train.py:train)\n /// The PredictKind describes how the function executes (sync/async/async_gen)\n StandaloneFunction(PredictKind),\n}\n\n/// A loaded Python predictor instance.\n///\n/// Input coercion (URL->Path/File) and FieldInfo default unwrapping are handled\n/// in Rust. The Python `_adt` and `_inspector` modules are no longer called.\npub struct PythonPredictor {\n instance: PyObject,\n /// The predictor's kind (class or standalone function) and method execution types\n kind: PredictorKind,\n}\n\n// PyObject is Send in PyO3 0.23+\n// Safety: We only access the instance through Python::attach()\nunsafe impl Send for PythonPredictor {}\nunsafe impl Sync for PythonPredictor {}\n\nimpl PythonPredictor {\n /// Load a predictor from a reference like \"predict.py:Predictor\".\n pub fn load(py: Python<'_>, predictor_ref: &str) -> PyResult {\n // Import the cog.predictor module to use its loading function\n let cog_predictor = py.import(\"cog.predictor\")?;\n let load_fn = cog_predictor.getattr(\"load_predictor_from_ref\")?;\n\n // Load the predictor class and instantiate it\n let instance: PyObject = load_fn.call1((predictor_ref,))?.unbind();\n\n // Check if this is a standalone function (train mode) or a Predictor instance\n let inspect = py.import(\"inspect\")?;\n let is_function: bool = inspect\n .call_method1(\"isfunction\", (instance.bind(py),))?\n .extract()?;\n\n let kind = if is_function {\n // Standalone function - detect its async nature\n let (is_async, is_async_gen) = Self::detect_async(py, &instance, \"\")?;\n let predict_kind = if is_async_gen {\n tracing::info!(\"Detected async generator train()\");\n PredictKind::AsyncGen\n } else if is_async {\n tracing::info!(\"Detected async train()\");\n PredictKind::Async\n } else {\n tracing::info!(\"Detected sync train()\");\n PredictKind::Sync\n };\n PredictorKind::StandaloneFunction(predict_kind)\n } else {\n // Class instance - detect predict() and train() methods\n let (is_async, is_async_gen) = Self::detect_async(py, &instance, \"predict\")?;\n let predict_kind = if is_async_gen {\n tracing::info!(\"Detected async generator predict()\");\n PredictKind::AsyncGen\n } else if is_async {\n tracing::info!(\"Detected async predict()\");\n PredictKind::Async\n } else {\n tracing::info!(\"Detected sync predict()\");\n PredictKind::Sync\n };\n\n // Check if train() method exists and if it's async\n let train_kind = if instance.bind(py).hasattr(\"train\")? {\n let (train_async, _) = Self::detect_async(py, &instance, \"train\")?;\n if train_async {\n tracing::info!(\"Detected async train()\");\n TrainKind::Async\n } else {\n tracing::info!(\"Detected sync train()\");\n TrainKind::Sync\n }\n } else {\n TrainKind::None\n };\n\n PredictorKind::Class {\n predict: predict_kind,\n train: train_kind,\n }\n };\n\n let predictor = Self { instance, kind };\n\n // Patch FieldInfo defaults on predict/train methods so Python uses actual\n // default values instead of FieldInfo wrapper objects for missing inputs.\n // Input(default=42, description=\"...\") creates a FieldInfo; without patching,\n // Python would pass the FieldInfo itself as the default value.\n if is_function {\n Self::unwrap_field_info_defaults(py, &predictor.instance, \"\")?;\n } else {\n Self::unwrap_field_info_defaults(py, &predictor.instance, \"predict\")?;\n if matches!(predictor.kind, PredictorKind::Class { train, .. } if train != TrainKind::None)\n {\n Self::unwrap_field_info_defaults(py, &predictor.instance, \"train\")?;\n }\n }\n\n Ok(predictor)\n }\n\n /// Replace FieldInfo defaults with their `.default` values on a method's signature.\n ///\n /// When users write `def predict(self, seed: int = Input(default=42, description=\"...\"))`,\n /// the Python default for `seed` is a `FieldInfo(default=42, ...)` object. If `seed` is\n /// missing from the input dict, Python would use this FieldInfo as the value — not `42`.\n ///\n /// This patches `__defaults__` on the underlying function so Python natively resolves\n /// to the actual default values.\n fn unwrap_field_info_defaults(\n py: Python<'_>,\n instance: &PyObject,\n method_name: &str,\n ) -> PyResult<()> {\n let field_info_class = py.import(\"cog.input\")?.getattr(\"FieldInfo\")?;\n\n // Get the underlying function object\n let func = if method_name.is_empty() {\n // Standalone function\n instance.bind(py).clone()\n } else {\n // Bound method — get __func__ for the raw function\n instance\n .bind(py)\n .getattr(method_name)?\n .getattr(\"__func__\")?\n };\n\n // Patch __defaults__ (positional parameter defaults)\n if let Ok(defaults) = func.getattr(\"__defaults__\")\n && !defaults.is_none()\n {\n let defaults_tuple = defaults.cast::()?;\n let mut new_defaults: Vec> = Vec::new();\n let mut changed = false;\n\n for item in defaults_tuple.iter() {\n if item.is_instance(&field_info_class)? {\n new_defaults.push(item.getattr(\"default\")?);\n changed = true;\n } else {\n new_defaults.push(item);\n }\n }\n\n if changed {\n let new_tuple = pyo3::types::PyTuple::new(py, &new_defaults)?;\n func.setattr(\"__defaults__\", new_tuple)?;\n tracing::debug!(\"Patched FieldInfo defaults on {}\", method_name);\n }\n }\n\n // Patch __kwdefaults__ (keyword-only parameter defaults)\n if let Ok(kwdefaults) = func.getattr(\"__kwdefaults__\")\n && !kwdefaults.is_none()\n {\n let kwdefaults_dict = kwdefaults.cast::()?;\n let mut changed = false;\n\n for (key, value) in kwdefaults_dict.iter() {\n if value.is_instance(&field_info_class)? {\n kwdefaults_dict.set_item(&key, value.getattr(\"default\")?)?;\n changed = true;\n }\n }\n\n if changed {\n tracing::debug!(\"Patched FieldInfo kwdefaults on {}\", method_name);\n }\n }\n\n Ok(())\n }\n\n /// Detect if a method is an async function.\n /// Returns (is_async, is_async_gen) tuple.\n ///\n /// If method_name is empty, checks the instance itself (for standalone functions).\n fn detect_async(\n py: Python<'_>,\n instance: &PyObject,\n method_name: &str,\n ) -> PyResult<(bool, bool)> {\n let inspect = py.import(\"inspect\")?;\n\n // If method_name is empty, check the instance itself (standalone function)\n let target = if method_name.is_empty() {\n instance.bind(py).clone()\n } else {\n instance.bind(py).getattr(method_name)?\n };\n\n // Check isasyncgenfunction first (it's more specific)\n let is_async_gen: bool = inspect\n .call_method1(\"isasyncgenfunction\", (&target,))?\n .extract()?;\n if is_async_gen {\n return Ok((true, true));\n }\n\n // Check iscoroutinefunction\n let is_coro: bool = inspect\n .call_method1(\"iscoroutinefunction\", (&target,))?\n .extract()?;\n Ok((is_coro, false))\n }\n\n /// Returns true if this predictor has an async predict() method.\n pub fn is_async(&self) -> bool {\n match &self.kind {\n PredictorKind::Class { predict, .. } => {\n matches!(predict, PredictKind::Async | PredictKind::AsyncGen)\n }\n PredictorKind::StandaloneFunction(predict_kind) => {\n matches!(predict_kind, PredictKind::Async | PredictKind::AsyncGen)\n }\n }\n }\n\n /// Returns true if this predictor has a train() method.\n pub fn has_train(&self) -> bool {\n match &self.kind {\n PredictorKind::Class { train, .. } => !matches!(train, TrainKind::None),\n PredictorKind::StandaloneFunction(_) => true,\n }\n }\n\n /// Returns true if the train() method is async.\n pub fn is_train_async(&self) -> bool {\n match &self.kind {\n PredictorKind::Class { train, .. } => matches!(train, TrainKind::Async),\n PredictorKind::StandaloneFunction(predict_kind) => {\n matches!(predict_kind, PredictKind::Async | PredictKind::AsyncGen)\n }\n }\n }\n\n /// Call setup() on the predictor, handling weights parameter if present.\n ///\n /// Uses cog.predictor helpers to detect and extract weights:\n /// - `has_setup_weights()` checks if setup() has a weights parameter\n /// - `extract_setup_weights()` reads from COG_WEIGHTS env or ./weights path\n pub fn setup(&self, py: Python<'_>) -> PyResult<()> {\n let instance = self.instance.bind(py);\n\n // Check if setup method exists\n if !instance.hasattr(\"setup\")? {\n return Ok(());\n }\n\n // Import cog.predictor helpers\n let cog_predictor = py.import(\"cog.predictor\")?;\n let has_setup_weights = cog_predictor.getattr(\"has_setup_weights\")?;\n let extract_setup_weights = cog_predictor.getattr(\"extract_setup_weights\")?;\n\n // Check if setup() has a weights parameter\n let needs_weights: bool = has_setup_weights.call1((&instance,))?.extract()?;\n\n if needs_weights {\n // Extract weights from COG_WEIGHTS env or ./weights path\n let weights = extract_setup_weights.call1((&instance,))?;\n instance.call_method1(\"setup\", (weights,))?;\n } else {\n instance.call_method0(\"setup\")?;\n }\n\n Ok(())\n }\n\n /// Get the predict function object for type annotation introspection.\n pub fn predict_func<'py>(&self, py: Python<'py>) -> PyResult> {\n let instance = self.instance.bind(py);\n match &self.kind {\n PredictorKind::Class { .. } => instance.getattr(\"predict\"),\n PredictorKind::StandaloneFunction(_) => Ok(instance.clone()),\n }\n }\n\n /// Get the train function object for type annotation introspection.\n pub fn train_func<'py>(&self, py: Python<'py>) -> PyResult> {\n let instance = self.instance.bind(py);\n match &self.kind {\n PredictorKind::Class { .. } => instance.getattr(\"train\"),\n PredictorKind::StandaloneFunction(_) => Ok(instance.clone()),\n }\n }\n\n /// Call predict() with the given input dict, returning raw Python output.\n ///\n /// For standalone functions, calls the function directly.\n pub fn predict_raw(&self, py: Python<'_>, input: &Bound<'_, PyDict>) -> PyResult {\n let (method_name, is_async) = match &self.kind {\n PredictorKind::Class { predict, .. } => (\n \"predict\",\n matches!(predict, PredictKind::Async | PredictKind::AsyncGen),\n ),\n PredictorKind::StandaloneFunction(predict_kind) => (\n \"\",\n matches!(predict_kind, PredictKind::Async | PredictKind::AsyncGen),\n ),\n };\n self.call_method_raw(py, method_name, is_async, input)\n }\n\n /// Call train() with the given input dict, returning raw Python output.\n ///\n /// For standalone train functions, calls the function directly.\n /// For Predictor classes with a train() method, calls instance.train().\n pub fn train_raw(&self, py: Python<'_>, input: &Bound<'_, PyDict>) -> PyResult {\n let (method_name, is_async) = match &self.kind {\n PredictorKind::Class { train, .. } => (\"train\", matches!(train, TrainKind::Async)),\n PredictorKind::StandaloneFunction(predict_kind) => (\n \"\",\n matches!(predict_kind, PredictKind::Async | PredictKind::AsyncGen),\n ),\n };\n self.call_method_raw(py, method_name, is_async, input)\n }\n\n /// Internal helper to call a method (predict or train) on the predictor.\n fn call_method_raw(\n &self,\n py: Python<'_>,\n method_name: &str,\n is_async: bool,\n input: &Bound<'_, PyDict>,\n ) -> PyResult {\n let instance = self.instance.bind(py);\n\n // Call the method - returns coroutine if async, result if sync\n // If method_name is empty, call the instance directly (standalone function)\n let method_result = if method_name.is_empty() {\n instance.call((), Some(input))?\n } else {\n instance.call_method(method_name, (), Some(input))?\n };\n\n // If async, run the coroutine with asyncio.run()\n let result = if is_async {\n let asyncio = py.import(\"asyncio\")?;\n asyncio.call_method1(\"run\", (&method_result,))?\n } else {\n method_result\n };\n\n Ok(result.unbind())\n }\n\n /// Worker mode predict - with input processing and output serialization.\n pub fn predict_worker(\n &self,\n input: serde_json::Value,\n slot_sender: Arc,\n ) -> Result {\n Python::attach(|py| {\n let json_module = py.import(\"json\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to import json module: {}\", e))\n })?;\n let types_module = py.import(\"types\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to import types module: {}\", e))\n })?;\n let generator_type = types_module.getattr(\"GeneratorType\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get GeneratorType: {}\", e))\n })?;\n\n let input_str = serde_json::to_string(&input)\n .map_err(|e| PredictionError::InvalidInput(e.to_string()))?;\n\n let py_input = json_module\n .call_method1(\"loads\", (input_str,))\n .map_err(|e| PredictionError::InvalidInput(format!(\"Invalid JSON input: {}\", e)))?;\n\n #[allow(deprecated)]\n let raw_input_dict = py_input.downcast::().map_err(|_| {\n PredictionError::InvalidInput(\"Input must be a JSON object\".to_string())\n })?;\n\n // PreparedInput cleans up temp files on drop (RAII)\n let func = self.predict_func(py).map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get predict function: {}\", e))\n })?;\n let prepared = input::prepare_input(py, raw_input_dict, &func)\n .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?;\n let input_dict = prepared.dict(py);\n\n // Call predict\n let result = self.predict_raw(py, &input_dict);\n\n // Handle errors (prepared drops here, cleaning up temp files)\n let result = match result {\n Ok(r) => r,\n Err(e) => {\n drop(prepared); // Explicit cleanup on error path\n if is_cancelation_exception(py, &e) {\n return Err(PredictionError::Cancelled);\n }\n return Err(PredictionError::Failed(format!(\"Prediction failed: {}\", e)));\n }\n };\n\n let result_bound = result.bind(py);\n let is_generator: bool = result_bound.is_instance(&generator_type).unwrap_or(false);\n\n let output = if is_generator {\n self.process_generator_output(py, result_bound, &json_module, &slot_sender)?\n } else {\n self.process_single_output(py, result_bound, &json_module, &slot_sender)?\n };\n\n // prepared drops here, cleaning up temp files via RAII\n drop(prepared);\n\n Ok(PredictionResult {\n output,\n predict_time: None,\n logs: String::new(),\n metrics: Default::default(),\n })\n })\n }\n\n /// Worker mode train - with input processing and output serialization.\n pub fn train_worker(\n &self,\n input: serde_json::Value,\n slot_sender: Arc,\n ) -> Result {\n Python::attach(|py| {\n let json_module = py.import(\"json\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to import json module: {}\", e))\n })?;\n let types_module = py.import(\"types\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to import types module: {}\", e))\n })?;\n let generator_type = types_module.getattr(\"GeneratorType\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get GeneratorType: {}\", e))\n })?;\n\n let input_str = serde_json::to_string(&input)\n .map_err(|e| PredictionError::InvalidInput(e.to_string()))?;\n\n let py_input = json_module\n .call_method1(\"loads\", (input_str,))\n .map_err(|e| PredictionError::InvalidInput(format!(\"Invalid JSON input: {}\", e)))?;\n\n #[allow(deprecated)]\n let raw_input_dict = py_input.downcast::().map_err(|_| {\n PredictionError::InvalidInput(\"Input must be a JSON object\".to_string())\n })?;\n\n // PreparedInput cleans up temp files on drop (RAII)\n let func = self.train_func(py).map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get train function: {}\", e))\n })?;\n let prepared = input::prepare_input(py, raw_input_dict, &func)\n .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?;\n let input_dict = prepared.dict(py);\n\n // Call train\n let result = self.train_raw(py, &input_dict);\n\n // Handle errors\n let result = match result {\n Ok(r) => r,\n Err(e) => {\n drop(prepared);\n if is_cancelation_exception(py, &e) {\n return Err(PredictionError::Cancelled);\n }\n return Err(PredictionError::Failed(format!(\"Training failed: {}\", e)));\n }\n };\n\n let result_bound = result.bind(py);\n let is_generator: bool = result_bound.is_instance(&generator_type).unwrap_or(false);\n\n let output = if is_generator {\n self.process_generator_output(py, result_bound, &json_module, &slot_sender)?\n } else {\n self.process_single_output(py, result_bound, &json_module, &slot_sender)?\n };\n\n drop(prepared);\n\n Ok(PredictionResult {\n output,\n predict_time: None,\n logs: String::new(),\n metrics: Default::default(),\n })\n })\n }\n\n /// Process generator output by streaming each yield over IPC.\n fn process_generator_output(\n &self,\n py: Python<'_>,\n result: &Bound<'_, PyAny>,\n json_module: &Bound<'_, PyAny>,\n slot_sender: &SlotSender,\n ) -> Result {\n let iter = result\n .try_iter()\n .map_err(|e| PredictionError::Failed(format!(\"Failed to iterate generator: {}\", e)))?;\n\n for item in iter {\n let item = item.map_err(|e| {\n if is_cancelation_exception(py, &e) {\n return PredictionError::Cancelled;\n }\n PredictionError::Failed(format!(\"Generator iteration error: {}\", e))\n })?;\n\n send_output_item(py, &item, json_module, slot_sender)?;\n }\n\n // Outputs already streamed over IPC — return empty stream\n Ok(PredictionOutput::Stream(vec![]))\n }\n\n /// Process single output into PredictionOutput::Single.\n ///\n /// For file outputs (Path/IOBase), the file is sent via slot_sender and\n /// an empty Single(Null) is returned since the output was already streamed.\n fn process_single_output(\n &self,\n py: Python<'_>,\n result: &Bound<'_, PyAny>,\n json_module: &Bound<'_, PyAny>,\n slot_sender: &SlotSender,\n ) -> Result {\n // Check for file-type outputs first\n let os = py\n .import(\"os\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to import os: {}\", e)))?;\n let io_mod = py\n .import(\"io\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to import io: {}\", e)))?;\n let pathlike = os\n .getattr(\"PathLike\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to get os.PathLike: {}\", e)))?;\n let iobase = io_mod\n .getattr(\"IOBase\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to get io.IOBase: {}\", e)))?;\n\n if result.is_instance(&pathlike).unwrap_or(false) {\n let path_str: String = result\n .call_method0(\"__fspath__\")\n .and_then(|p| p.extract())\n .map_err(|e| PredictionError::Failed(format!(\"Failed to get fspath: {}\", e)))?;\n slot_sender\n .send_file_output(std::path::PathBuf::from(path_str), None)\n .map_err(|e| {\n PredictionError::Failed(format!(\"Failed to send file output: {}\", e))\n })?;\n return Ok(PredictionOutput::Single(serde_json::Value::Null));\n }\n\n if result.is_instance(&iobase).unwrap_or(false) {\n if result\n .call_method0(\"seekable\")\n .and_then(|r| r.extract::())\n .unwrap_or(false)\n {\n let _ = result.call_method1(\"seek\", (0,));\n }\n let data: Vec = result\n .call_method0(\"read\")\n .and_then(|d| d.extract())\n .map_err(|e| PredictionError::Failed(format!(\"Failed to read IOBase: {}\", e)))?;\n let ext = result\n .getattr(\"name\")\n .and_then(|n| n.extract::())\n .ok()\n .and_then(|name| {\n std::path::Path::new(&name)\n .extension()\n .and_then(|e| e.to_str())\n .map(|s| s.to_string())\n })\n .unwrap_or_else(|| \"bin\".to_string());\n slot_sender\n .write_file_output(&data, &ext, None)\n .map_err(|e| {\n PredictionError::Failed(format!(\"Failed to write file output: {}\", e))\n })?;\n return Ok(PredictionOutput::Single(serde_json::Value::Null));\n }\n\n // List/tuple output — iterate items so file outputs (Path, IOBase)\n // go through the FileOutput IPC path for upload instead of being\n // base64-encoded inline by process_output.\n if let Ok(list) = result.cast::() {\n for item in list.iter() {\n send_output_item(py, &item, json_module, slot_sender)?;\n }\n return Ok(PredictionOutput::Stream(vec![]));\n }\n if let Ok(tuple) = result.cast::() {\n for item in tuple.iter() {\n send_output_item(py, &item, json_module, slot_sender)?;\n }\n return Ok(PredictionOutput::Stream(vec![]));\n }\n\n // Non-file output — process normally\n let processed = output::process_output(py, result)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to process output: {}\", e)))?;\n\n let result_str: String = json_module\n .call_method1(\"dumps\", (&processed,))\n .map_err(|e| PredictionError::Failed(format!(\"Failed to serialize output: {}\", e)))?\n .extract()\n .map_err(|e| {\n PredictionError::Failed(format!(\"Failed to extract output string: {}\", e))\n })?;\n\n let output_json: serde_json::Value = serde_json::from_str(&result_str)\n .map_err(|e| PredictionError::Failed(format!(\"Failed to parse output JSON: {}\", e)))?;\n\n Ok(PredictionOutput::Single(output_json))\n }\n\n /// Worker mode async predict - submits to shared event loop.\n ///\n /// Uses run_coroutine_threadsafe to submit the coroutine to the provided event loop.\n /// Returns the concurrent.futures.Future, is_async_gen flag, and PreparedInput for cleanup.\n /// Caller should block on future.result() to get the result, then drop PreparedInput.\n ///\n /// The prediction_id is used to set up log routing in the event loop thread.\n pub fn predict_async_worker(\n &self,\n input: serde_json::Value,\n event_loop: &Py,\n prediction_id: &str,\n ) -> Result<(Py, bool, PreparedInput), PredictionError> {\n Python::attach(|py| {\n let json_module = py.import(\"json\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to import json module: {}\", e))\n })?;\n let asyncio = py\n .import(\"asyncio\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to import asyncio: {}\", e)))?;\n\n let input_str = serde_json::to_string(&input)\n .map_err(|e| PredictionError::InvalidInput(e.to_string()))?;\n let py_input = json_module\n .call_method1(\"loads\", (input_str,))\n .map_err(|e| PredictionError::InvalidInput(format!(\"Invalid JSON input: {}\", e)))?;\n\n #[allow(deprecated)]\n let raw_input_dict = py_input.downcast::().map_err(|_| {\n PredictionError::InvalidInput(\"Input must be a JSON object\".to_string())\n })?;\n\n let func = self.predict_func(py).map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get predict function: {}\", e))\n })?;\n let prepared = input::prepare_input(py, raw_input_dict, &func)\n .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?;\n let input_dict = prepared.dict(py);\n\n // Call predict - returns coroutine\n let instance = self.instance.bind(py);\n let coro = instance\n .call_method(\"predict\", (), Some(&input_dict))\n .map_err(|e| PredictionError::Failed(format!(\"Failed to call predict: {}\", e)))?;\n\n // For async generators, wrap to collect all values\n let is_async_gen = matches!(\n &self.kind,\n PredictorKind::Class {\n predict: PredictKind::AsyncGen,\n ..\n } | PredictorKind::StandaloneFunction(PredictKind::AsyncGen)\n );\n let coro = if is_async_gen {\n let collect_fn = get_collect_async_gen(py)?;\n collect_fn\n .call1(py, (&coro,))\n .map_err(|e| {\n PredictionError::Failed(format!(\"Failed to wrap async generator: {}\", e))\n })?\n .into_bound(py)\n } else {\n coro\n };\n\n // Wrap coroutine to set up log routing in the event loop thread\n let ctx_wrapper = get_ctx_wrapper(py)?;\n\n // Get the same ContextVar instance used by SlotLogWriter for log routing\n let contextvar = crate::log_writer::get_prediction_contextvar(py).map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get prediction ContextVar: {}\", e))\n })?;\n\n // Wrap the coroutine with context setup\n let wrapped_coro = ctx_wrapper\n .call1(py, (&coro, prediction_id, contextvar.bind(py)))\n .map_err(|e| {\n PredictionError::Failed(format!(\"Failed to wrap coroutine with context: {}\", e))\n })?;\n\n // Submit wrapped coroutine to shared event loop via run_coroutine_threadsafe\n let future = asyncio\n .call_method1(\n \"run_coroutine_threadsafe\",\n (wrapped_coro.bind(py), event_loop.bind(py)),\n )\n .map_err(|e| {\n PredictionError::Failed(format!(\"Failed to submit coroutine: {}\", e))\n })?;\n\n Ok((future.unbind(), is_async_gen, prepared))\n })\n }\n\n /// Process the result from an async prediction future.\n ///\n /// Call this after future.result() returns to convert the Python result\n /// to a PredictionResult.\n pub fn process_async_result(\n &self,\n py: Python<'_>,\n result: &Bound<'_, PyAny>,\n is_async_gen: bool,\n slot_sender: &SlotSender,\n ) -> Result {\n let json_module = py\n .import(\"json\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to import json module: {}\", e)))?;\n let types_module = py.import(\"types\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to import types module: {}\", e))\n })?;\n\n // Process output\n let output = if is_async_gen {\n // Result is a pre-collected list — stream each item over IPC\n if let Ok(list) = result.extract::>>() {\n for item in list {\n send_output_item(py, &item, &json_module, slot_sender)?;\n }\n }\n PredictionOutput::Stream(vec![])\n } else {\n // Check if result is a generator (sync generator from async predict)\n let generator_type = types_module.getattr(\"GeneratorType\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get GeneratorType: {}\", e))\n })?;\n let is_generator: bool = result.is_instance(&generator_type).unwrap_or(false);\n\n if is_generator {\n self.process_generator_output(py, result, &json_module, slot_sender)?\n } else {\n self.process_single_output(py, result, &json_module, slot_sender)?\n }\n };\n\n Ok(PredictionResult {\n output,\n predict_time: None,\n logs: String::new(),\n metrics: Default::default(),\n })\n }\n\n /// Worker mode async train - submits to shared event loop.\n pub fn train_async_worker(\n &self,\n input: serde_json::Value,\n event_loop: &Py,\n prediction_id: &str,\n ) -> Result<(Py, bool, PreparedInput), PredictionError> {\n Python::attach(|py| {\n let json_module = py.import(\"json\").map_err(|e| {\n PredictionError::Failed(format!(\"Failed to import json module: {}\", e))\n })?;\n let asyncio = py\n .import(\"asyncio\")\n .map_err(|e| PredictionError::Failed(format!(\"Failed to import asyncio: {}\", e)))?;\n\n let input_str = serde_json::to_string(&input)\n .map_err(|e| PredictionError::InvalidInput(e.to_string()))?;\n let py_input = json_module\n .call_method1(\"loads\", (input_str,))\n .map_err(|e| PredictionError::InvalidInput(format!(\"Invalid JSON input: {}\", e)))?;\n\n #[allow(deprecated)]\n let raw_input_dict = py_input.downcast::().map_err(|_| {\n PredictionError::InvalidInput(\"Input must be a JSON object\".to_string())\n })?;\n\n let func = self.train_func(py).map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get train function: {}\", e))\n })?;\n let prepared = input::prepare_input(py, raw_input_dict, &func)\n .map_err(|e| PredictionError::InvalidInput(format_validation_error(py, &e)))?;\n let input_dict = prepared.dict(py);\n\n // Call train - returns coroutine\n let instance = self.instance.bind(py);\n let coro = match &self.kind {\n PredictorKind::StandaloneFunction(_) => instance.call((), Some(&input_dict)),\n PredictorKind::Class { .. } => instance.call_method(\"train\", (), Some(&input_dict)),\n }\n .map_err(|e| PredictionError::Failed(format!(\"Failed to call train: {}\", e)))?;\n\n // Wrap coroutine to set up log routing\n let ctx_wrapper = get_ctx_wrapper(py)?;\n\n // Get the same ContextVar instance used by SlotLogWriter\n let contextvar = crate::log_writer::get_prediction_contextvar(py).map_err(|e| {\n PredictionError::Failed(format!(\"Failed to get prediction ContextVar: {}\", e))\n })?;\n\n // Wrap the coroutine with context setup\n let wrapped_coro = ctx_wrapper\n .call1(py, (&coro, prediction_id, contextvar.bind(py)))\n .map_err(|e| {\n PredictionError::Failed(format!(\"Failed to wrap coroutine with context: {}\", e))\n })?;\n\n // Submit wrapped coroutine to shared event loop\n let future = asyncio\n .call_method1(\n \"run_coroutine_threadsafe\",\n (wrapped_coro.bind(py), event_loop.bind(py)),\n )\n .map_err(|e| {\n PredictionError::Failed(format!(\"Failed to submit coroutine: {}\", e))\n })?;\n\n // Train doesn't typically use async generators, but we return false for consistency\n Ok((future.unbind(), false, prepared))\n })\n }\n\n // =========================================================================\n // Healthcheck methods\n // =========================================================================\n\n /// Healthcheck timeout in seconds.\n const HEALTHCHECK_TIMEOUT: f64 = 5.0;\n\n /// Check if the predictor has a healthcheck() method.\n pub fn has_healthcheck(&self, py: Python<'_>) -> bool {\n match &self.kind {\n PredictorKind::Class { .. } => {\n let instance = self.instance.bind(py);\n instance.hasattr(\"healthcheck\").unwrap_or(false)\n }\n PredictorKind::StandaloneFunction(_) => false,\n }\n }\n\n /// Check if the healthcheck() method is async.\n pub fn is_healthcheck_async(&self, py: Python<'_>) -> bool {\n match &self.kind {\n PredictorKind::Class { .. } => {\n let instance = self.instance.bind(py);\n if let Ok(healthcheck) = instance.getattr(\"healthcheck\") {\n let inspect = py.import(\"inspect\").ok();\n if let Some(inspect) = inspect {\n inspect\n .call_method1(\"iscoroutinefunction\", (&healthcheck,))\n .ok()\n .and_then(|r| r.extract::().ok())\n .unwrap_or(false)\n } else {\n false\n }\n } else {\n false\n }\n }\n PredictorKind::StandaloneFunction(_) => false,\n }\n }\n\n /// Run a synchronous healthcheck with timeout.\n ///\n /// Runs the healthcheck in a thread pool executor with a 5 second timeout.\n pub fn healthcheck_sync(&self, py: Python<'_>) -> coglet_core::orchestrator::HealthcheckResult {\n use coglet_core::orchestrator::HealthcheckResult;\n\n let instance = self.instance.bind(py);\n\n // Run healthcheck in executor with timeout, mirroring Python impl\n let result: PyResult = (|| {\n let concurrent_futures = py.import(\"concurrent.futures\")?;\n let thread_pool = concurrent_futures.getattr(\"ThreadPoolExecutor\")?;\n\n // Create a small executor just for this healthcheck\n let executor = thread_pool.call1((1,))?;\n\n // Get the healthcheck method\n let healthcheck_fn = instance.getattr(\"healthcheck\")?;\n\n // Submit to executor\n let future = executor.call_method1(\"submit\", (healthcheck_fn,))?;\n\n // Wait with timeout\n let result = future.call_method1(\"result\", (Self::HEALTHCHECK_TIMEOUT,));\n\n // Shutdown executor\n let _ = executor.call_method1(\"shutdown\", (false,));\n\n match result {\n Ok(r) => Ok(r.extract::().unwrap_or(true)),\n Err(e) => {\n let err_str = e.to_string();\n if err_str.contains(\"TimeoutError\") {\n Err(pyo3::exceptions::PyTimeoutError::new_err(\n \"Healthcheck timed out\",\n ))\n } else {\n Err(e)\n }\n }\n }\n })();\n\n match result {\n Ok(true) => HealthcheckResult::healthy(),\n Ok(false) => HealthcheckResult::unhealthy(\n \"Healthcheck failed: user-defined healthcheck returned False\",\n ),\n Err(e) => {\n let err_str = e.to_string();\n if err_str.contains(\"TimeoutError\") {\n HealthcheckResult::unhealthy(format!(\n \"Healthcheck failed: user-defined healthcheck timed out after {:.1} seconds\",\n Self::HEALTHCHECK_TIMEOUT\n ))\n } else {\n HealthcheckResult::unhealthy(format!(\"Healthcheck failed: {}\", e))\n }\n }\n }\n }\n\n /// Run an async healthcheck with timeout.\n ///\n /// Runs the healthcheck in the async event loop with a 5 second timeout.\n pub fn healthcheck_async(\n &self,\n py: Python<'_>,\n event_loop: &Py,\n ) -> coglet_core::orchestrator::HealthcheckResult {\n use coglet_core::orchestrator::HealthcheckResult;\n\n let instance = self.instance.bind(py);\n\n let result: PyResult = (|| {\n let asyncio = py.import(\"asyncio\")?;\n\n // Get the healthcheck coroutine\n let healthcheck_fn = instance.getattr(\"healthcheck\")?;\n let coro = healthcheck_fn.call0()?;\n\n // Wrap with timeout\n let wait_for = asyncio.getattr(\"wait_for\")?;\n let timeout_coro = wait_for.call1((&coro, Self::HEALTHCHECK_TIMEOUT))?;\n\n // Submit to event loop\n let future = asyncio.call_method1(\n \"run_coroutine_threadsafe\",\n (&timeout_coro, event_loop.bind(py)),\n )?;\n\n // Block on result with extra buffer time for event loop overhead\n let result = future.call_method1(\"result\", (Self::HEALTHCHECK_TIMEOUT + 1.0,));\n\n match result {\n Ok(r) => Ok(r.extract::().unwrap_or(true)),\n Err(e) => {\n let err_str = e.to_string();\n if err_str.contains(\"TimeoutError\") || err_str.contains(\"timed out\") {\n Err(pyo3::exceptions::PyTimeoutError::new_err(\n \"Healthcheck timed out\",\n ))\n } else {\n Err(e)\n }\n }\n }\n })();\n\n match result {\n Ok(true) => HealthcheckResult::healthy(),\n Ok(false) => HealthcheckResult::unhealthy(\n \"Healthcheck failed: user-defined healthcheck returned False\",\n ),\n Err(e) => {\n let err_str = e.to_string();\n if err_str.contains(\"TimeoutError\") {\n HealthcheckResult::unhealthy(format!(\n \"Healthcheck failed: user-defined healthcheck timed out after {:.1} seconds\",\n Self::HEALTHCHECK_TIMEOUT\n ))\n } else {\n HealthcheckResult::unhealthy(format!(\"Healthcheck failed: {}\", e))\n }\n }\n }\n }\n}\n" }, { "path": "crates/coglet-python/src/worker_bridge.rs", "content": "//! Bridge between coglet-worker's PredictHandler trait and PythonPredictor.\n\nuse std::collections::HashMap;\nuse std::sync::{Arc, Mutex};\nuse std::thread::JoinHandle;\n\nuse pyo3::prelude::*;\n\nuse coglet_core::bridge::protocol::SlotId;\nuse coglet_core::worker::{PredictHandler, PredictResult, SetupError, SlotSender};\n\nuse crate::predictor::PythonPredictor;\n\n/// What operation the handler performs\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum HandlerMode {\n /// Calls predict() method\n Predict,\n /// Calls train() method\n Train,\n}\n\n/// SDK implementation type detected from the Python predictor.\n///\n/// This enum allows for future extensibility if additional SDK\n/// implementations are needed (e.g., Node.js).\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum SdkImplementation {\n /// Standard cog Python SDK\n Cog,\n /// Unable to detect SDK type\n Unknown,\n}\n\nimpl std::fmt::Display for SdkImplementation {\n fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {\n match self {\n Self::Cog => write!(f, \"cog\"),\n Self::Unknown => write!(f, \"unknown\"),\n }\n }\n}\n\n/// Current state of a prediction slot\n#[derive(Debug, Default)]\npub enum SlotState {\n /// No prediction running\n #[default]\n Idle,\n /// Sync prediction in progress\n SyncPrediction {\n cancelled: bool,\n /// Python thread identifier (for `PyThreadState_SetAsyncExc`)\n py_thread_id: std::ffi::c_long,\n },\n /// Async prediction in progress\n AsyncPrediction {\n /// Future for cancellation\n future: Py,\n cancelled: bool,\n },\n}\n\nimpl SlotState {\n pub fn is_cancelled(&self) -> bool {\n match self {\n SlotState::SyncPrediction { cancelled, .. } => *cancelled,\n SlotState::AsyncPrediction { cancelled, .. } => *cancelled,\n SlotState::Idle => false,\n }\n }\n\n pub fn mark_cancelled(&mut self) {\n match self {\n SlotState::SyncPrediction { cancelled, .. } => *cancelled = true,\n SlotState::AsyncPrediction { cancelled, .. } => *cancelled = true,\n SlotState::Idle => { /* no-op */ }\n }\n }\n}\n\n/// Wraps PythonPredictor to implement the PredictHandler trait.\n///\n/// The `is_train` flag determines whether predict() calls the Python\n/// predict() or train() method. This is set at construction time.\n///\n/// BUG-FOR-BUG COMPATIBILITY: In cog mainline, training routes use a worker\n/// that was created with is_train=false, so training routes actually call\n/// predict() instead of train(). We replicate this by always creating the\n/// handler with is_train=false. To fix this bug, pass is_train=true when\n/// creating a handler for training routes.\npub struct PythonPredictHandler {\n predictor_ref: String,\n predictor: Mutex>>,\n /// Per-slot cancellation state (keyed by SlotId).\n slots: Mutex>,\n /// What operation this handler performs (predict or train).\n /// BUG: cog mainline always uses Predict mode, even for training routes.\n mode: HandlerMode,\n /// Shared asyncio event loop for async predictions (runs in dedicated thread).\n async_loop: Mutex>>,\n /// Handle to the asyncio loop thread for joining on shutdown.\n async_thread: Mutex>>,\n}\n\nimpl PythonPredictHandler {\n /// Create a handler in prediction mode.\n pub fn new(predictor_ref: String) -> Result {\n let (loop_obj, thread) = Self::init_async_loop()?;\n Ok(Self {\n predictor_ref,\n predictor: Mutex::new(None),\n slots: Mutex::new(HashMap::new()),\n mode: HandlerMode::Predict,\n async_loop: Mutex::new(Some(loop_obj)),\n async_thread: Mutex::new(Some(thread)),\n })\n }\n\n /// Create a handler in training mode.\n ///\n /// NOTE: For bug-for-bug compatibility with cog mainline, use new() instead.\n /// Cog mainline's training routes incorrectly use a predict-mode worker.\n #[allow(dead_code)]\n pub fn new_train(predictor_ref: String) -> Result {\n let (loop_obj, thread) = Self::init_async_loop()?;\n Ok(Self {\n predictor_ref,\n predictor: Mutex::new(None),\n slots: Mutex::new(HashMap::new()),\n mode: HandlerMode::Train,\n async_loop: Mutex::new(Some(loop_obj)),\n async_thread: Mutex::new(Some(thread)),\n })\n }\n\n /// Initialize the shared asyncio event loop in a dedicated thread.\n fn init_async_loop() -> Result<(Py, JoinHandle<()>), SetupError> {\n Python::attach(|py| {\n let asyncio = py\n .import(\"asyncio\")\n .map_err(|e| SetupError::internal(format!(\"Failed to import asyncio: {}\", e)))?;\n let loop_obj = asyncio\n .call_method0(\"new_event_loop\")\n .map_err(|e| SetupError::internal(format!(\"Failed to create event loop: {}\", e)))?;\n\n // Clone for the thread\n let loop_for_thread = loop_obj.clone().unbind();\n let loop_result = loop_obj.unbind();\n\n // Spawn thread running loop.run_forever()\n let thread = std::thread::spawn(move || {\n Python::attach(|py| {\n let loop_ref = loop_for_thread.bind(py);\n // These errors in the thread are logged but can't be propagated\n // The thread dying will cause async predictions to fail\n let Ok(asyncio) = py.import(\"asyncio\") else {\n tracing::error!(\"Failed to import asyncio in loop thread\");\n return;\n };\n if let Err(e) = asyncio.call_method1(\"set_event_loop\", (loop_ref,)) {\n tracing::error!(error = %e, \"Failed to set event loop\");\n return;\n }\n tracing::trace!(\"Asyncio event loop thread starting\");\n if let Err(e) = loop_ref.call_method0(\"run_forever\") {\n tracing::error!(error = %e, \"Asyncio event loop error\");\n }\n tracing::trace!(\"Asyncio event loop thread exiting\");\n });\n });\n\n Ok((loop_result, thread))\n })\n }\n\n // NOTE: All mutex locks in this file use .expect().\n // See log_writer.rs for the full rationale. Short version: poisoned mutex\n // means slot isolation is compromised. The panic hook installed by\n // coglet_core::worker sends a Fatal IPC message and aborts.\n\n /// Check the cancelled flag for a slot without clearing it.\n fn is_cancelled(&self, slot: SlotId) -> bool {\n let slots = self.slots.lock().expect(\"slots mutex poisoned\");\n slots.get(&slot).is_some_and(|s| s.is_cancelled())\n }\n\n /// Check and clear the cancelled flag for a slot.\n fn take_cancelled(&self, slot: SlotId) -> bool {\n let mut slots = self.slots.lock().expect(\"slots mutex poisoned\");\n let state = slots.entry(slot).or_default();\n let was_cancelled = state.is_cancelled();\n // Reset to idle after checking cancellation\n if was_cancelled {\n *state = SlotState::Idle;\n }\n was_cancelled\n }\n\n /// Mark a slot as having a sync prediction in progress.\n ///\n /// `py_thread_id` is the Python thread identifier of the thread that will\n /// run the prediction, for use with `PyThreadState_SetAsyncExc` on cancel.\n fn start_sync_prediction(&self, slot: SlotId, py_thread_id: std::ffi::c_long) {\n let mut slots = self.slots.lock().expect(\"slots mutex poisoned\");\n slots.insert(\n slot,\n SlotState::SyncPrediction {\n cancelled: false,\n py_thread_id,\n },\n );\n }\n\n /// Mark a slot as having an async prediction in progress.\n fn start_async_prediction(&self, slot: SlotId, future: Py) {\n let mut slots = self.slots.lock().expect(\"slots mutex poisoned\");\n slots.insert(\n slot,\n SlotState::AsyncPrediction {\n future,\n cancelled: false,\n },\n );\n }\n\n /// Clear prediction state for a slot.\n fn finish_prediction(&self, slot: SlotId) {\n let mut slots = self.slots.lock().expect(\"slots mutex poisoned\");\n slots.insert(slot, SlotState::Idle);\n }\n\n /// Cancel an async prediction using future.cancel().\n /// Returns true if cancellation was requested, false if no future found.\n fn cancel_async_future(&self, slot: SlotId) -> bool {\n Python::attach(|py| {\n let future = {\n let slots = self.slots.lock().expect(\"slots mutex poisoned\");\n if let Some(SlotState::AsyncPrediction { future, .. }) = slots.get(&slot) {\n Some(future.clone_ref(py))\n } else {\n None\n }\n };\n\n if let Some(future) = future {\n match future.call_method0(py, \"cancel\") {\n Ok(_) => {\n tracing::trace!(%slot, \"Cancelled async future\");\n true\n }\n Err(e) => {\n tracing::warn!(%slot, error = %e, \"Failed to cancel async future\");\n false\n }\n }\n } else {\n tracing::trace!(%slot, \"No async future to cancel\");\n false\n }\n })\n }\n\n /// Get a reference to the shared asyncio event loop.\n fn get_async_loop(&self) -> Option> {\n Python::attach(|py| {\n self.async_loop\n .lock()\n .expect(\"async_loop mutex poisoned\")\n .as_ref()\n .map(|l| l.clone_ref(py))\n })\n }\n}\n\n#[async_trait::async_trait]\nimpl PredictHandler for PythonPredictHandler {\n async fn setup(&self) -> Result<(), SetupError> {\n Python::attach(|py| {\n tracing::info!(predictor_ref = %self.predictor_ref, \"Loading predictor\");\n\n let pred = PythonPredictor::load(py, &self.predictor_ref)\n .map_err(|e| SetupError::load(e.to_string()))?;\n\n // Detect SDK implementation\n let sdk_impl = match py.import(\"cog\") {\n Ok(cog) => match cog.getattr(\"BasePredictor\") {\n Ok(_) => SdkImplementation::Cog,\n Err(_) => SdkImplementation::Unknown,\n },\n Err(_) => SdkImplementation::Unknown,\n };\n tracing::info!(sdk_implementation = %sdk_impl, \"Detected Cog SDK implementation\");\n\n tracing::info!(\"Running setup\");\n pred.setup(py)\n .map_err(|e| SetupError::setup(e.to_string()))?;\n\n let mut guard = self.predictor.lock().expect(\"predictor mutex poisoned\");\n *guard = Some(Arc::new(pred));\n\n tracing::info!(\"Setup complete\");\n Ok(())\n })\n }\n\n async fn predict(\n &self,\n slot: SlotId,\n id: String,\n input: serde_json::Value,\n slot_sender: Arc,\n context: HashMap,\n ) -> PredictResult {\n tracing::trace!(%slot, %id, \"PythonPredictHandler::predict starting\");\n\n // Get predictor\n let pred = {\n let guard = self.predictor.lock().expect(\"predictor mutex poisoned\");\n match guard.as_ref() {\n Some(p) => Arc::clone(p),\n None => {\n return PredictResult::failed(\"Predictor not initialized\".to_string(), 0.0);\n }\n }\n };\n let is_async = pred.is_async();\n tracing::trace!(%slot, %id, is_async, \"Got predictor\");\n\n // Track that we're starting a prediction on this slot.\n // Capture the Python thread ID for this thread (used by\n // PyThreadState_SetAsyncExc to inject CancelationException on cancel).\n // For async predictions, the slot state is updated later with the future.\n let py_thread_id = crate::cancel::current_py_thread_id();\n self.start_sync_prediction(slot, py_thread_id);\n\n // Check cancellation first (in case cancel was called before we started)\n if self.take_cancelled(slot) {\n self.finish_prediction(slot);\n return PredictResult::cancelled(0.0);\n }\n\n // Enter prediction context - sets cog_prediction_id ContextVar for log routing\n let prediction_id = id.clone();\n let slot_sender_clone = slot_sender.clone();\n let log_guard = Python::attach(|py| {\n crate::log_writer::PredictionLogGuard::enter(\n py,\n prediction_id.clone(),\n slot_sender_clone,\n )\n });\n let log_guard = match log_guard {\n Ok(g) => Some(g),\n Err(e) => {\n tracing::warn!(error = %e, \"Failed to enter prediction context\");\n None\n }\n };\n\n // Enter metric scope - sets Scope ContextVar for metric recording\n let slot_sender_for_metrics = slot_sender.clone();\n let scope_guard = Python::attach(|py| {\n crate::metric_scope::ScopeGuard::enter(py, slot_sender_for_metrics, context)\n });\n let scope_guard = match scope_guard {\n Ok(g) => Some(g),\n Err(e) => {\n tracing::warn!(error = %e, \"Failed to enter metric scope\");\n None\n }\n };\n\n tracing::trace!(%slot, %id, \"Prediction context entered\");\n\n // Run prediction or training based on mode.\n let start = std::time::Instant::now();\n\n let result = match self.mode {\n HandlerMode::Train => {\n // Training mode - check if train() exists\n if !pred.has_train() {\n self.finish_prediction(slot);\n return PredictResult::failed(\n \"Training not supported by this predictor\".to_string(),\n 0.0,\n );\n }\n // Use worker-mode train\n if pred.is_train_async() {\n // Async train - submit to shared event loop\n let loop_obj = match self.get_async_loop() {\n Some(l) => l,\n None => {\n return PredictResult::failed(\n \"Async event loop not initialized\".to_string(),\n start.elapsed().as_secs_f64(),\n );\n }\n };\n\n // Submit coroutine and get future + prepared input for cleanup\n let (future, is_async_gen, prepared) = match pred\n .train_async_worker(input, &loop_obj, &id)\n {\n Ok(f) => f,\n Err(e) => {\n self.finish_prediction(slot);\n drop(log_guard);\n return if matches!(e, coglet_core::PredictionError::Cancelled) {\n PredictResult::cancelled(start.elapsed().as_secs_f64())\n } else {\n PredictResult::failed(e.to_string(), start.elapsed().as_secs_f64())\n };\n }\n };\n\n // Update slot state with future for cancellation\n Python::attach(|py| {\n self.start_async_prediction(slot, future.clone_ref(py));\n });\n\n // Block on future.result()\n let sender_for_async = slot_sender.clone();\n let result = Python::attach(|py| match future.call_method0(py, \"result\") {\n Ok(result) => pred.process_async_result(\n py,\n result.bind(py),\n is_async_gen,\n &sender_for_async,\n ),\n Err(e) => {\n let err_str = e.to_string();\n if err_str.contains(\"CancelledError\") || err_str.contains(\"cancelled\") {\n Err(coglet_core::PredictionError::Cancelled)\n } else {\n Err(coglet_core::PredictionError::Failed(format!(\n \"Async training failed: {}\",\n e\n )))\n }\n }\n });\n\n // Cleanup temp files via RAII\n drop(prepared);\n\n result\n } else {\n // Sync train - set sync prediction ID for log routing\n crate::log_writer::set_sync_prediction_id(Some(&id));\n let r = pred.train_worker(input, slot_sender.clone());\n crate::log_writer::set_sync_prediction_id(None);\n\n // Upgrade to Cancelled if the slot was marked cancelled\n // (same logic as sync predict above)\n match r {\n Err(_) if self.is_cancelled(slot) => {\n Err(coglet_core::PredictionError::Cancelled)\n }\n other => other,\n }\n }\n }\n HandlerMode::Predict => {\n // Prediction mode\n tracing::trace!(%slot, %id, is_async = pred.is_async(), \"Running prediction\");\n if pred.is_async() {\n // Async predict - submit to shared event loop\n let loop_obj = match self.get_async_loop() {\n Some(l) => l,\n None => {\n return PredictResult::failed(\n \"Async event loop not initialized\".to_string(),\n start.elapsed().as_secs_f64(),\n );\n }\n };\n\n // Submit coroutine and get future + prepared input for cleanup\n let (future, is_async_gen, prepared) = match pred\n .predict_async_worker(input, &loop_obj, &id)\n {\n Ok(f) => f,\n Err(e) => {\n self.finish_prediction(slot);\n drop(log_guard);\n return if matches!(e, coglet_core::PredictionError::Cancelled) {\n PredictResult::cancelled(start.elapsed().as_secs_f64())\n } else {\n PredictResult::failed(e.to_string(), start.elapsed().as_secs_f64())\n };\n }\n };\n\n // Update slot state with future for cancellation\n Python::attach(|py| {\n self.start_async_prediction(slot, future.clone_ref(py));\n });\n\n // Block on future.result()\n let sender_for_async = slot_sender.clone();\n let result = Python::attach(|py| match future.call_method0(py, \"result\") {\n Ok(result) => pred.process_async_result(\n py,\n result.bind(py),\n is_async_gen,\n &sender_for_async,\n ),\n Err(e) => {\n let err_str = e.to_string();\n if err_str.contains(\"CancelledError\") || err_str.contains(\"cancelled\") {\n Err(coglet_core::PredictionError::Cancelled)\n } else {\n Err(coglet_core::PredictionError::Failed(format!(\n \"Async prediction failed: {}\",\n e\n )))\n }\n }\n });\n\n // Cleanup temp files via RAII\n drop(prepared);\n\n result\n } else {\n // Sync predict - set sync prediction ID for log routing\n crate::log_writer::set_sync_prediction_id(Some(&id));\n tracing::trace!(%slot, %id, \"Calling predict_worker\");\n let r = pred.predict_worker(input, slot_sender.clone());\n tracing::trace!(%slot, %id, \"predict_worker returned\");\n crate::log_writer::set_sync_prediction_id(None);\n\n // If the prediction failed AND the slot was marked cancelled,\n // treat it as a cancellation. PyThreadState_SetAsyncExc injects\n // CancelationException which predict_worker sees as a generic\n // PyErr — we upgrade it to Cancelled here.\n match r {\n Err(_) if self.is_cancelled(slot) => {\n Err(coglet_core::PredictionError::Cancelled)\n }\n other => other,\n }\n }\n }\n };\n tracing::trace!(%slot, %id, \"Prediction completed\");\n\n self.finish_prediction(slot);\n\n // Exit prediction context\n drop(scope_guard);\n drop(log_guard);\n\n match result {\n Ok(r) => {\n let is_stream = r.output.is_stream();\n PredictResult::success(\n output_to_json(r.output),\n start.elapsed().as_secs_f64(),\n is_stream,\n )\n }\n Err(e) => {\n if matches!(e, coglet_core::PredictionError::Cancelled) {\n PredictResult::cancelled(start.elapsed().as_secs_f64())\n } else {\n PredictResult::failed(e.to_string(), start.elapsed().as_secs_f64())\n }\n }\n }\n }\n\n fn cancel(&self, slot: SlotId) {\n // Mark slot as cancelled and determine how to cancel based on state\n let mut slots = self.slots.lock().expect(\"slots mutex poisoned\");\n\n if let Some(state) = slots.get_mut(&slot) {\n state.mark_cancelled();\n\n match state {\n SlotState::AsyncPrediction { .. } => {\n drop(slots); // Release lock before calling cancel_async_future\n // Async: cancel via future.cancel()\n if !self.cancel_async_future(slot) {\n tracing::trace!(%slot, \"No async future to cancel (prediction may have completed)\");\n }\n }\n SlotState::SyncPrediction { py_thread_id, .. } => {\n let py_thread_id = *py_thread_id;\n drop(slots); // Release lock\n // Sync: inject CancelationException into the Python thread\n // via PyThreadState_SetAsyncExc (fires at next bytecode boundary)\n crate::cancel::cancel_sync_thread(py_thread_id);\n }\n SlotState::Idle => {\n // Already idle, nothing to cancel\n tracing::trace!(%slot, \"Cancel called on idle slot\");\n }\n }\n } else {\n tracing::trace!(%slot, \"Cancel called on unknown slot\");\n }\n }\n\n async fn healthcheck(&self) -> coglet_core::orchestrator::HealthcheckResult {\n // Get predictor\n let pred = {\n let guard = self.predictor.lock().expect(\"predictor mutex poisoned\");\n match guard.as_ref() {\n Some(p) => Arc::clone(p),\n None => {\n return coglet_core::orchestrator::HealthcheckResult::unhealthy(\n \"Predictor not initialized\",\n );\n }\n }\n };\n\n // Check if predictor has a healthcheck method\n let has_healthcheck = Python::attach(|py| pred.has_healthcheck(py));\n if !has_healthcheck {\n // No healthcheck defined = healthy\n return coglet_core::orchestrator::HealthcheckResult::healthy();\n }\n\n // Run healthcheck with timeout\n let is_async = Python::attach(|py| pred.is_healthcheck_async(py));\n\n if is_async {\n // Async healthcheck - run in event loop with timeout\n let loop_obj = match self.get_async_loop() {\n Some(l) => l,\n None => {\n return coglet_core::orchestrator::HealthcheckResult::unhealthy(\n \"Async event loop not initialized\",\n );\n }\n };\n\n Python::attach(|py| pred.healthcheck_async(py, &loop_obj))\n } else {\n // Sync healthcheck - run in thread pool with timeout\n Python::attach(|py| pred.healthcheck_sync(py))\n }\n }\n}\n\n/// Shutdown the asyncio event loop and join the thread.\nimpl Drop for PythonPredictHandler {\n fn drop(&mut self) {\n // Stop the event loop\n if let Some(loop_obj) = self\n .async_loop\n .lock()\n .expect(\"async_loop mutex poisoned\")\n .take()\n {\n Python::attach(|py| {\n let loop_ref = loop_obj.bind(py);\n // Get the stop method and schedule it via call_soon_threadsafe\n match loop_ref.getattr(\"stop\") {\n Ok(stop_method) => {\n if let Err(e) =\n loop_ref.call_method1(\"call_soon_threadsafe\", (stop_method,))\n {\n tracing::warn!(error = %e, \"Failed to stop asyncio loop\");\n }\n }\n Err(e) => {\n tracing::warn!(error = %e, \"Failed to get loop.stop method\");\n }\n }\n });\n }\n\n // Join the thread\n if let Some(thread) = self\n .async_thread\n .lock()\n .expect(\"async_thread mutex poisoned\")\n .take()\n && let Err(e) = thread.join()\n {\n tracing::warn!(\"Failed to join asyncio loop thread: {:?}\", e);\n }\n }\n}\n\n/// Convert PredictionOutput to serde_json::Value\nfn output_to_json(output: coglet_core::PredictionOutput) -> serde_json::Value {\n match output {\n coglet_core::PredictionOutput::Single(v) => v,\n coglet_core::PredictionOutput::Stream(v) => serde_json::Value::Array(v),\n }\n}\n" }, { "path": "crates/coglet-python/tests/test_coglet.py", "content": "\"\"\"Tests for coglet Python bindings.\"\"\"\n\nimport queue\nimport re\nimport socket\nimport subprocess\nimport sys\nimport threading\nimport time\nfrom pathlib import Path\n\nimport coglet\nimport pytest\nimport requests\n\n# =============================================================================\n# Module structure tests (no server needed)\n# =============================================================================\n\n\nclass TestModuleStructure:\n \"\"\"Tests for coglet module public API and structure.\"\"\"\n\n def test_version_is_pep440(self) -> None:\n \"\"\"__version__ must be a valid PEP 440 version string.\"\"\"\n # PEP 440: N.N.N, N.N.NaN, N.N.NbN, N.N.NrcN, N.N.N.devN, etc.\n assert re.match(\n r\"^\\d+\\.\\d+\\.\\d+(\\.dev\\d+|a\\d+|b\\d+|rc\\d+)?(\\+.+)?$\",\n coglet.__version__,\n ), f\"Not PEP 440: {coglet.__version__!r}\"\n\n def test_version_is_str(self) -> None:\n assert isinstance(coglet.__version__, str)\n\n def test_build_info_exists(self) -> None:\n build = coglet.__build__\n assert hasattr(build, \"version\")\n assert hasattr(build, \"git_sha\")\n assert hasattr(build, \"build_time\")\n assert hasattr(build, \"rustc_version\")\n\n def test_build_info_fields_are_strings(self) -> None:\n build = coglet.__build__\n assert isinstance(build.version, str)\n assert isinstance(build.git_sha, str)\n assert isinstance(build.build_time, str)\n assert isinstance(build.rustc_version, str)\n\n def test_build_info_version_matches_module_version(self) -> None:\n assert coglet.__build__.version == coglet.__version__\n\n def test_build_info_repr(self) -> None:\n r = repr(coglet.__build__)\n assert r.startswith(\"BuildInfo(\")\n assert \"version=\" in r\n assert \"git_sha=\" in r\n\n def test_build_info_frozen(self) -> None:\n with pytest.raises(AttributeError):\n coglet.__build__.version = \"hacked\" # type: ignore[misc]\n\n def test_server_exists(self) -> None:\n assert hasattr(coglet, \"server\")\n\n def test_server_active_is_false(self) -> None:\n \"\"\"Outside a worker subprocess, active should be False.\"\"\"\n assert coglet.server.active is False\n\n def test_server_active_is_property(self) -> None:\n \"\"\"active should be a property (no parens needed), not callable.\"\"\"\n assert isinstance(coglet.server.active, bool)\n\n def test_server_frozen(self) -> None:\n with pytest.raises(AttributeError):\n coglet.server.foo = \"bar\" # type: ignore[attr-defined]\n\n def test_server_active_not_settable(self) -> None:\n with pytest.raises(AttributeError):\n coglet.server.active = True # type: ignore[misc]\n\n def test_server_repr(self) -> None:\n assert repr(coglet.server) == \"coglet.server\"\n\n def test_sdk_submodule_exists(self) -> None:\n assert hasattr(coglet, \"_sdk\")\n\n def test_sdk_has_slot_log_writer(self) -> None:\n assert hasattr(coglet._sdk, \"_SlotLogWriter\")\n\n def test_sdk_has_tee_writer(self) -> None:\n assert hasattr(coglet._sdk, \"_TeeWriter\")\n\n def test_all_excludes_internals(self) -> None:\n \"\"\"__all__ should only list public API.\"\"\"\n assert \"__version__\" in coglet.__all__\n assert \"__build__\" in coglet.__all__\n assert \"server\" in coglet.__all__\n # _sdk should not be in __all__ (underscore = private)\n assert \"_sdk\" not in coglet.__all__\n assert \"_impl\" not in coglet.__all__\n\n\n@pytest.fixture\ndef sync_predictor(tmp_path: Path) -> Path:\n \"\"\"Create a simple sync predictor.\"\"\"\n predictor = tmp_path / \"predict.py\"\n predictor.write_text(\"\"\"\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def setup(self):\n self.prefix = \"Hello, \"\n\n def predict(self, name: str = \"World\") -> str:\n return self.prefix + name + \"!\"\n\"\"\")\n\n # Create cog.yaml\n cog_yaml = tmp_path / \"cog.yaml\"\n cog_yaml.write_text(\"\"\"\npredict: \"predict.py:Predictor\"\n\"\"\")\n\n return predictor\n\n\n@pytest.fixture\ndef generator_predictor(tmp_path: Path) -> Path:\n \"\"\"Create a generator predictor.\"\"\"\n predictor = tmp_path / \"predict.py\"\n predictor.write_text(\"\"\"\nfrom cog import BasePredictor\nfrom typing import Iterator\n\nclass Predictor(BasePredictor):\n def setup(self):\n pass\n\n def predict(self, count: int = 3) -> Iterator[str]:\n for i in range(count):\n yield f\"chunk {i}\"\n\"\"\")\n\n # Create cog.yaml\n cog_yaml = tmp_path / \"cog.yaml\"\n cog_yaml.write_text(\"\"\"\npredict: \"predict.py:Predictor\"\n\"\"\")\n\n return predictor\n\n\n@pytest.fixture\ndef async_predictor(tmp_path: Path) -> Path:\n \"\"\"Create an async predictor.\"\"\"\n predictor = tmp_path / \"predict.py\"\n predictor.write_text(\"\"\"\nimport asyncio\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def setup(self):\n self.call_count = 0\n\n async def predict(self, delay: float = 0.1, name: str = \"test\") -> str:\n self.call_count += 1\n await asyncio.sleep(delay)\n return f\"{name}: done after {delay}s (call #{self.call_count})\"\n\"\"\")\n\n # Create cog.yaml\n cog_yaml = tmp_path / \"cog.yaml\"\n cog_yaml.write_text(\"\"\"\npredict: \"predict.py:Predictor\"\n\"\"\")\n\n return predictor\n\n\n@pytest.fixture\ndef async_generator_predictor(tmp_path: Path) -> Path:\n \"\"\"Create an async generator predictor.\"\"\"\n predictor = tmp_path / \"predict.py\"\n predictor.write_text(\"\"\"\nimport asyncio\nfrom cog import BasePredictor\nfrom typing import AsyncIterator\n\nclass Predictor(BasePredictor):\n def setup(self):\n pass\n\n async def predict(self, count: int = 3, delay: float = 0.05) -> AsyncIterator[str]:\n for i in range(count):\n await asyncio.sleep(delay)\n yield f\"async chunk {i}\"\n\"\"\")\n\n # Create cog.yaml\n cog_yaml = tmp_path / \"cog.yaml\"\n cog_yaml.write_text(\"\"\"\npredict: \"predict.py:Predictor\"\n\"\"\")\n\n return predictor\n\n\nclass CogletServer:\n \"\"\"Context manager for running coglet server.\"\"\"\n\n def __init__(self, predictor_path: Path, port: int = 0):\n self.predictor_path = predictor_path\n self.requested_port = port\n self.port = None\n self.process = None\n self.stderr_lines = []\n self.stderr_queue = queue.Queue()\n self.stderr_thread = None\n\n def __enter__(self):\n cmd = [\n sys.executable,\n \"-c\",\n f\"import coglet; coglet.server.serve('{self.predictor_path}:Predictor', port={self.requested_port})\",\n ]\n self.process = subprocess.Popen(\n cmd,\n stdout=subprocess.PIPE,\n stderr=subprocess.PIPE,\n text=True,\n bufsize=1, # Line buffered\n cwd=str(\n self.predictor_path.parent\n ), # Run from predictor directory to find cog.yaml\n )\n\n # Start background thread to read stderr\n self.stderr_thread = threading.Thread(target=self._read_stderr, daemon=True)\n self.stderr_thread.start()\n\n # Discover actual port from logs\n self._discover_port()\n # Wait for server to become ready\n self._wait_for_ready()\n return self\n\n def __exit__(self, exc_type, exc_val, exc_tb):\n if self.process:\n self.process.terminate()\n self.process.wait(timeout=5)\n\n def _read_stderr(self):\n \"\"\"Background thread that reads stderr and queues lines.\"\"\"\n try:\n for line in self.process.stderr:\n self.stderr_lines.append(line)\n self.stderr_queue.put(line)\n except Exception:\n pass # Process terminated\n\n def _discover_port(self, timeout: float = 5.0):\n \"\"\"Read stderr until we find the port the server bound to.\"\"\"\n start = time.time()\n while time.time() - start < timeout:\n try:\n line = self.stderr_queue.get(timeout=0.1)\n except queue.Empty:\n if self.process.poll() is not None:\n # Process died\n raise RuntimeError(\n f\"Server process died during startup\\nSTDERR:\\n{''.join(self.stderr_lines)}\"\n )\n continue\n\n # Look for: \"Starting coglet server on 0.0.0.0:PORT\"\n match = re.search(r\"Starting coglet server on [\\d.]+:(\\d+)\", line)\n if match:\n self.port = int(match.group(1))\n return\n\n raise TimeoutError(\n f\"Could not discover server port within {timeout}s\\nSTDERR:\\n{''.join(self.stderr_lines)}\"\n )\n\n def _wait_for_ready(self, timeout: float = 10.0):\n start = time.time()\n while time.time() - start < timeout:\n try:\n resp = requests.get(\n f\"http://localhost:{self.port}/health-check\", timeout=1\n )\n if resp.status_code == 200 and resp.json().get(\"status\") == \"READY\":\n return\n except requests.exceptions.ConnectionError:\n pass\n time.sleep(0.1)\n\n # Terminate on failure\n if self.process and self.process.poll() is None:\n self.process.terminate()\n self.process.wait(timeout=2)\n\n raise TimeoutError(\n f\"Server did not become ready within {timeout}s (port={self.port})\\n\"\n f\"STDERR:\\n{''.join(self.stderr_lines)}\"\n )\n\n @property\n def base_url(self) -> str:\n return f\"http://localhost:{self.port}\"\n\n def health_check(self) -> dict:\n resp = requests.get(f\"{self.base_url}/health-check\")\n resp.raise_for_status()\n return resp.json()\n\n def predict(self, input_data: dict) -> dict:\n resp = requests.post(\n f\"{self.base_url}/predictions\",\n json={\"input\": input_data},\n )\n return resp.json()\n\n\nclass TestHealthCheck:\n \"\"\"Tests for health check endpoint.\"\"\"\n\n def test_returns_ready_status(self, sync_predictor: Path):\n with CogletServer(sync_predictor) as server:\n health = server.health_check()\n assert health[\"status\"] == \"READY\"\n\n def test_returns_version_info(self, sync_predictor: Path):\n with CogletServer(sync_predictor) as server:\n health = server.health_check()\n assert \"version\" in health\n assert \"coglet\" in health[\"version\"]\n assert \"python\" in health[\"version\"]\n assert \"python_sdk\" in health[\"version\"]\n\n\nclass TestSyncPredictor:\n \"\"\"Tests for sync predictor.\"\"\"\n\n def test_basic_prediction(self, sync_predictor: Path):\n with CogletServer(sync_predictor) as server:\n result = server.predict({\"name\": \"Claude\"})\n assert result[\"status\"] == \"succeeded\"\n assert result[\"output\"] == \"Hello, Claude!\"\n\n def test_default_input(self, sync_predictor: Path):\n with CogletServer(sync_predictor) as server:\n result = server.predict({})\n assert result[\"status\"] == \"succeeded\"\n assert result[\"output\"] == \"Hello, World!\"\n\n def test_includes_predict_time(self, sync_predictor: Path):\n with CogletServer(sync_predictor) as server:\n result = server.predict({\"name\": \"test\"})\n assert \"metrics\" in result\n assert \"predict_time\" in result[\"metrics\"]\n assert result[\"metrics\"][\"predict_time\"] >= 0\n\n\nclass TestGeneratorPredictor:\n \"\"\"Tests for generator predictor.\"\"\"\n\n def test_returns_array_output(self, generator_predictor: Path):\n with CogletServer(generator_predictor) as server:\n result = server.predict({\"count\": 3})\n assert result[\"status\"] == \"succeeded\"\n assert result[\"output\"] == [\"chunk 0\", \"chunk 1\", \"chunk 2\"]\n\n def test_custom_count(self, generator_predictor: Path):\n with CogletServer(generator_predictor) as server:\n result = server.predict({\"count\": 5})\n assert len(result[\"output\"]) == 5\n\n\nclass TestAsyncPredictor:\n \"\"\"Tests for async predictor.\"\"\"\n\n def test_basic_prediction(self, async_predictor: Path):\n with CogletServer(async_predictor) as server:\n result = server.predict({\"delay\": 0.1, \"name\": \"async\"})\n assert result[\"status\"] == \"succeeded\"\n assert \"async: done\" in result[\"output\"]\n\n def test_sequential_requests(self, async_predictor: Path):\n \"\"\"Sequential requests both succeed (subprocess isolation means no concurrency).\"\"\"\n with CogletServer(async_predictor) as server:\n # Run two sequential requests\n result1 = server.predict({\"delay\": 0.1, \"name\": \"req1\"})\n result2 = server.predict({\"delay\": 0.1, \"name\": \"req2\"})\n\n assert result1[\"status\"] == \"succeeded\"\n assert result2[\"status\"] == \"succeeded\"\n assert \"req1\" in result1[\"output\"]\n assert \"req2\" in result2[\"output\"]\n\n\nclass TestAsyncGeneratorPredictor:\n \"\"\"Tests for async generator predictor.\"\"\"\n\n def test_returns_array_output(self, async_generator_predictor: Path):\n with CogletServer(async_generator_predictor) as server:\n result = server.predict({\"count\": 3, \"delay\": 0.01})\n assert result[\"status\"] == \"succeeded\"\n assert result[\"output\"] == [\n \"async chunk 0\",\n \"async chunk 1\",\n \"async chunk 2\",\n ]\n\n\n@pytest.fixture\ndef slow_sync_predictor(tmp_path: Path) -> Path:\n \"\"\"Create a sync predictor that busy-loops (cancellable at bytecode boundaries).\"\"\"\n predictor = tmp_path / \"predict.py\"\n predictor.write_text(\"\"\"\nimport time\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def setup(self):\n pass\n\n def predict(self, duration: float = 60.0) -> str:\n # Busy-loop in Python (hits bytecode boundaries, so PyThreadState_SetAsyncExc works)\n deadline = time.monotonic() + duration\n while time.monotonic() < deadline:\n pass\n return \"completed\"\n\"\"\")\n\n cog_yaml = tmp_path / \"cog.yaml\"\n cog_yaml.write_text(\"\"\"\npredict: \"predict.py:Predictor\"\n\"\"\")\n\n return predictor\n\n\n@pytest.fixture\ndef blocking_sleep_predictor(tmp_path: Path) -> Path:\n \"\"\"Create a sync predictor that blocks in time.sleep() (C-level nanosleep).\"\"\"\n predictor = tmp_path / \"predict.py\"\n predictor.write_text(\"\"\"\nimport time\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def setup(self):\n pass\n\n def predict(self, duration: float = 60.0) -> str:\n # C-level blocking sleep — PyThreadState_SetAsyncExc fires after sleep returns\n time.sleep(duration)\n return \"completed\"\n\"\"\")\n\n cog_yaml = tmp_path / \"cog.yaml\"\n cog_yaml.write_text(\"\"\"\npredict: \"predict.py:Predictor\"\n\"\"\")\n\n return predictor\n\n\n@pytest.fixture\ndef slow_async_predictor(tmp_path: Path) -> Path:\n \"\"\"Create an async predictor that sleeps for a long time (cancellable).\"\"\"\n predictor = tmp_path / \"predict.py\"\n predictor.write_text(\"\"\"\nimport asyncio\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def setup(self):\n pass\n\n async def predict(self, sleep_time: float = 60.0) -> str:\n await asyncio.sleep(sleep_time)\n return \"completed\"\n\"\"\")\n\n cog_yaml = tmp_path / \"cog.yaml\"\n cog_yaml.write_text(\"\"\"\npredict: \"predict.py:Predictor\"\n\"\"\")\n\n return predictor\n\n\ndef _wait_for_health_status(\n server: \"CogletServer\", status: str, timeout: float = 5.0\n) -> None:\n \"\"\"Poll health check until the expected status is reached, or fail.\"\"\"\n deadline = time.time() + timeout\n last_status = \"\"\n while time.time() < deadline:\n health = server.health_check()\n last_status = health[\"status\"]\n if last_status == status:\n return\n time.sleep(0.1)\n stderr = \"\".join(server.stderr_lines)\n pytest.fail(\n f\"Server did not reach status {status!r} within {timeout}s\\n\"\n f\"Last status: {last_status!r}\\n\"\n f\"STDERR:\\n{stderr}\"\n )\n\n\nclass TestCancellation:\n \"\"\"Tests for prediction cancellation.\"\"\"\n\n def test_cancel_endpoint_returns_404_for_unknown_id(self, sync_predictor: Path):\n \"\"\"Test that cancelling an unknown prediction returns 404.\"\"\"\n with CogletServer(sync_predictor) as server:\n resp = requests.post(f\"{server.base_url}/predictions/unknown-id/cancel\")\n assert resp.status_code == 404\n result = resp.json()\n assert result == {}\n\n def test_prediction_response_includes_id(self, sync_predictor: Path):\n \"\"\"Test that prediction responses include an ID.\"\"\"\n with CogletServer(sync_predictor) as server:\n result = server.predict({\"name\": \"test\"})\n assert \"id\" in result\n assert result[\"id\"].startswith(\"pred_\")\n\n def test_cancel_running_sync_prediction(self, slow_sync_predictor: Path):\n \"\"\"Test that cancelling a running sync prediction actually terminates it.\"\"\"\n with CogletServer(slow_sync_predictor) as server:\n # Start a long-running prediction asynchronously\n prediction_id = \"cancel-sync-test\"\n resp = requests.put(\n f\"{server.base_url}/predictions/{prediction_id}\",\n json={\"input\": {\"duration\": 60.0}},\n headers={\"Prefer\": \"respond-async\"},\n )\n assert resp.status_code == 202\n\n # Wait for the prediction to actually be processing (slot occupied)\n _wait_for_health_status(server, \"BUSY\", timeout=5.0)\n\n # Cancel the prediction\n cancel_resp = requests.post(\n f\"{server.base_url}/predictions/{prediction_id}/cancel\"\n )\n assert cancel_resp.status_code == 200\n\n # Wait for the server to return to READY (slot freed after cancel)\n _wait_for_health_status(server, \"READY\", timeout=10.0)\n\n def test_cancel_running_async_prediction(self, slow_async_predictor: Path):\n \"\"\"Test that cancelling a running async prediction actually terminates it.\"\"\"\n with CogletServer(slow_async_predictor) as server:\n # Start a long-running async prediction\n prediction_id = \"cancel-async-test\"\n resp = requests.put(\n f\"{server.base_url}/predictions/{prediction_id}\",\n json={\"input\": {\"sleep_time\": 60.0}},\n headers={\"Prefer\": \"respond-async\"},\n )\n assert resp.status_code == 202\n\n # Wait for the prediction to actually be processing (slot occupied)\n _wait_for_health_status(server, \"BUSY\", timeout=5.0)\n\n # Cancel the prediction\n cancel_resp = requests.post(\n f\"{server.base_url}/predictions/{prediction_id}/cancel\"\n )\n assert cancel_resp.status_code == 200\n\n # Wait for the server to return to READY (slot freed after cancel)\n _wait_for_health_status(server, \"READY\", timeout=10.0)\n\n @pytest.mark.parametrize(\n (\"predictor_fixture\", \"duration\", \"ready_timeout\"),\n [\n # Busy-loop: cancels immediately at the next bytecode boundary\n (\"slow_sync_predictor\", 60.0, 10.0),\n # time.sleep (nanosleep): blocks in C; cancel fires once sleep returns\n (\"blocking_sleep_predictor\", 5.0, 15.0),\n ],\n ids=[\"busy_loop\", \"nanosleep\"],\n )\n def test_repeated_cancel_is_idempotent(\n self,\n predictor_fixture: str,\n duration: float,\n ready_timeout: float,\n request: pytest.FixtureRequest,\n ):\n \"\"\"Test that cancelling the same prediction multiple times doesn't panic or break.\n\n Covers both busy-loop (bytecode boundaries) and time.sleep (C-level nanosleep).\n For nanosleep the cancel is deferred until the sleep returns, so we use a short\n duration and a longer timeout for the server to recover.\n \"\"\"\n predictor_path: Path = request.getfixturevalue(predictor_fixture)\n with CogletServer(predictor_path) as server:\n prediction_id = \"cancel-repeat-test\"\n resp = requests.put(\n f\"{server.base_url}/predictions/{prediction_id}\",\n json={\"input\": {\"duration\": duration}},\n headers={\"Prefer\": \"respond-async\"},\n )\n assert resp.status_code == 202\n\n # Wait for the prediction to actually be processing\n _wait_for_health_status(server, \"BUSY\", timeout=5.0)\n\n # Cancel the same prediction multiple times in rapid succession\n for i in range(5):\n cancel_resp = requests.post(\n f\"{server.base_url}/predictions/{prediction_id}/cancel\"\n )\n # First cancel returns 200 (found), subsequent may return 200 or\n # 404 depending on timing — but must never panic or 500.\n assert cancel_resp.status_code in (200, 404), (\n f\"Cancel attempt {i + 1} returned unexpected {cancel_resp.status_code}\"\n )\n\n # Server should recover to READY\n _wait_for_health_status(server, \"READY\", timeout=ready_timeout)\n\n def test_cancel_sync_prediction_connection_drop(self, slow_sync_predictor: Path):\n \"\"\"Test that dropping a sync connection cancels the prediction.\"\"\"\n with CogletServer(slow_sync_predictor) as server:\n # Start a sync (non-async) prediction with a short timeout\n # The connection drop should trigger cancellation via SyncPredictionGuard\n sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)\n sock.connect((\"localhost\", server.port))\n\n request_body = '{\"input\": {\"duration\": 60.0}}'\n http_request = (\n f\"POST /predictions HTTP/1.1\\r\\n\"\n f\"Host: localhost:{server.port}\\r\\n\"\n f\"Content-Type: application/json\\r\\n\"\n f\"Content-Length: {len(request_body)}\\r\\n\"\n f\"\\r\\n\"\n f\"{request_body}\"\n )\n sock.sendall(http_request.encode())\n\n # Wait for the prediction to be processing (slot occupied)\n _wait_for_health_status(server, \"BUSY\", timeout=5.0)\n\n # Drop the connection abruptly\n sock.close()\n\n # Wait for the server to return to READY (slot freed after cancel)\n _wait_for_health_status(server, \"READY\", timeout=10.0)\n" }, { "path": "crates/deny.toml", "content": "# cargo-deny configuration for coglet crates\n# See: https://embarkstudios.github.io/cargo-deny/\n\n[graph]\nall-features = false\nno-default-features = false\n\n[output]\nfeature-depth = 1\n\n[advisories]\nignore = [\n # Unmaintained unic-* crates from rustpython-parser (transitive dep of pyo3-stub-gen)\n # No fix available - these are only used at build time for stub generation\n \"RUSTSEC-2025-0075\", # unic-char-range\n \"RUSTSEC-2025-0080\", # unic-common\n \"RUSTSEC-2025-0081\", # unic-char-property\n \"RUSTSEC-2025-0090\", # unic-emoji-char\n \"RUSTSEC-2025-0098\", # unic-ucd-version\n \"RUSTSEC-2025-0100\", # unic-ucd-ident\n]\n\n[licenses]\n# Apache-2.0 compatible licenses only (no GPL/copyleft)\nallow = [\n \"MIT\",\n \"MIT-0\", # MIT No Attribution (more permissive than MIT)\n \"Apache-2.0\",\n \"Apache-2.0 WITH LLVM-exception\",\n \"BSD-2-Clause\",\n \"BSD-3-Clause\",\n \"ISC\",\n \"Zlib\",\n \"0BSD\",\n \"Unicode-3.0\",\n \"Unicode-DFS-2016\",\n \"CC0-1.0\",\n \"MPL-2.0\", # Weak copyleft, Apache compatible for our use\n \"CDLA-Permissive-2.0\", # Community Data License, for Mozilla CA certificate data\n]\nconfidence-threshold = 0.8\nexceptions = []\n\n[licenses.private]\nignore = false\nregistries = []\n\n[bans]\nmultiple-versions = \"warn\"\nwildcards = \"allow\"\nhighlight = \"all\"\nworkspace-default-features = \"allow\"\nexternal-default-features = \"allow\"\nallow = []\nallow-workspace = false\ndeny = []\nskip = []\nskip-tree = []\n\n[sources]\nunknown-registry = \"warn\"\nunknown-git = \"warn\"\nallow-registry = [\"https://github.com/rust-lang/crates.io-index\"]\nallow-git = []\n\n[sources.allow-org]\ngithub = []\ngitlab = []\nbitbucket = []\n" }, { "path": "docs/CNAME", "content": "cog.run\n" }, { "path": "docs/cli.md", "content": "# CLI reference\n\n\n\n## `cog`\n\nContainers for machine learning.\n\nTo get started, take a look at the documentation:\nhttps://github.com/replicate/cog\n\n**Examples**\n\n```\n To run a command inside a Docker environment defined with Cog:\n $ cog run echo hello world\n```\n\n**Options**\n\n```\n --debug Show debugging output\n -h, --help help for cog\n --no-color Disable colored output\n --version Show version of Cog\n```\n## `cog build`\n\nBuild a Docker image from the cog.yaml in the current directory.\n\nThe generated image contains your model code, dependencies, and the Cog\nruntime. It can be run locally with 'cog predict' or pushed to a registry\nwith 'cog push'.\n\n```\ncog build [flags]\n```\n\n**Examples**\n\n```\n # Build with default settings\n cog build\n\n # Build and tag the image\n cog build -t my-model:latest\n\n # Build without using the cache\n cog build --no-cache\n\n # Build with model weights in a separate layer\n cog build --separate-weights -t my-model:v1\n```\n\n**Options**\n\n```\n -f, --file string The name of the config file. (default \"cog.yaml\")\n -h, --help help for build\n --no-cache Do not use cache when building the image\n --openapi-schema string Load OpenAPI schema from a file\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n --secret stringArray Secrets to pass to the build environment in the form 'id=foo,src=/path/to/file'\n --separate-weights Separate model weights from code in image layers\n -t, --tag string A name for the built image in the form 'repository:tag'\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n```\n## `cog init`\n\nCreate a cog.yaml and predict.py in the current directory.\n\nThese files provide a starting template for defining your model's environment\nand prediction interface. Edit them to match your model's requirements.\n\n```\ncog init [flags]\n```\n\n**Examples**\n\n```\n # Set up a new Cog project in the current directory\n cog init\n```\n\n**Options**\n\n```\n -h, --help help for init\n```\n## `cog login`\n\nLog in to a container registry.\n\nFor Replicate's registry (r8.im), this command handles authentication\nthrough Replicate's token-based flow.\n\nFor other registries, this command prompts for username and password,\nthen stores credentials using Docker's credential system.\n\n```\ncog login [flags]\n```\n\n**Options**\n\n```\n -h, --help help for login\n --token-stdin Pass login token on stdin instead of opening a browser. You can find your Replicate login token at https://replicate.com/auth/token\n```\n## `cog predict`\n\nRun a prediction.\n\nIf 'image' is passed, it will run the prediction on that Docker image.\nIt must be an image that has been built by Cog.\n\nOtherwise, it will build the model in the current directory and run\nthe prediction on that.\n\n```\ncog predict [image] [flags]\n```\n\n**Examples**\n\n```\n # Run a prediction with named inputs\n cog predict -i prompt=\"a photo of a cat\"\n\n # Pass a file as input\n cog predict -i image=@photo.jpg\n\n # Save output to a file\n cog predict -i image=@input.jpg -o output.png\n\n # Pass multiple inputs\n cog predict -i prompt=\"sunset\" -i width=1024 -i height=768\n\n # Run against a pre-built image\n cog predict r8.im/your-username/my-model -i prompt=\"hello\"\n\n # Pass inputs as JSON\n echo '{\"prompt\": \"a cat\"}' | cog predict --json @-\n```\n\n**Options**\n\n```\n -e, --env stringArray Environment variables, in the form name=value\n -f, --file string The name of the config file. (default \"cog.yaml\")\n --gpus docker run --gpus GPU devices to add to the container, in the same format as docker run --gpus.\n -h, --help help for predict\n -i, --input stringArray Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg\n --json string Pass inputs as JSON object, read from file (@inputs.json) or via stdin (@-)\n -o, --output string Output path\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n --setup-timeout uint32 The timeout for a container to setup (in seconds). (default 300)\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n --use-replicate-token Pass REPLICATE_API_TOKEN from local environment into the model context\n```\n## `cog push`\n\nBuild a Docker image from cog.yaml and push it to a container registry.\n\nCog can push to any OCI-compliant registry. When pushing to Replicate's\nregistry (r8.im), run 'cog login' first to authenticate.\n\n```\ncog push [IMAGE] [flags]\n```\n\n**Examples**\n\n```\n # Push to Replicate\n cog push r8.im/your-username/my-model\n\n # Push to any OCI registry\n cog push registry.example.com/your-username/model-name\n\n # Push with model weights in a separate layer (Replicate only)\n cog push r8.im/your-username/my-model --separate-weights\n```\n\n**Options**\n\n```\n -f, --file string The name of the config file. (default \"cog.yaml\")\n -h, --help help for push\n --no-cache Do not use cache when building the image\n --openapi-schema string Load OpenAPI schema from a file\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n --secret stringArray Secrets to pass to the build environment in the form 'id=foo,src=/path/to/file'\n --separate-weights Separate model weights from code in image layers\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n```\n## `cog run`\n\nRun a command inside a Docker environment defined by cog.yaml.\n\nCog builds a temporary image from your cog.yaml configuration and runs the\ngiven command inside it. This is useful for debugging, running scripts, or\nexploring the environment your model will run in.\n\n```\ncog run [arg...] [flags]\n```\n\n**Examples**\n\n```\n # Open a Python interpreter inside the model environment\n cog run python\n\n # Run a script\n cog run python train.py\n\n # Run with environment variables\n cog run -e HUGGING_FACE_HUB_TOKEN=abc123 python download.py\n\n # Expose a port (e.g. for Jupyter)\n cog run -p 8888 jupyter notebook\n```\n\n**Options**\n\n```\n -e, --env stringArray Environment variables, in the form name=value\n -f, --file string The name of the config file. (default \"cog.yaml\")\n --gpus docker run --gpus GPU devices to add to the container, in the same format as docker run --gpus.\n -h, --help help for run\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n -p, --publish stringArray Publish a container's port to the host, e.g. -p 8000\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n```\n## `cog serve`\n\nRun a prediction HTTP server.\n\nBuilds the model and starts an HTTP server that exposes the model's inputs\nand outputs as a REST API. Compatible with the Cog HTTP protocol.\n\n```\ncog serve [flags]\n```\n\n**Examples**\n\n```\n # Start the server on the default port (8393)\n cog serve\n\n # Start on a custom port\n cog serve -p 5000\n\n # Test the server\n curl http://localhost:8393/predictions \\\n -X POST \\\n -H 'Content-Type: application/json' \\\n -d '{\"input\": {\"prompt\": \"a cat\"}}'\n```\n\n**Options**\n\n```\n -f, --file string The name of the config file. (default \"cog.yaml\")\n --gpus docker run --gpus GPU devices to add to the container, in the same format as docker run --gpus.\n -h, --help help for serve\n -p, --port int Port on which to listen (default 8393)\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n --upload-url string Upload URL for file outputs (e.g. https://example.com/upload/)\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n```\n" }, { "path": "docs/deploy.md", "content": "# Deploy models with Cog\n\nCog containers are Docker containers that serve an HTTP server\nfor running predictions on your model.\nYou can deploy them anywhere that Docker containers run.\n\nThe server inside Cog containers is **coglet**, a Rust-based prediction server\nthat handles HTTP requests, worker process management, and prediction execution.\n\nThis guide assumes you have a model packaged with Cog.\nIf you don't, [follow our getting started guide](getting-started-own-model.md),\nor use [an example model](https://github.com/replicate/cog-examples).\n\n## Getting started\n\nFirst, build your model:\n\n```console\ncog build -t my-model\n```\n\nYou can serve predictions locally with `cog serve`:\n\n```console\ncog serve\n# or, from a built image:\ncog serve my-model\n```\n\nAlternatively, start the Docker container directly:\n\n```shell\n# If your model uses a CPU:\ndocker run -d -p 5001:5000 my-model\n\n# If your model uses a GPU:\ndocker run -d -p 5001:5000 --gpus all my-model\n```\n\nThe server listens on port 5000 inside the container (mapped to 5001 above).\n\nTo view the OpenAPI schema,\nopen [localhost:5001/openapi.json](http://localhost:5001/openapi.json)\nin your browser\nor use cURL to make a request:\n\n```console\ncurl http://localhost:5001/openapi.json\n```\n\nTo stop the server, run:\n\n```console\ndocker kill my-model\n```\n\nTo run a prediction on the model,\ncall the `/predictions` endpoint,\npassing input in the format expected by your model:\n\n```console\ncurl http://localhost:5001/predictions -X POST \\\n --header \"Content-Type: application/json\" \\\n --data '{\"input\": {\"image\": \"https://.../input.jpg\"}}'\n```\n\nFor more details about the HTTP API,\nsee the [HTTP API reference documentation](http.md).\n\n## Health checks\n\nThe server exposes a `GET /health-check` endpoint that returns the current status of the model container. Use this for readiness probes in orchestration systems like Kubernetes.\n\n```console\ncurl http://localhost:5001/health-check\n```\n\nThe response includes a `status` field with values like `STARTING`, `READY`, `BUSY`, `SETUP_FAILED`, or `DEFUNCT`. See the [HTTP API reference](http.md#get-health-check) for full details.\n\n## Concurrency\n\nBy default, the server processes one prediction at a time. To enable concurrent predictions, set the `concurrency.max` option in `cog.yaml`:\n\n```yaml\nconcurrency:\n max: 4\n```\n\nSee the [`cog.yaml` reference](yaml.md#concurrency) for more details.\n\n## Environment variables\n\nYou can configure runtime behavior with environment variables:\n\n- `COG_SETUP_TIMEOUT`: Maximum time in seconds for the `setup()` method (default: no timeout).\n\nSee the [environment variables reference](environment.md) for the full list.\n" }, { "path": "docs/environment.md", "content": "# Environment variables\n\nThis guide lists the environment variables that change how Cog functions.\n\n## Build-time variables\n\n### `COG_SDK_WHEEL`\n\nControls which cog Python SDK wheel is installed in the Docker image during `cog build`. Takes precedence over `build.sdk_version` in `cog.yaml`.\n\n**Supported values:**\n\n| Value | Description |\n| -------------------- | ---------------------------------------------------- |\n| `pypi` | Install latest version from PyPI |\n| `pypi:0.12.0` | Install specific version from PyPI |\n| `dist` | Use wheel from `dist/` directory (requires git repo) |\n| `https://...` | Install from URL |\n| `/path/to/wheel.whl` | Install from local file path |\n\n**Default behavior:**\n\n- **Release builds**: Installs latest cog from PyPI\n- **Development builds**: Auto-detects wheel in `dist/` directory, falls back to latest PyPI\n\n**Examples:**\n\n```console\n# Use specific PyPI version\n$ COG_SDK_WHEEL=pypi:0.11.0 cog build\n\n# Use local development wheel\n$ COG_SDK_WHEEL=dist cog build\n\n# Use wheel from URL\n$ COG_SDK_WHEEL=https://example.com/cog-0.12.0-py3-none-any.whl cog build\n```\n\nThe `dist` option searches for wheels in:\n\n1. `./dist/` (current directory)\n2. `$REPO_ROOT/dist/` (if REPO_ROOT is set)\n3. `/dist/` (via `git rev-parse`, useful when running from subdirectories)\n\n### `COGLET_WHEEL`\n\nControls which coglet wheel is installed in the Docker image. Coglet is the Rust-based prediction server.\n\n**Supported values:** Same as `COG_SDK_WHEEL`\n\n**Default behavior:** For development builds, auto-detects a wheel in `dist/`. For release builds, installs the latest version from PyPI. Can be overridden with an explicit value.\n\n**Examples:**\n\n```console\n# Use local development wheel\n$ COGLET_WHEEL=dist cog build\n\n# Use specific version from PyPI\n$ COGLET_WHEEL=pypi:0.1.0 cog build\n```\n\n## Runtime variables\n\n### `COG_NO_UPDATE_CHECK`\n\nBy default, Cog automatically checks for updates\nand notifies you if there is a new version available.\n\nTo disable this behavior,\nset the `COG_NO_UPDATE_CHECK` environment variable to any value.\n\n```console\n$ COG_NO_UPDATE_CHECK=1 cog build # runs without automatic update check\n```\n\n### `COG_SETUP_TIMEOUT`\n\nControls the maximum time (in seconds) allowed for the model's `setup()` method to complete. If setup exceeds this timeout, the server will report a setup failure.\n\nBy default, there is no timeout — setup runs indefinitely.\n\nSet to `0` to disable the timeout (same as default). Invalid values are ignored with a warning.\n\n```console\n$ COG_SETUP_TIMEOUT=300 docker run -p 5000:5000 my-model # 5-minute setup timeout\n```\n\n### `COG_CA_CERT`\n\nInjects a custom CA certificate into the Docker image during `cog build`. This is useful when building behind a corporate proxy or VPN that uses custom certificate authorities (e.g. Cloudflare WARP).\n\n**Supported values:**\n\n| Value | Description |\n| -------------------------------- | ----------------------------------------------------------- |\n| `/path/to/cert.crt` | Path to a single PEM certificate file |\n| `/path/to/certs/` | Directory of `.crt` and `.pem` files (all are concatenated) |\n| `-----BEGIN CERTIFICATE-----...` | Inline PEM certificate |\n| `LS0tLS1CRUdJTi...` | Base64-encoded PEM certificate |\n\nThe certificate is installed into the system CA store and the `SSL_CERT_FILE` and `REQUESTS_CA_BUNDLE` environment variables are set automatically in the built image.\n\n**Examples:**\n\n```console\n# From a file\n$ COG_CA_CERT=/usr/local/share/ca-certificates/corporate-ca.crt cog build\n\n# From a directory of certs\n$ COG_CA_CERT=/etc/custom-certs/ cog build\n\n# Inline (e.g. from a CI secret)\n$ COG_CA_CERT=\"$(cat /path/to/cert.pem)\" cog build\n```\n" }, { "path": "docs/getting-started-own-model.md", "content": "# Getting started with your own model\n\nThis guide will show you how to put your own machine learning model in a Docker image using Cog. If you haven't got a model to try out, you'll want to follow the [main getting started guide](getting-started.md).\n\n## Prerequisites\n\n- **macOS or Linux**. Cog works on macOS and Linux, but does not currently support Windows.\n- **Docker**. Cog uses Docker to create a container for your model. You'll need to [install Docker](https://docs.docker.com/get-docker/) before you can run Cog.\n\n## Initialization\n\nFirst, install Cog if you haven't already:\n\n**macOS (recommended):**\n\n```sh\nbrew install replicate/tap/cog\n```\n\n**Linux or macOS (manual):**\n\n```sh\nsudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`\nsudo chmod +x /usr/local/bin/cog\n```\n\nTo configure your project for use with Cog, you'll need to add two files:\n\n- [`cog.yaml`](yaml.md) defines system requirements, Python package dependencies, etc\n- [`predict.py`](python.md) describes the prediction interface for your model\n\nUse the `cog init` command to generate these files in your project:\n\n```sh\n$ cd path/to/your/model\n$ cog init\n```\n\n## Define the Docker environment\n\nThe `cog.yaml` file defines all the different things that need to be installed for your model to run. You can think of it as a simple way of defining a Docker image.\n\nFor example:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n python_requirements: requirements.txt\n```\n\nWith a `requirements.txt` containing your dependencies:\n\n```\ntorch==2.6.0\n```\n\nThis will generate a Docker image with Python 3.13 and PyTorch 2 installed, for both CPU and GPU, with the correct version of CUDA, and various other sensible best-practices.\n\nTo run a command inside this environment, prefix it with `cog run`:\n\n```\n$ cog run python\n✓ Building Docker image from cog.yaml... Successfully built 8f54020c8981\nRunning 'python' in Docker with the current directory mounted as a volume...\n────────────────────────────────────────────────────────────────────────────────────────\n\nPython 3.13.x (main, ...)\n[GCC 12.2.0] on linux\nType \"help\", \"copyright\", \"credits\" or \"license\" for more information.\n>>>\n```\n\nThis is handy for ensuring a consistent environment for development or training.\n\nWith `cog.yaml`, you can also install system packages and other things. [Take a look at the full reference to see what else you can do.](yaml.md)\n\n## Define how to run predictions\n\nThe next step is to update `predict.py` to define the interface for running predictions on your model. The `predict.py` generated by `cog init` looks something like this:\n\n```python\nfrom cog import BasePredictor, Path, Input\nimport torch\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n self.net = torch.load(\"weights.pth\")\n\n def predict(self,\n image: Path = Input(description=\"Image to enlarge\"),\n scale: float = Input(description=\"Factor to scale image by\", default=1.5)\n ) -> Path:\n \"\"\"Run a single prediction on the model\"\"\"\n # ... pre-processing ...\n output = self.net(input)\n # ... post-processing ...\n return output\n```\n\nEdit your `predict.py` file and fill in the functions with your own model's setup and prediction code. You might need to import parts of your model from another file.\n\nYou also need to define the inputs to your model as arguments to the `predict()` function, as demonstrated above. For each argument, you need to annotate with a type. The supported types are:\n\n- `str`: a string\n- `int`: an integer\n- `float`: a floating point number\n- `bool`: a boolean\n- `cog.File`: a file-like object representing a file (deprecated — use `cog.Path` instead)\n- `cog.Path`: a path to a file on disk\n\nYou can provide more information about the input with the `Input()` function, as shown above. It takes these basic arguments:\n\n- `description`: A description of what to pass to this input for users of the model\n- `default`: A default value to set the input to. If this argument is not passed, the input is required. If it is explicitly set to `None`, the input is optional.\n- `ge`: For `int` or `float` types, the value should be greater than or equal to this number.\n- `le`: For `int` or `float` types, the value should be less than or equal to this number.\n- `min_length`: For `str` types, the minimum length of the string.\n- `max_length`: For `str` types, the maximum length of the string.\n- `regex`: For `str` types, the string must match this regular expression.\n- `choices`: For `str` or `int` types, a list of possible values for this input.\n- `deprecated`: Mark this input as deprecated with a message explaining what to use instead.\n\nThere are some more advanced options you can pass, too. For more details, [take a look at the prediction interface documentation](python.md).\n\nNext, add the line `predict: \"predict.py:Predictor\"` to your `cog.yaml`, so it looks something like this:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n python_requirements: requirements.txt\npredict: \"predict.py:Predictor\"\n```\n\nThat's it! To test this works, try running a prediction on the model:\n\n```\n$ cog predict -i image=@input.jpg\n✓ Building Docker image from cog.yaml... Successfully built 664ef88bc1f4\n✓ Model running in Docker image 664ef88bc1f4\n\nWritten output to output.png\n```\n\nTo pass more inputs to the model, you can add more `-i` options:\n\n```\n$ cog predict -i image=@image.jpg -i scale=2.0\n```\n\nIn this case it is just a number, not a file, so you don't need the `@` prefix.\n\n## Using GPUs\n\nTo use GPUs with Cog, add the `gpu: true` option to the `build` section of your `cog.yaml`:\n\n```yaml\nbuild:\n gpu: true\n ...\n```\n\nCog will use the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image and automatically figure out what versions of CUDA and cuDNN to use based on the version of Python, PyTorch, and Tensorflow that you are using.\n\nFor more details, [see the `gpu` section of the `cog.yaml` reference](yaml.md#gpu).\n\n## Next steps\n\nNext, you might want to take a look at:\n\n- [A guide explaining how to deploy a model.](deploy.md)\n- [The reference for `cog.yaml`](yaml.md)\n- [The reference for the Python library](python.md)\n" }, { "path": "docs/getting-started.md", "content": "# Getting started\n\nThis guide will walk you through what you can do with Cog by using an example model.\n\n> [!TIP]\n> Using a language model to help you write the code for your new Cog model?\n>\n> Feed it [https://cog.run/llms.txt](https://cog.run/llms.txt), which has all of Cog's documentation bundled into a single file. To learn more about this format, check out [llmstxt.org](https://llmstxt.org).\n\n## Prerequisites\n\n- **macOS or Linux**. Cog works on macOS and Linux, but does not currently support Windows.\n- **Docker**. Cog uses Docker to create a container for your model. You'll need to [install Docker](https://docs.docker.com/get-docker/) before you can run Cog.\n\n## Install Cog\n\n**macOS (recommended):**\n\n```bash\nbrew install replicate/tap/cog\n```\n\n**Linux or macOS (manual):**\n\n```bash\nsudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`\nsudo chmod +x /usr/local/bin/cog\nsudo xattr -d com.apple.quarantine /usr/local/bin/cog 2>/dev/null || true\n\n```\n\n> [!NOTE]\n> **macOS: \"cannot be opened because the developer cannot be verified\"**\n>\n> If you downloaded the binary manually (via `curl` or a browser) and see this Gatekeeper warning, run:\n>\n> ```bash\n> sudo xattr -d com.apple.quarantine /usr/local/bin/cog\n> ```\n>\n> Installing via `brew install replicate/tap/cog` handles this automatically.\n\n## Create a project\n\nLet's make a directory to work in:\n\n```bash\nmkdir cog-quickstart\ncd cog-quickstart\n\n```\n\n## Run commands\n\nThe simplest thing you can do with Cog is run a command inside a Docker environment.\n\nThe first thing you need to do is create a file called `cog.yaml`:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n```\n\nThen, you can run any command inside this environment. For example, enter\n\n```bash\ncog run python\n\n```\n\nand you'll get an interactive Python shell:\n\n```none\n✓ Building Docker image from cog.yaml... Successfully built 8f54020c8981\nRunning 'python' in Docker with the current directory mounted as a volume...\n───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n\nPython 3.13.x (main, ...)\n[GCC 12.2.0] on linux\nType \"help\", \"copyright\", \"credits\" or \"license\" for more information.\n>>>\n```\n\n(Hit Ctrl-D to exit the Python shell.)\n\nInside this Docker environment you can do anything – run a Jupyter notebook, your training script, your evaluation script, and so on.\n\n## Run predictions on a model\n\nLet's pretend we've trained a model. With Cog, we can define how to run predictions on it in a standard way, so other people can easily run predictions on it without having to hunt around for a prediction script.\n\nWe need to write some code to describe how predictions are run on the model.\n\nSave this to `predict.py`:\n\n```python\nimport os\nos.environ[\"TORCH_HOME\"] = \".\"\n\nimport torch\nfrom cog import BasePredictor, Input, Path\nfrom PIL import Image\nfrom torchvision import models\n\nWEIGHTS = models.ResNet50_Weights.IMAGENET1K_V1\n\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n self.model = models.resnet50(weights=WEIGHTS).to(self.device)\n self.model.eval()\n\n def predict(self, image: Path = Input(description=\"Image to classify\")) -> dict:\n \"\"\"Run a single prediction on the model\"\"\"\n img = Image.open(image).convert(\"RGB\")\n preds = self.model(WEIGHTS.transforms()(img).unsqueeze(0).to(self.device))\n top3 = preds[0].softmax(0).topk(3)\n categories = WEIGHTS.meta[\"categories\"]\n return {categories[i]: p.detach().item() for p, i in zip(*top3)}\n```\n\nWe also need to point Cog at this, and tell it what Python dependencies to install.\n\nSave this to `requirements.txt`:\n\n```\npillow==11.1.0\ntorch==2.6.0\ntorchvision==0.21.0\n```\n\nThen update `cog.yaml` to look like this:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n python_requirements: requirements.txt\npredict: \"predict.py:Predictor\"\n```\n\n> [!TIP]\n> If you have a machine with an NVIDIA GPU attached, add `gpu: true` to the `build` section of your `cog.yaml` to enable GPU acceleration.\n\nLet's grab an image to test the model with:\n\n```bash\nIMAGE_URL=https://gist.githubusercontent.com/bfirsh/3c2115692682ae260932a67d93fd94a8/raw/56b19f53f7643bb6c0b822c410c366c3a6244de2/mystery.jpg\ncurl $IMAGE_URL > input.jpg\n\n```\n\nNow, let's run the model using Cog:\n\n```bash\ncog predict -i image=@input.jpg\n\n```\n\nIf you see the following output\n\n```json\n{\n \"tiger_cat\": 0.4874822497367859,\n \"tabby\": 0.23169134557247162,\n \"Egyptian_cat\": 0.09728282690048218\n}\n```\n\nthen it worked!\n\nNote: The first time you run `cog predict`, the build process will be triggered to generate a Docker container that can run your model. The next time you run `cog predict` the pre-built container will be used.\n\n## Build an image\n\nWe can bake your model's code, the trained weights, and the Docker environment into a Docker image. This image serves predictions with an HTTP server, and can be deployed to anywhere that Docker runs to serve real-time predictions.\n\n```bash\ncog build -t resnet\n# Building Docker image...\n# Built resnet:latest\n\n```\n\nYou can run this image with `cog predict` by passing the filename as an argument:\n\n```bash\ncog predict resnet -i image=@input.jpg\n\n```\n\nOr, you can run it with Docker directly, and it'll serve an HTTP server:\n\n```bash\ndocker run -d --rm -p 5000:5000 resnet\n\n```\n\nWe can send inputs directly with `curl`:\n\n```bash\ncurl http://localhost:5000/predictions -X POST \\\n -H 'Content-Type: application/json' \\\n -d '{\"input\": {\"image\": \"https://gist.githubusercontent.com/bfirsh/3c2115692682ae260932a67d93fd94a8/raw/56b19f53f7643bb6c0b822c410c366c3a6244de2/mystery.jpg\"}}'\n\n```\n\nAs a shorthand, you can add the Docker image's name as an extra line in `cog.yaml`:\n\n```yaml\nimage: \"r8.im/replicate/resnet\"\n```\n\nOnce you've done this, you can use `cog push` to build and push the image to a Docker registry:\n\n```bash\ncog push\n# Building r8.im/replicate/resnet...\n# Pushing r8.im/replicate/resnet...\n# Pushed!\n```\n\nThe Docker image is now accessible to anyone or any system that has access to this Docker registry.\n\n## Next steps\n\nThose are the basics! Next, you might want to take a look at:\n\n- [A guide to help you set up your own model on Cog.](getting-started-own-model.md)\n- [A guide explaining how to deploy a model.](deploy.md)\n- [Reference for `cog.yaml`](yaml.md)\n- [Reference for the Python library](python.md)\n" }, { "path": "docs/http.md", "content": "# HTTP API\n\n> [!TIP]\n> For information about how to run the HTTP server,\n> see [our documentation on deploying models](deploy.md).\n\nWhen you run a Docker image built by Cog,\nit serves an HTTP API for making predictions.\n\nThe server supports both synchronous and asynchronous prediction creation:\n\n- **Synchronous**:\n The server waits until the prediction is completed\n and responds with the result.\n- **Asynchronous**:\n The server immediately returns a response\n and processes the prediction in the background.\n\nThe client can create a prediction asynchronously\nby setting the `Prefer: respond-async` header in their request.\nWhen provided, the server responds immediately after starting the prediction\nwith `202 Accepted` status and a prediction object in status `processing`.\n\n> [!NOTE]\n> The only supported way to receive updates on the status of predictions\n> started asynchronously is using [webhooks](#webhooks).\n> Polling for prediction status is not currently supported.\n\nYou can also use certain server endpoints to create predictions idempotently,\nsuch that if a client calls this endpoint more than once with the same ID\n(for example, due to a network interruption)\nwhile the prediction is still running,\nno new prediction is created.\nInstead, the client receives a `202 Accepted` response\nwith the initial state of the prediction.\n\n---\n\nHere's a summary of the prediction creation endpoints:\n\n| Endpoint | Header | Behavior |\n| ---------------------------------- | ----------------------- | ---------------------------- |\n| `POST /predictions` | - | Synchronous, non-idempotent |\n| `POST /predictions` | `Prefer: respond-async` | Asynchronous, non-idempotent |\n| `PUT /predictions/` | - | Synchronous, idempotent |\n| `PUT /predictions/` | `Prefer: respond-async` | Asynchronous, idempotent |\n\nChoose the endpoint that best fits your needs:\n\n- Use synchronous endpoints when you want to wait for the prediction result.\n- Use asynchronous endpoints when you want to start a prediction\n and receive updates via webhooks.\n- Use idempotent endpoints when you need to safely retry requests\n without creating duplicate predictions.\n\n## Webhooks\n\nYou can provide a `webhook` parameter in the client request body\nwhen creating a prediction.\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n \"webhook\": \"https://example.com/webhook/prediction\"\n}\n```\n\nThe server makes requests to the provided URL\nwith the current state of the prediction object in the request body\nat the following times.\n\n- `start`:\n Once, when the prediction starts\n (`status` is `starting`).\n- `output`:\n Each time a predict function generates an output\n (either once using `return` or multiple times using `yield`)\n- `logs`:\n Each time the predict function writes to `stdout`\n- `completed`:\n Once, when the prediction reaches a terminal state\n (`status` is `succeeded`, `canceled`, or `failed`)\n\nWebhook requests for `start` and `completed` event types\nare sent immediately.\nWebhook requests for `output` and `logs` event types\nare sent at most once every 500ms.\nThis interval is not configurable.\n\nBy default, the server sends requests for all event types.\nClients can specify which events trigger webhook requests\nwith the `webhook_events_filter` parameter in the prediction request body.\nFor example,\nthe following request specifies that webhooks are sent by the server\nonly at the start and end of the prediction:\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n \"webhook\": \"https://example.com/webhook/prediction\",\n \"webhook_events_filter\": [\"start\", \"completed\"]\n}\n```\n\n## Generating unique prediction IDs\n\nEndpoints for creating and canceling a prediction idempotently\naccept a `prediction_id` parameter in their path.\nBy default, the server runs one prediction at a time,\nbut this can be increased with the [`concurrency.max`](yaml.md#concurrency) setting.\nWhen all prediction slots are in use, the server returns `409 Conflict`.\nThe client should ensure prediction slots are available\nbefore creating a new prediction with a different ID.\n\nClients are responsible for providing unique prediction IDs.\nWe recommend generating a UUIDv4 or [UUIDv7](https://uuid7.com),\nbase32-encoding that value,\nand removing padding characters (`==`).\nThis produces a random identifier that is 26 ASCII characters long.\n\n```python\n>> from uuid import uuid4\n>> from base64 import b32encode\n>> b32encode(uuid4().bytes).decode('utf-8').lower().rstrip('=')\n'wjx3whax6rf4vphkegkhcvpv6a'\n```\n\n## File uploads\n\nA model's `predict` function can produce file output by yielding or returning\na `cog.Path` or `cog.File` value.\n\nBy default,\nfiles are returned as a base64-encoded\n[data URL](https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs).\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n}\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"succeeded\",\n \"output\": \"data:image/png;base64,...\"\n}\n```\n\nWhen creating a prediction synchronously,\nthe client can configure a base URL to upload output files to instead\nby setting the `output_file_prefix` parameter in the request body:\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n \"output_file_prefix\": \"https://example.com/upload\",\n}\n```\n\nWhen the model produces a file output,\nthe server sends the following request to upload the file to the configured URL:\n\n```http\nPUT /upload HTTP/1.1\nHost: example.com\nContent-Type: multipart/form-data\n\n--boundary\nContent-Disposition: form-data; name=\"file\"; filename=\"image.png\"\nContent-Type: image/png\n\n\n--boundary--\n```\n\nIf the upload succeeds, the server responds with output:\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"succeeded\",\n \"output\": \"http://example.com/upload/image.png\"\n}\n```\n\nIf the upload fails, the server responds with an error.\n\n> [!IMPORTANT] \n> File uploads for predictions created asynchronously\n> require `--upload-url` to be specified when starting the HTTP server.\n\n\n\n## Endpoints\n\n### `GET /`\n\nReturns a discovery document listing available API endpoints, the OpenAPI schema URL, and version information.\n\n```http\nGET / HTTP/1.1\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"cog_version\": \"0.17.0\",\n \"docs_url\": \"/openapi.json\",\n \"openapi_url\": \"/openapi.json\",\n \"predictions_url\": \"/predictions\",\n \"health_check_url\": \"/health-check\"\n}\n```\n\nIf training is configured, the response also includes a `trainings_url` field.\n\n### `GET /health-check`\n\nReturns the current health status of the model container.\nThis endpoint always responds with `200 OK` —\ncheck the `status` field in the response body to determine readiness.\n\nThe response body is a JSON object with the following fields:\n\n- `status`: One of the following values:\n - `STARTING`: The model's `setup()` method is still running.\n - `READY`: The model is ready to accept predictions.\n - `BUSY`: The model is ready but all prediction slots are in use.\n - `SETUP_FAILED`: The model's `setup()` method raised an exception.\n - `DEFUNCT`: The model encountered an unrecoverable error.\n - `UNHEALTHY`: The model is ready\n but a user-defined `healthcheck()` method returned `False`.\n- `setup`: Setup phase details (included once setup has started):\n - `started_at`: ISO 8601 timestamp of when setup began.\n - `completed_at`: ISO 8601 timestamp of when setup finished (if complete).\n - `status`: One of `starting`, `succeeded`, or `failed`.\n - `logs`: Output captured during setup.\n- `version`: Runtime version information:\n - `coglet`: Coglet version.\n - `cog`: Cog Python SDK version (if available).\n - `python`: Python version (if available).\n- `user_healthcheck_error`:\n Error message from a user-defined `healthcheck()` method (if applicable).\n\n```http\nGET /health-check HTTP/1.1\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"READY\",\n \"setup\": {\n \"started_at\": \"2025-01-01T00:00:00.000000+00:00\",\n \"completed_at\": \"2025-01-01T00:00:05.000000+00:00\",\n \"status\": \"succeeded\",\n \"logs\": \"\"\n },\n \"version\": {\n \"coglet\": \"0.17.0\",\n \"cog\": \"0.14.0\",\n \"python\": \"3.13.0\"\n }\n}\n```\n\n### `GET /openapi.json`\n\nThe [OpenAPI](https://swagger.io/specification/) specification of the API,\nwhich is derived from the input and output types specified in your model's\n[Predictor](python.md) and [Training](training.md) objects.\n\n### `POST /predictions`\n\nMakes a single prediction.\n\nThe request body is a JSON object with the following fields:\n\n- `input`:\n A JSON object with the same keys as the\n [arguments to the `predict()` function](python.md).\n Any `File` or `Path` inputs are passed as URLs.\n\nThe response body is a JSON object with the following fields:\n\n- `status`: Either `succeeded` or `failed`.\n- `output`: The return value of the `predict()` function.\n- `error`: If `status` is `failed`, the error message.\n- `metrics`: An object containing prediction metrics.\n Always includes `predict_time` (elapsed seconds).\n May also include custom metrics recorded by the model\n using [`self.record_metric()`](python.md#metrics).\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\n\n{\n \"input\": {\n \"image\": \"https://example.com/image.jpg\",\n \"text\": \"Hello world!\"\n }\n}\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"succeeded\",\n \"output\": \"data:image/png;base64,...\",\n \"metrics\": {\n \"predict_time\": 4.52\n }\n}\n```\n\nIf the client sets the `Prefer: respond-async` header in their request,\nthe server responds immediately after starting the prediction\nwith `202 Accepted` status and a prediction object in status `processing`.\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"}\n}\n```\n\n```http\nHTTP/1.1 202 Accepted\nContent-Type: application/json\n\n{\n \"status\": \"starting\",\n}\n```\n\n### `PUT /predictions/`\n\nMake a single prediction.\nThis is the idempotent version of the `POST /predictions` endpoint.\n\n```http\nPUT /predictions/wjx3whax6rf4vphkegkhcvpv6a HTTP/1.1\nContent-Type: application/json; charset=utf-8\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"}\n}\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"succeeded\",\n \"output\": \"data:image/png;base64,...\"\n}\n```\n\nIf the client sets the `Prefer: respond-async` header in their request,\nthe server responds immediately after starting the prediction\nwith `202 Accepted` status and a prediction object in status `processing`.\n\n```http\nPUT /predictions/wjx3whax6rf4vphkegkhcvpv6a HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"}\n}\n```\n\n```http\nHTTP/1.1 202 Accepted\nContent-Type: application/json\n\n{\n \"id\": \"wjx3whax6rf4vphkegkhcvpv6a\",\n \"status\": \"starting\"\n}\n```\n\n### `POST /predictions//cancel`\n\nA client can cancel an asynchronous prediction by making a\n`POST /predictions//cancel` request\nusing the prediction `id` provided when the prediction was created.\n\nFor example,\nif the client creates a prediction by sending the request:\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"id\": \"abcd1234\",\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n}\n```\n\nThe client can cancel the prediction by sending the request:\n\n```http\nPOST /predictions/abcd1234/cancel HTTP/1.1\n```\n\nA prediction cannot be canceled if it's\ncreated synchronously, without the `Prefer: respond-async` header,\nor created without a provided `id`.\n\nIf a prediction exists with the provided `id`,\nthe server responds with status `200 OK`.\nOtherwise, the server responds with status `404 Not Found`.\n\nWhen a prediction is canceled,\nCog raises [`CancelationException`](python.md#cancelationexception)\nin sync predictors (or `asyncio.CancelledError` in async predictors).\nThis exception may be caught by the model to perform necessary cleanup.\nThe cleanup should be brief, ideally completing within a few seconds.\nAfter cleanup, the exception must be re-raised using a bare `raise` statement.\nFailure to re-raise the exception may result in the termination of the container.\n\n```python\nfrom cog import BasePredictor, CancelationException, Input, Path\n\nclass Predictor(BasePredictor):\n def predict(self, image: Path = Input(description=\"Image to process\")) -> Path:\n try:\n return self.process(image)\n except CancelationException:\n self.cleanup()\n raise # always re-raise\n```\n" }, { "path": "docs/llms.txt", "content": "# Cog: Containers for machine learning\n\nCog is an open-source tool that lets you package machine learning models in a standard, production-ready container.\n\nYou can deploy your packaged model to your own infrastructure, or to [Replicate](https://replicate.com/).\n\n## Highlights\n\n- 📦 **Docker containers without the pain.** Writing your own `Dockerfile` can be a bewildering process. With Cog, you define your environment with a [simple configuration file](#how-it-works) and it generates a Docker image with all the best practices: Nvidia base images, efficient caching of dependencies, installing specific Python versions, sensible environment variable defaults, and so on.\n\n- 🤬️ **No more CUDA hell.** Cog knows which CUDA/cuDNN/PyTorch/Tensorflow/Python combos are compatible and will set it all up correctly for you.\n\n- ✅ **Define the inputs and outputs for your model with standard Python.** Then, Cog generates an OpenAPI schema and validates the inputs and outputs.\n\n- 🎁 **Automatic HTTP prediction server**: Your model's types are used to dynamically generate a RESTful HTTP API using a high-performance Rust/Axum server.\n\n- 🚀 **Ready for production.** Deploy your model anywhere that Docker images run. Your own infrastructure, or [Replicate](https://replicate.com).\n\n## How it works\n\nDefine the Docker environment your model runs in with `cog.yaml`:\n\n```yaml\nbuild:\n gpu: true\n system_packages:\n - \"libgl1-mesa-glx\"\n - \"libglib2.0-0\"\n python_version: \"3.13\"\n python_requirements: requirements.txt\npredict: \"predict.py:Predictor\"\n```\n\nDefine how predictions are run on your model with `predict.py`:\n\n```python\nfrom cog import BasePredictor, Input, Path\nimport torch\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n self.model = torch.load(\"./weights.pth\")\n\n # The arguments and types the model takes as input\n def predict(self,\n image: Path = Input(description=\"Grayscale input image\")\n ) -> Path:\n \"\"\"Run a single prediction on the model\"\"\"\n processed_image = preprocess(image)\n output = self.model(processed_image)\n return postprocess(output)\n```\n\nIn the above we accept a path to the image as an input, and return a path to our transformed image after running it through our model.\n\nNow, you can run predictions on this model:\n\n```console\n$ cog predict -i image=@input.jpg\n--> Building Docker image...\n--> Running Prediction...\n--> Output written to output.jpg\n```\n\nOr, build a Docker image for deployment:\n\n```console\n$ cog build -t my-classification-model\n--> Building Docker image...\n--> Built my-classification-model:latest\n\n$ docker run -d -p 5000:5000 --gpus all my-classification-model\n\n$ curl http://localhost:5000/predictions -X POST \\\n -H 'Content-Type: application/json' \\\n -d '{\"input\": {\"image\": \"https://.../input.jpg\"}}'\n```\n\nOr, combine build and run via the `serve` command:\n\n```console\n$ cog serve -p 8080\n\n$ curl http://localhost:8080/predictions -X POST \\\n -H 'Content-Type: application/json' \\\n -d '{\"input\": {\"image\": \"https://.../input.jpg\"}}'\n```\n\n\n\n## Why are we building this?\n\nIt's really hard for researchers to ship machine learning models to production.\n\nPart of the solution is Docker, but it is so complex to get it to work: Dockerfiles, pre-/post-processing, Flask servers, CUDA versions. More often than not the researcher has to sit down with an engineer to get the damn thing deployed.\n\n[Andreas](https://github.com/andreasjansson) and [Ben](https://github.com/bfirsh) created Cog. Andreas used to work at Spotify, where he built tools for building and deploying ML models with Docker. Ben worked at Docker, where he created [Docker Compose](https://github.com/docker/compose).\n\nWe realized that, in addition to Spotify, other companies were also using Docker to build and deploy machine learning models. [Uber](https://eng.uber.com/michelangelo-pyml/) and others have built similar systems. So, we're making an open source version so other people can do this too.\n\nHit us up if you're interested in using it or want to collaborate with us. [We're on Discord](https://discord.gg/replicate) or email us at [team@replicate.com](mailto:team@replicate.com).\n\n## Prerequisites\n\n- **macOS, Linux or Windows 11**. Cog works on macOS, Linux and Windows 11 with [WSL 2](docs/wsl2/wsl2.md)\n- **Docker**. Cog uses Docker to create a container for your model. You'll need to [install Docker](https://docs.docker.com/get-docker/) before you can run Cog. If you install Docker Engine instead of Docker Desktop, you will need to [install Buildx](https://docs.docker.com/build/architecture/#buildx) as well.\n\n## Install\n\nIf you're using macOS, you can install Cog using Homebrew:\n\n```console\nbrew install replicate/tap/cog\n```\n\nYou can also download and install the latest release using our\n[install script](https://cog.run/install):\n\n```sh\n# bash, zsh, and other shells\nsh <(curl -fsSL https://cog.run/install.sh)\n\n# fish shell\nsh (curl -fsSL https://cog.run/install.sh | psub)\n\n# download with wget and run in a separate command\nwget -qO- https://cog.run/install.sh\nsh ./install.sh\n```\n\nYou can manually install the latest release of Cog directly from GitHub\nby running the following commands in a terminal:\n\n```console\nsudo curl -o /usr/local/bin/cog -L \"https://github.com/replicate/cog/releases/latest/download/cog_$(uname -s)_$(uname -m)\"\nsudo chmod +x /usr/local/bin/cog\n```\n\nOr if you are on docker:\n\n```\nRUN sh -c \"INSTALL_DIR=\\\"/usr/local/bin\\\" SUDO=\\\"\\\" $(curl -fsSL https://cog.run/install.sh)\"\n```\n\n## Upgrade\n\nIf you're using macOS and you previously installed Cog with Homebrew, run the following:\n\n```console\nbrew upgrade replicate/tap/cog\n```\n\nOtherwise, you can upgrade to the latest version by running the same commands you used to install it.\n\n## Development\n\nSee [CONTRIBUTING.md](CONTRIBUTING.md) for how to set up a development environment and build from source.\n\n## Next steps\n\n- [Get started with an example model](docs/getting-started.md)\n- [Get started with your own model](docs/getting-started-own-model.md)\n- [Using Cog with notebooks](docs/notebooks.md)\n- [Using Cog with Windows 11](docs/wsl2/wsl2.md)\n- [Take a look at some examples of using Cog](https://github.com/replicate/cog-examples)\n- [Deploy models with Cog](docs/deploy.md)\n- [`cog.yaml` reference](docs/yaml.md) to learn how to define your model's environment\n- [Prediction interface reference](docs/python.md) to learn how the `Predictor` interface works\n- [Training interface reference](docs/training.md) to learn how to add a fine-tuning API to your model\n- [HTTP API reference](docs/http.md) to learn how to use the HTTP API that models serve\n\n## Need help?\n\n[Join us in #cog on Discord.](https://discord.gg/replicate)\n\n[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/replicate/cog)\n\n\n\n---\n\n# CLI reference\n\n\n\n## `cog`\n\nContainers for machine learning.\n\nTo get started, take a look at the documentation:\nhttps://github.com/replicate/cog\n\n**Examples**\n\n```\n To run a command inside a Docker environment defined with Cog:\n $ cog run echo hello world\n```\n\n**Options**\n\n```\n --debug Show debugging output\n -h, --help help for cog\n --no-color Disable colored output\n --version Show version of Cog\n```\n## `cog build`\n\nBuild a Docker image from the cog.yaml in the current directory.\n\nThe generated image contains your model code, dependencies, and the Cog\nruntime. It can be run locally with 'cog predict' or pushed to a registry\nwith 'cog push'.\n\n```\ncog build [flags]\n```\n\n**Examples**\n\n```\n # Build with default settings\n cog build\n\n # Build and tag the image\n cog build -t my-model:latest\n\n # Build without using the cache\n cog build --no-cache\n\n # Build with model weights in a separate layer\n cog build --separate-weights -t my-model:v1\n```\n\n**Options**\n\n```\n -f, --file string The name of the config file. (default \"cog.yaml\")\n -h, --help help for build\n --no-cache Do not use cache when building the image\n --openapi-schema string Load OpenAPI schema from a file\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n --secret stringArray Secrets to pass to the build environment in the form 'id=foo,src=/path/to/file'\n --separate-weights Separate model weights from code in image layers\n -t, --tag string A name for the built image in the form 'repository:tag'\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n```\n## `cog init`\n\nCreate a cog.yaml and predict.py in the current directory.\n\nThese files provide a starting template for defining your model's environment\nand prediction interface. Edit them to match your model's requirements.\n\n```\ncog init [flags]\n```\n\n**Examples**\n\n```\n # Set up a new Cog project in the current directory\n cog init\n```\n\n**Options**\n\n```\n -h, --help help for init\n```\n## `cog login`\n\nLog in to a container registry.\n\nFor Replicate's registry (r8.im), this command handles authentication\nthrough Replicate's token-based flow.\n\nFor other registries, this command prompts for username and password,\nthen stores credentials using Docker's credential system.\n\n```\ncog login [flags]\n```\n\n**Options**\n\n```\n -h, --help help for login\n --token-stdin Pass login token on stdin instead of opening a browser. You can find your Replicate login token at https://replicate.com/auth/token\n```\n## `cog predict`\n\nRun a prediction.\n\nIf 'image' is passed, it will run the prediction on that Docker image.\nIt must be an image that has been built by Cog.\n\nOtherwise, it will build the model in the current directory and run\nthe prediction on that.\n\n```\ncog predict [image] [flags]\n```\n\n**Examples**\n\n```\n # Run a prediction with named inputs\n cog predict -i prompt=\"a photo of a cat\"\n\n # Pass a file as input\n cog predict -i image=@photo.jpg\n\n # Save output to a file\n cog predict -i image=@input.jpg -o output.png\n\n # Pass multiple inputs\n cog predict -i prompt=\"sunset\" -i width=1024 -i height=768\n\n # Run against a pre-built image\n cog predict r8.im/your-username/my-model -i prompt=\"hello\"\n\n # Pass inputs as JSON\n echo '{\"prompt\": \"a cat\"}' | cog predict --json @-\n```\n\n**Options**\n\n```\n -e, --env stringArray Environment variables, in the form name=value\n -f, --file string The name of the config file. (default \"cog.yaml\")\n --gpus docker run --gpus GPU devices to add to the container, in the same format as docker run --gpus.\n -h, --help help for predict\n -i, --input stringArray Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg\n --json string Pass inputs as JSON object, read from file (@inputs.json) or via stdin (@-)\n -o, --output string Output path\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n --setup-timeout uint32 The timeout for a container to setup (in seconds). (default 300)\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n --use-replicate-token Pass REPLICATE_API_TOKEN from local environment into the model context\n```\n## `cog push`\n\nBuild a Docker image from cog.yaml and push it to a container registry.\n\nCog can push to any OCI-compliant registry. When pushing to Replicate's\nregistry (r8.im), run 'cog login' first to authenticate.\n\n```\ncog push [IMAGE] [flags]\n```\n\n**Examples**\n\n```\n # Push to Replicate\n cog push r8.im/your-username/my-model\n\n # Push to any OCI registry\n cog push registry.example.com/your-username/model-name\n\n # Push with model weights in a separate layer (Replicate only)\n cog push r8.im/your-username/my-model --separate-weights\n```\n\n**Options**\n\n```\n -f, --file string The name of the config file. (default \"cog.yaml\")\n -h, --help help for push\n --no-cache Do not use cache when building the image\n --openapi-schema string Load OpenAPI schema from a file\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n --secret stringArray Secrets to pass to the build environment in the form 'id=foo,src=/path/to/file'\n --separate-weights Separate model weights from code in image layers\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n```\n## `cog run`\n\nRun a command inside a Docker environment defined by cog.yaml.\n\nCog builds a temporary image from your cog.yaml configuration and runs the\ngiven command inside it. This is useful for debugging, running scripts, or\nexploring the environment your model will run in.\n\n```\ncog run [arg...] [flags]\n```\n\n**Examples**\n\n```\n # Open a Python interpreter inside the model environment\n cog run python\n\n # Run a script\n cog run python train.py\n\n # Run with environment variables\n cog run -e HUGGING_FACE_HUB_TOKEN=abc123 python download.py\n\n # Expose a port (e.g. for Jupyter)\n cog run -p 8888 jupyter notebook\n```\n\n**Options**\n\n```\n -e, --env stringArray Environment variables, in the form name=value\n -f, --file string The name of the config file. (default \"cog.yaml\")\n --gpus docker run --gpus GPU devices to add to the container, in the same format as docker run --gpus.\n -h, --help help for run\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n -p, --publish stringArray Publish a container's port to the host, e.g. -p 8000\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n```\n## `cog serve`\n\nRun a prediction HTTP server.\n\nBuilds the model and starts an HTTP server that exposes the model's inputs\nand outputs as a REST API. Compatible with the Cog HTTP protocol.\n\n```\ncog serve [flags]\n```\n\n**Examples**\n\n```\n # Start the server on the default port (8393)\n cog serve\n\n # Start on a custom port\n cog serve -p 5000\n\n # Test the server\n curl http://localhost:8393/predictions \\\n -X POST \\\n -H 'Content-Type: application/json' \\\n -d '{\"input\": {\"prompt\": \"a cat\"}}'\n```\n\n**Options**\n\n```\n -f, --file string The name of the config file. (default \"cog.yaml\")\n --gpus docker run --gpus GPU devices to add to the container, in the same format as docker run --gpus.\n -h, --help help for serve\n -p, --port int Port on which to listen (default 8393)\n --progress string Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet' (default \"auto\")\n --upload-url string Upload URL for file outputs (e.g. https://example.com/upload/)\n --use-cog-base-image Use pre-built Cog base image for faster cold boots (default true)\n --use-cuda-base-image string Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects (default \"auto\")\n```\n\n\n---\n\n# Deploy models with Cog\n\nCog containers are Docker containers that serve an HTTP server\nfor running predictions on your model.\nYou can deploy them anywhere that Docker containers run.\n\nThe server inside Cog containers is **coglet**, a Rust-based prediction server\nthat handles HTTP requests, worker process management, and prediction execution.\n\nThis guide assumes you have a model packaged with Cog.\nIf you don't, [follow our getting started guide](getting-started-own-model.md),\nor use [an example model](https://github.com/replicate/cog-examples).\n\n## Getting started\n\nFirst, build your model:\n\n```console\ncog build -t my-model\n```\n\nYou can serve predictions locally with `cog serve`:\n\n```console\ncog serve\n# or, from a built image:\ncog serve my-model\n```\n\nAlternatively, start the Docker container directly:\n\n```shell\n# If your model uses a CPU:\ndocker run -d -p 5001:5000 my-model\n\n# If your model uses a GPU:\ndocker run -d -p 5001:5000 --gpus all my-model\n```\n\nThe server listens on port 5000 inside the container (mapped to 5001 above).\n\nTo view the OpenAPI schema,\nopen [localhost:5001/openapi.json](http://localhost:5001/openapi.json)\nin your browser\nor use cURL to make a request:\n\n```console\ncurl http://localhost:5001/openapi.json\n```\n\nTo stop the server, run:\n\n```console\ndocker kill my-model\n```\n\nTo run a prediction on the model,\ncall the `/predictions` endpoint,\npassing input in the format expected by your model:\n\n```console\ncurl http://localhost:5001/predictions -X POST \\\n --header \"Content-Type: application/json\" \\\n --data '{\"input\": {\"image\": \"https://.../input.jpg\"}}'\n```\n\nFor more details about the HTTP API,\nsee the [HTTP API reference documentation](http.md).\n\n## Health checks\n\nThe server exposes a `GET /health-check` endpoint that returns the current status of the model container. Use this for readiness probes in orchestration systems like Kubernetes.\n\n```console\ncurl http://localhost:5001/health-check\n```\n\nThe response includes a `status` field with values like `STARTING`, `READY`, `BUSY`, `SETUP_FAILED`, or `DEFUNCT`. See the [HTTP API reference](http.md#get-health-check) for full details.\n\n## Concurrency\n\nBy default, the server processes one prediction at a time. To enable concurrent predictions, set the `concurrency.max` option in `cog.yaml`:\n\n```yaml\nconcurrency:\n max: 4\n```\n\nSee the [`cog.yaml` reference](yaml.md#concurrency) for more details.\n\n## Environment variables\n\nYou can configure runtime behavior with environment variables:\n\n- `COG_SETUP_TIMEOUT`: Maximum time in seconds for the `setup()` method (default: no timeout).\n\nSee the [environment variables reference](environment.md) for the full list.\n\n\n---\n\n# Environment variables\n\nThis guide lists the environment variables that change how Cog functions.\n\n## Build-time variables\n\n### `COG_SDK_WHEEL`\n\nControls which cog Python SDK wheel is installed in the Docker image during `cog build`. Takes precedence over `build.sdk_version` in `cog.yaml`.\n\n**Supported values:**\n\n| Value | Description |\n| -------------------- | ---------------------------------------------------- |\n| `pypi` | Install latest version from PyPI |\n| `pypi:0.12.0` | Install specific version from PyPI |\n| `dist` | Use wheel from `dist/` directory (requires git repo) |\n| `https://...` | Install from URL |\n| `/path/to/wheel.whl` | Install from local file path |\n\n**Default behavior:**\n\n- **Release builds**: Installs latest cog from PyPI\n- **Development builds**: Auto-detects wheel in `dist/` directory, falls back to latest PyPI\n\n**Examples:**\n\n```console\n# Use specific PyPI version\n$ COG_SDK_WHEEL=pypi:0.11.0 cog build\n\n# Use local development wheel\n$ COG_SDK_WHEEL=dist cog build\n\n# Use wheel from URL\n$ COG_SDK_WHEEL=https://example.com/cog-0.12.0-py3-none-any.whl cog build\n```\n\nThe `dist` option searches for wheels in:\n\n1. `./dist/` (current directory)\n2. `$REPO_ROOT/dist/` (if REPO_ROOT is set)\n3. `/dist/` (via `git rev-parse`, useful when running from subdirectories)\n\n### `COGLET_WHEEL`\n\nControls which coglet wheel is installed in the Docker image. Coglet is the Rust-based prediction server.\n\n**Supported values:** Same as `COG_SDK_WHEEL`\n\n**Default behavior:** For development builds, auto-detects a wheel in `dist/`. For release builds, installs the latest version from PyPI. Can be overridden with an explicit value.\n\n**Examples:**\n\n```console\n# Use local development wheel\n$ COGLET_WHEEL=dist cog build\n\n# Use specific version from PyPI\n$ COGLET_WHEEL=pypi:0.1.0 cog build\n```\n\n## Runtime variables\n\n### `COG_NO_UPDATE_CHECK`\n\nBy default, Cog automatically checks for updates\nand notifies you if there is a new version available.\n\nTo disable this behavior,\nset the `COG_NO_UPDATE_CHECK` environment variable to any value.\n\n```console\n$ COG_NO_UPDATE_CHECK=1 cog build # runs without automatic update check\n```\n\n### `COG_SETUP_TIMEOUT`\n\nControls the maximum time (in seconds) allowed for the model's `setup()` method to complete. If setup exceeds this timeout, the server will report a setup failure.\n\nBy default, there is no timeout — setup runs indefinitely.\n\nSet to `0` to disable the timeout (same as default). Invalid values are ignored with a warning.\n\n```console\n$ COG_SETUP_TIMEOUT=300 docker run -p 5000:5000 my-model # 5-minute setup timeout\n```\n\n### `COG_CA_CERT`\n\nInjects a custom CA certificate into the Docker image during `cog build`. This is useful when building behind a corporate proxy or VPN that uses custom certificate authorities (e.g. Cloudflare WARP).\n\n**Supported values:**\n\n| Value | Description |\n| -------------------------------- | ----------------------------------------------------------- |\n| `/path/to/cert.crt` | Path to a single PEM certificate file |\n| `/path/to/certs/` | Directory of `.crt` and `.pem` files (all are concatenated) |\n| `-----BEGIN CERTIFICATE-----...` | Inline PEM certificate |\n| `LS0tLS1CRUdJTi...` | Base64-encoded PEM certificate |\n\nThe certificate is installed into the system CA store and the `SSL_CERT_FILE` and `REQUESTS_CA_BUNDLE` environment variables are set automatically in the built image.\n\n**Examples:**\n\n```console\n# From a file\n$ COG_CA_CERT=/usr/local/share/ca-certificates/corporate-ca.crt cog build\n\n# From a directory of certs\n$ COG_CA_CERT=/etc/custom-certs/ cog build\n\n# Inline (e.g. from a CI secret)\n$ COG_CA_CERT=\"$(cat /path/to/cert.pem)\" cog build\n```\n\n\n---\n\n# Getting started with your own model\n\nThis guide will show you how to put your own machine learning model in a Docker image using Cog. If you haven't got a model to try out, you'll want to follow the [main getting started guide](getting-started.md).\n\n## Prerequisites\n\n- **macOS or Linux**. Cog works on macOS and Linux, but does not currently support Windows.\n- **Docker**. Cog uses Docker to create a container for your model. You'll need to [install Docker](https://docs.docker.com/get-docker/) before you can run Cog.\n\n## Initialization\n\nFirst, install Cog if you haven't already:\n\n**macOS (recommended):**\n\n```sh\nbrew install replicate/tap/cog\n```\n\n**Linux or macOS (manual):**\n\n```sh\nsudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`\nsudo chmod +x /usr/local/bin/cog\n```\n\nTo configure your project for use with Cog, you'll need to add two files:\n\n- [`cog.yaml`](yaml.md) defines system requirements, Python package dependencies, etc\n- [`predict.py`](python.md) describes the prediction interface for your model\n\nUse the `cog init` command to generate these files in your project:\n\n```sh\n$ cd path/to/your/model\n$ cog init\n```\n\n## Define the Docker environment\n\nThe `cog.yaml` file defines all the different things that need to be installed for your model to run. You can think of it as a simple way of defining a Docker image.\n\nFor example:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n python_requirements: requirements.txt\n```\n\nWith a `requirements.txt` containing your dependencies:\n\n```\ntorch==2.6.0\n```\n\nThis will generate a Docker image with Python 3.13 and PyTorch 2 installed, for both CPU and GPU, with the correct version of CUDA, and various other sensible best-practices.\n\nTo run a command inside this environment, prefix it with `cog run`:\n\n```\n$ cog run python\n✓ Building Docker image from cog.yaml... Successfully built 8f54020c8981\nRunning 'python' in Docker with the current directory mounted as a volume...\n────────────────────────────────────────────────────────────────────────────────────────\n\nPython 3.13.x (main, ...)\n[GCC 12.2.0] on linux\nType \"help\", \"copyright\", \"credits\" or \"license\" for more information.\n>>>\n```\n\nThis is handy for ensuring a consistent environment for development or training.\n\nWith `cog.yaml`, you can also install system packages and other things. [Take a look at the full reference to see what else you can do.](yaml.md)\n\n## Define how to run predictions\n\nThe next step is to update `predict.py` to define the interface for running predictions on your model. The `predict.py` generated by `cog init` looks something like this:\n\n```python\nfrom cog import BasePredictor, Path, Input\nimport torch\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n self.net = torch.load(\"weights.pth\")\n\n def predict(self,\n image: Path = Input(description=\"Image to enlarge\"),\n scale: float = Input(description=\"Factor to scale image by\", default=1.5)\n ) -> Path:\n \"\"\"Run a single prediction on the model\"\"\"\n # ... pre-processing ...\n output = self.net(input)\n # ... post-processing ...\n return output\n```\n\nEdit your `predict.py` file and fill in the functions with your own model's setup and prediction code. You might need to import parts of your model from another file.\n\nYou also need to define the inputs to your model as arguments to the `predict()` function, as demonstrated above. For each argument, you need to annotate with a type. The supported types are:\n\n- `str`: a string\n- `int`: an integer\n- `float`: a floating point number\n- `bool`: a boolean\n- `cog.File`: a file-like object representing a file (deprecated — use `cog.Path` instead)\n- `cog.Path`: a path to a file on disk\n\nYou can provide more information about the input with the `Input()` function, as shown above. It takes these basic arguments:\n\n- `description`: A description of what to pass to this input for users of the model\n- `default`: A default value to set the input to. If this argument is not passed, the input is required. If it is explicitly set to `None`, the input is optional.\n- `ge`: For `int` or `float` types, the value should be greater than or equal to this number.\n- `le`: For `int` or `float` types, the value should be less than or equal to this number.\n- `min_length`: For `str` types, the minimum length of the string.\n- `max_length`: For `str` types, the maximum length of the string.\n- `regex`: For `str` types, the string must match this regular expression.\n- `choices`: For `str` or `int` types, a list of possible values for this input.\n- `deprecated`: Mark this input as deprecated with a message explaining what to use instead.\n\nThere are some more advanced options you can pass, too. For more details, [take a look at the prediction interface documentation](python.md).\n\nNext, add the line `predict: \"predict.py:Predictor\"` to your `cog.yaml`, so it looks something like this:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n python_requirements: requirements.txt\npredict: \"predict.py:Predictor\"\n```\n\nThat's it! To test this works, try running a prediction on the model:\n\n```\n$ cog predict -i image=@input.jpg\n✓ Building Docker image from cog.yaml... Successfully built 664ef88bc1f4\n✓ Model running in Docker image 664ef88bc1f4\n\nWritten output to output.png\n```\n\nTo pass more inputs to the model, you can add more `-i` options:\n\n```\n$ cog predict -i image=@image.jpg -i scale=2.0\n```\n\nIn this case it is just a number, not a file, so you don't need the `@` prefix.\n\n## Using GPUs\n\nTo use GPUs with Cog, add the `gpu: true` option to the `build` section of your `cog.yaml`:\n\n```yaml\nbuild:\n gpu: true\n ...\n```\n\nCog will use the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image and automatically figure out what versions of CUDA and cuDNN to use based on the version of Python, PyTorch, and Tensorflow that you are using.\n\nFor more details, [see the `gpu` section of the `cog.yaml` reference](yaml.md#gpu).\n\n## Next steps\n\nNext, you might want to take a look at:\n\n- [A guide explaining how to deploy a model.](deploy.md)\n- [The reference for `cog.yaml`](yaml.md)\n- [The reference for the Python library](python.md)\n\n\n---\n\n# Getting started\n\nThis guide will walk you through what you can do with Cog by using an example model.\n\n> [!TIP]\n> Using a language model to help you write the code for your new Cog model?\n>\n> Feed it [https://cog.run/llms.txt](https://cog.run/llms.txt), which has all of Cog's documentation bundled into a single file. To learn more about this format, check out [llmstxt.org](https://llmstxt.org).\n\n## Prerequisites\n\n- **macOS or Linux**. Cog works on macOS and Linux, but does not currently support Windows.\n- **Docker**. Cog uses Docker to create a container for your model. You'll need to [install Docker](https://docs.docker.com/get-docker/) before you can run Cog.\n\n## Install Cog\n\n**macOS (recommended):**\n\n```bash\nbrew install replicate/tap/cog\n```\n\n**Linux or macOS (manual):**\n\n```bash\nsudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`\nsudo chmod +x /usr/local/bin/cog\nsudo xattr -d com.apple.quarantine /usr/local/bin/cog 2>/dev/null || true\n\n```\n\n> [!NOTE]\n> **macOS: \"cannot be opened because the developer cannot be verified\"**\n>\n> If you downloaded the binary manually (via `curl` or a browser) and see this Gatekeeper warning, run:\n>\n> ```bash\n> sudo xattr -d com.apple.quarantine /usr/local/bin/cog\n> ```\n>\n> Installing via `brew install replicate/tap/cog` handles this automatically.\n\n## Create a project\n\nLet's make a directory to work in:\n\n```bash\nmkdir cog-quickstart\ncd cog-quickstart\n\n```\n\n## Run commands\n\nThe simplest thing you can do with Cog is run a command inside a Docker environment.\n\nThe first thing you need to do is create a file called `cog.yaml`:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n```\n\nThen, you can run any command inside this environment. For example, enter\n\n```bash\ncog run python\n\n```\n\nand you'll get an interactive Python shell:\n\n```none\n✓ Building Docker image from cog.yaml... Successfully built 8f54020c8981\nRunning 'python' in Docker with the current directory mounted as a volume...\n───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n\nPython 3.13.x (main, ...)\n[GCC 12.2.0] on linux\nType \"help\", \"copyright\", \"credits\" or \"license\" for more information.\n>>>\n```\n\n(Hit Ctrl-D to exit the Python shell.)\n\nInside this Docker environment you can do anything – run a Jupyter notebook, your training script, your evaluation script, and so on.\n\n## Run predictions on a model\n\nLet's pretend we've trained a model. With Cog, we can define how to run predictions on it in a standard way, so other people can easily run predictions on it without having to hunt around for a prediction script.\n\nWe need to write some code to describe how predictions are run on the model.\n\nSave this to `predict.py`:\n\n```python\nimport os\nos.environ[\"TORCH_HOME\"] = \".\"\n\nimport torch\nfrom cog import BasePredictor, Input, Path\nfrom PIL import Image\nfrom torchvision import models\n\nWEIGHTS = models.ResNet50_Weights.IMAGENET1K_V1\n\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n self.model = models.resnet50(weights=WEIGHTS).to(self.device)\n self.model.eval()\n\n def predict(self, image: Path = Input(description=\"Image to classify\")) -> dict:\n \"\"\"Run a single prediction on the model\"\"\"\n img = Image.open(image).convert(\"RGB\")\n preds = self.model(WEIGHTS.transforms()(img).unsqueeze(0).to(self.device))\n top3 = preds[0].softmax(0).topk(3)\n categories = WEIGHTS.meta[\"categories\"]\n return {categories[i]: p.detach().item() for p, i in zip(*top3)}\n```\n\nWe also need to point Cog at this, and tell it what Python dependencies to install.\n\nSave this to `requirements.txt`:\n\n```\npillow==11.1.0\ntorch==2.6.0\ntorchvision==0.21.0\n```\n\nThen update `cog.yaml` to look like this:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n python_requirements: requirements.txt\npredict: \"predict.py:Predictor\"\n```\n\n> [!TIP]\n> If you have a machine with an NVIDIA GPU attached, add `gpu: true` to the `build` section of your `cog.yaml` to enable GPU acceleration.\n\nLet's grab an image to test the model with:\n\n```bash\nIMAGE_URL=https://gist.githubusercontent.com/bfirsh/3c2115692682ae260932a67d93fd94a8/raw/56b19f53f7643bb6c0b822c410c366c3a6244de2/mystery.jpg\ncurl $IMAGE_URL > input.jpg\n\n```\n\nNow, let's run the model using Cog:\n\n```bash\ncog predict -i image=@input.jpg\n\n```\n\nIf you see the following output\n\n```json\n{\n \"tiger_cat\": 0.4874822497367859,\n \"tabby\": 0.23169134557247162,\n \"Egyptian_cat\": 0.09728282690048218\n}\n```\n\nthen it worked!\n\nNote: The first time you run `cog predict`, the build process will be triggered to generate a Docker container that can run your model. The next time you run `cog predict` the pre-built container will be used.\n\n## Build an image\n\nWe can bake your model's code, the trained weights, and the Docker environment into a Docker image. This image serves predictions with an HTTP server, and can be deployed to anywhere that Docker runs to serve real-time predictions.\n\n```bash\ncog build -t resnet\n# Building Docker image...\n# Built resnet:latest\n\n```\n\nYou can run this image with `cog predict` by passing the filename as an argument:\n\n```bash\ncog predict resnet -i image=@input.jpg\n\n```\n\nOr, you can run it with Docker directly, and it'll serve an HTTP server:\n\n```bash\ndocker run -d --rm -p 5000:5000 resnet\n\n```\n\nWe can send inputs directly with `curl`:\n\n```bash\ncurl http://localhost:5000/predictions -X POST \\\n -H 'Content-Type: application/json' \\\n -d '{\"input\": {\"image\": \"https://gist.githubusercontent.com/bfirsh/3c2115692682ae260932a67d93fd94a8/raw/56b19f53f7643bb6c0b822c410c366c3a6244de2/mystery.jpg\"}}'\n\n```\n\nAs a shorthand, you can add the Docker image's name as an extra line in `cog.yaml`:\n\n```yaml\nimage: \"r8.im/replicate/resnet\"\n```\n\nOnce you've done this, you can use `cog push` to build and push the image to a Docker registry:\n\n```bash\ncog push\n# Building r8.im/replicate/resnet...\n# Pushing r8.im/replicate/resnet...\n# Pushed!\n```\n\nThe Docker image is now accessible to anyone or any system that has access to this Docker registry.\n\n## Next steps\n\nThose are the basics! Next, you might want to take a look at:\n\n- [A guide to help you set up your own model on Cog.](getting-started-own-model.md)\n- [A guide explaining how to deploy a model.](deploy.md)\n- [Reference for `cog.yaml`](yaml.md)\n- [Reference for the Python library](python.md)\n\n\n---\n\n# HTTP API\n\n> [!TIP]\n> For information about how to run the HTTP server,\n> see [our documentation on deploying models](deploy.md).\n\nWhen you run a Docker image built by Cog,\nit serves an HTTP API for making predictions.\n\nThe server supports both synchronous and asynchronous prediction creation:\n\n- **Synchronous**:\n The server waits until the prediction is completed\n and responds with the result.\n- **Asynchronous**:\n The server immediately returns a response\n and processes the prediction in the background.\n\nThe client can create a prediction asynchronously\nby setting the `Prefer: respond-async` header in their request.\nWhen provided, the server responds immediately after starting the prediction\nwith `202 Accepted` status and a prediction object in status `processing`.\n\n> [!NOTE]\n> The only supported way to receive updates on the status of predictions\n> started asynchronously is using [webhooks](#webhooks).\n> Polling for prediction status is not currently supported.\n\nYou can also use certain server endpoints to create predictions idempotently,\nsuch that if a client calls this endpoint more than once with the same ID\n(for example, due to a network interruption)\nwhile the prediction is still running,\nno new prediction is created.\nInstead, the client receives a `202 Accepted` response\nwith the initial state of the prediction.\n\n---\n\nHere's a summary of the prediction creation endpoints:\n\n| Endpoint | Header | Behavior |\n| ---------------------------------- | ----------------------- | ---------------------------- |\n| `POST /predictions` | - | Synchronous, non-idempotent |\n| `POST /predictions` | `Prefer: respond-async` | Asynchronous, non-idempotent |\n| `PUT /predictions/` | - | Synchronous, idempotent |\n| `PUT /predictions/` | `Prefer: respond-async` | Asynchronous, idempotent |\n\nChoose the endpoint that best fits your needs:\n\n- Use synchronous endpoints when you want to wait for the prediction result.\n- Use asynchronous endpoints when you want to start a prediction\n and receive updates via webhooks.\n- Use idempotent endpoints when you need to safely retry requests\n without creating duplicate predictions.\n\n## Webhooks\n\nYou can provide a `webhook` parameter in the client request body\nwhen creating a prediction.\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n \"webhook\": \"https://example.com/webhook/prediction\"\n}\n```\n\nThe server makes requests to the provided URL\nwith the current state of the prediction object in the request body\nat the following times.\n\n- `start`:\n Once, when the prediction starts\n (`status` is `starting`).\n- `output`:\n Each time a predict function generates an output\n (either once using `return` or multiple times using `yield`)\n- `logs`:\n Each time the predict function writes to `stdout`\n- `completed`:\n Once, when the prediction reaches a terminal state\n (`status` is `succeeded`, `canceled`, or `failed`)\n\nWebhook requests for `start` and `completed` event types\nare sent immediately.\nWebhook requests for `output` and `logs` event types\nare sent at most once every 500ms.\nThis interval is not configurable.\n\nBy default, the server sends requests for all event types.\nClients can specify which events trigger webhook requests\nwith the `webhook_events_filter` parameter in the prediction request body.\nFor example,\nthe following request specifies that webhooks are sent by the server\nonly at the start and end of the prediction:\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n \"webhook\": \"https://example.com/webhook/prediction\",\n \"webhook_events_filter\": [\"start\", \"completed\"]\n}\n```\n\n## Generating unique prediction IDs\n\nEndpoints for creating and canceling a prediction idempotently\naccept a `prediction_id` parameter in their path.\nBy default, the server runs one prediction at a time,\nbut this can be increased with the [`concurrency.max`](yaml.md#concurrency) setting.\nWhen all prediction slots are in use, the server returns `409 Conflict`.\nThe client should ensure prediction slots are available\nbefore creating a new prediction with a different ID.\n\nClients are responsible for providing unique prediction IDs.\nWe recommend generating a UUIDv4 or [UUIDv7](https://uuid7.com),\nbase32-encoding that value,\nand removing padding characters (`==`).\nThis produces a random identifier that is 26 ASCII characters long.\n\n```python\n>> from uuid import uuid4\n>> from base64 import b32encode\n>> b32encode(uuid4().bytes).decode('utf-8').lower().rstrip('=')\n'wjx3whax6rf4vphkegkhcvpv6a'\n```\n\n## File uploads\n\nA model's `predict` function can produce file output by yielding or returning\na `cog.Path` or `cog.File` value.\n\nBy default,\nfiles are returned as a base64-encoded\n[data URL](https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs).\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n}\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"succeeded\",\n \"output\": \"data:image/png;base64,...\"\n}\n```\n\nWhen creating a prediction synchronously,\nthe client can configure a base URL to upload output files to instead\nby setting the `output_file_prefix` parameter in the request body:\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n \"output_file_prefix\": \"https://example.com/upload\",\n}\n```\n\nWhen the model produces a file output,\nthe server sends the following request to upload the file to the configured URL:\n\n```http\nPUT /upload HTTP/1.1\nHost: example.com\nContent-Type: multipart/form-data\n\n--boundary\nContent-Disposition: form-data; name=\"file\"; filename=\"image.png\"\nContent-Type: image/png\n\n\n--boundary--\n```\n\nIf the upload succeeds, the server responds with output:\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"succeeded\",\n \"output\": \"http://example.com/upload/image.png\"\n}\n```\n\nIf the upload fails, the server responds with an error.\n\n> [!IMPORTANT] \n> File uploads for predictions created asynchronously\n> require `--upload-url` to be specified when starting the HTTP server.\n\n\n\n## Endpoints\n\n### `GET /`\n\nReturns a discovery document listing available API endpoints, the OpenAPI schema URL, and version information.\n\n```http\nGET / HTTP/1.1\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"cog_version\": \"0.17.0\",\n \"docs_url\": \"/openapi.json\",\n \"openapi_url\": \"/openapi.json\",\n \"predictions_url\": \"/predictions\",\n \"health_check_url\": \"/health-check\"\n}\n```\n\nIf training is configured, the response also includes a `trainings_url` field.\n\n### `GET /health-check`\n\nReturns the current health status of the model container.\nThis endpoint always responds with `200 OK` —\ncheck the `status` field in the response body to determine readiness.\n\nThe response body is a JSON object with the following fields:\n\n- `status`: One of the following values:\n - `STARTING`: The model's `setup()` method is still running.\n - `READY`: The model is ready to accept predictions.\n - `BUSY`: The model is ready but all prediction slots are in use.\n - `SETUP_FAILED`: The model's `setup()` method raised an exception.\n - `DEFUNCT`: The model encountered an unrecoverable error.\n - `UNHEALTHY`: The model is ready\n but a user-defined `healthcheck()` method returned `False`.\n- `setup`: Setup phase details (included once setup has started):\n - `started_at`: ISO 8601 timestamp of when setup began.\n - `completed_at`: ISO 8601 timestamp of when setup finished (if complete).\n - `status`: One of `starting`, `succeeded`, or `failed`.\n - `logs`: Output captured during setup.\n- `version`: Runtime version information:\n - `coglet`: Coglet version.\n - `cog`: Cog Python SDK version (if available).\n - `python`: Python version (if available).\n- `user_healthcheck_error`:\n Error message from a user-defined `healthcheck()` method (if applicable).\n\n```http\nGET /health-check HTTP/1.1\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"READY\",\n \"setup\": {\n \"started_at\": \"2025-01-01T00:00:00.000000+00:00\",\n \"completed_at\": \"2025-01-01T00:00:05.000000+00:00\",\n \"status\": \"succeeded\",\n \"logs\": \"\"\n },\n \"version\": {\n \"coglet\": \"0.17.0\",\n \"cog\": \"0.14.0\",\n \"python\": \"3.13.0\"\n }\n}\n```\n\n### `GET /openapi.json`\n\nThe [OpenAPI](https://swagger.io/specification/) specification of the API,\nwhich is derived from the input and output types specified in your model's\n[Predictor](python.md) and [Training](training.md) objects.\n\n### `POST /predictions`\n\nMakes a single prediction.\n\nThe request body is a JSON object with the following fields:\n\n- `input`:\n A JSON object with the same keys as the\n [arguments to the `predict()` function](python.md).\n Any `File` or `Path` inputs are passed as URLs.\n\nThe response body is a JSON object with the following fields:\n\n- `status`: Either `succeeded` or `failed`.\n- `output`: The return value of the `predict()` function.\n- `error`: If `status` is `failed`, the error message.\n- `metrics`: An object containing prediction metrics.\n Always includes `predict_time` (elapsed seconds).\n May also include custom metrics recorded by the model\n using [`self.record_metric()`](python.md#metrics).\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\n\n{\n \"input\": {\n \"image\": \"https://example.com/image.jpg\",\n \"text\": \"Hello world!\"\n }\n}\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"succeeded\",\n \"output\": \"data:image/png;base64,...\",\n \"metrics\": {\n \"predict_time\": 4.52\n }\n}\n```\n\nIf the client sets the `Prefer: respond-async` header in their request,\nthe server responds immediately after starting the prediction\nwith `202 Accepted` status and a prediction object in status `processing`.\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"}\n}\n```\n\n```http\nHTTP/1.1 202 Accepted\nContent-Type: application/json\n\n{\n \"status\": \"starting\",\n}\n```\n\n### `PUT /predictions/`\n\nMake a single prediction.\nThis is the idempotent version of the `POST /predictions` endpoint.\n\n```http\nPUT /predictions/wjx3whax6rf4vphkegkhcvpv6a HTTP/1.1\nContent-Type: application/json; charset=utf-8\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"}\n}\n```\n\n```http\nHTTP/1.1 200 OK\nContent-Type: application/json\n\n{\n \"status\": \"succeeded\",\n \"output\": \"data:image/png;base64,...\"\n}\n```\n\nIf the client sets the `Prefer: respond-async` header in their request,\nthe server responds immediately after starting the prediction\nwith `202 Accepted` status and a prediction object in status `processing`.\n\n```http\nPUT /predictions/wjx3whax6rf4vphkegkhcvpv6a HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"}\n}\n```\n\n```http\nHTTP/1.1 202 Accepted\nContent-Type: application/json\n\n{\n \"id\": \"wjx3whax6rf4vphkegkhcvpv6a\",\n \"status\": \"starting\"\n}\n```\n\n### `POST /predictions//cancel`\n\nA client can cancel an asynchronous prediction by making a\n`POST /predictions//cancel` request\nusing the prediction `id` provided when the prediction was created.\n\nFor example,\nif the client creates a prediction by sending the request:\n\n```http\nPOST /predictions HTTP/1.1\nContent-Type: application/json; charset=utf-8\nPrefer: respond-async\n\n{\n \"id\": \"abcd1234\",\n \"input\": {\"prompt\": \"A picture of an onion with sunglasses\"},\n}\n```\n\nThe client can cancel the prediction by sending the request:\n\n```http\nPOST /predictions/abcd1234/cancel HTTP/1.1\n```\n\nA prediction cannot be canceled if it's\ncreated synchronously, without the `Prefer: respond-async` header,\nor created without a provided `id`.\n\nIf a prediction exists with the provided `id`,\nthe server responds with status `200 OK`.\nOtherwise, the server responds with status `404 Not Found`.\n\nWhen a prediction is canceled,\nCog raises [`CancelationException`](python.md#cancelationexception)\nin sync predictors (or `asyncio.CancelledError` in async predictors).\nThis exception may be caught by the model to perform necessary cleanup.\nThe cleanup should be brief, ideally completing within a few seconds.\nAfter cleanup, the exception must be re-raised using a bare `raise` statement.\nFailure to re-raise the exception may result in the termination of the container.\n\n```python\nfrom cog import BasePredictor, CancelationException, Input, Path\n\nclass Predictor(BasePredictor):\n def predict(self, image: Path = Input(description=\"Image to process\")) -> Path:\n try:\n return self.process(image)\n except CancelationException:\n self.cleanup()\n raise # always re-raise\n```\n\n\n---\n\n# Notebooks\n\nCog plays nicely with Jupyter notebooks.\n\n## Install the jupyterlab Python package\n\nFirst, add `jupyterlab` to your `requirements.txt` file and reference it in [`cog.yaml`](yaml.md):\n\n`requirements.txt`:\n\n```\njupyterlab\n```\n\n`cog.yaml`:\n\n```yaml\nbuild:\n python_requirements: requirements.txt\n```\n\n## Run a notebook\n\nCog can run notebooks in the environment you've defined in `cog.yaml` with the following command:\n\n```sh\ncog run -p 8888 jupyter lab --allow-root --ip=0.0.0.0\n```\n\n## Use notebook code in your predictor\n\nYou can also import a notebook into your Cog [Predictor](python.md) file.\n\nFirst, export your notebook to a Python file:\n\n```sh\njupyter nbconvert --to script my_notebook.ipynb # creates my_notebook.py\n```\n\nThen import the exported Python script into your `predict.py` file. Any functions or variables defined in your notebook will be available to your predictor:\n\n```python\nfrom cog import BasePredictor, Input\n\nimport my_notebook\n\nclass Predictor(BasePredictor):\n def predict(self, prompt: str = Input(description=\"string prompt\")) -> str:\n output = my_notebook.do_stuff(prompt)\n return output\n```\n\n\n---\n\n# Private package registry\n\nThis guide describes how to build a Docker image with Cog that fetches Python packages from a private registry during setup.\n\n## `pip.conf`\n\nIn a directory outside your Cog project, create a `pip.conf` file with an `index-url` set to the registry's URL with embedded credentials.\n\n```conf\n[global]\nindex-url = https://username:password@my-private-registry.com\n```\n\n> **Warning**\n> Be careful not to commit secrets in Git or include them in Docker images. If your Cog project contains any sensitive files, make sure they're listed in `.gitignore` and `.dockerignore`.\n\n## `cog.yaml`\n\nIn your project's [`cog.yaml`](yaml.md) file, add a setup command to run `pip install` with a secret configuration file mounted to `/etc/pip.conf`.\n\n```yaml\nbuild:\n run:\n - command: pip install\n mounts:\n - type: secret\n id: pip\n target: /etc/pip.conf\n```\n\n## Build\n\nWhen building or pushing your model with Cog, pass the `--secret` option with an `id` matching the one specified in `cog.yaml`, along with a path to your local `pip.conf` file.\n\n```console\n$ cog build --secret id=pip,source=/path/to/pip.conf\n```\n\nUsing a secret mount allows the private registry credentials to be securely passed to the `pip install` setup command, without baking them into the Docker image.\n\n> **Warning**\n> If you run `cog build` or `cog push` and then change the contents of a secret source file, the cached version of the file will be used on subsequent builds, ignoring any changes you made. To update the contents of the target secret file, either change the `id` value in `cog.yaml` and the `--secret` option, or pass the `--no-cache` option to bypass the cache entirely.\n\n\n---\n\n# Prediction interface reference\n\nThis document defines the API of the `cog` Python module, which is used to define the interface for running predictions on your model.\n\n> [!TIP]\n> Run [`cog init`](getting-started-own-model.md#initialization) to generate an annotated `predict.py` file that can be used as a starting point for setting up your model.\n\n> [!TIP]\n> Using a language model to help you write the code for your new Cog model?\n>\n> Feed it [https://cog.run/llms.txt](https://cog.run/llms.txt), which has all of Cog's documentation bundled into a single file. To learn more about this format, check out [llmstxt.org](https://llmstxt.org).\n\n## Contents\n\n- [Contents](#contents)\n- [`BasePredictor`](#basepredictor)\n - [`Predictor.setup()`](#predictorsetup)\n - [`Predictor.predict(**kwargs)`](#predictorpredictkwargs)\n- [`async` predictors and concurrency](#async-predictors-and-concurrency)\n- [`Input(**kwargs)`](#inputkwargs)\n - [Deprecating inputs](#deprecating-inputs)\n- [Output](#output)\n - [Returning an object](#returning-an-object)\n - [Returning a list](#returning-a-list)\n - [Optional properties](#optional-properties)\n - [Streaming output](#streaming-output)\n- [Metrics](#metrics)\n - [Recording metrics](#recording-metrics)\n - [Accumulation modes](#accumulation-modes)\n - [Dot-path keys](#dot-path-keys)\n - [Type safety](#type-safety)\n- [Cancellation](#cancellation)\n - [`CancelationException`](#cancelationexception)\n- [Input and output types](#input-and-output-types)\n- [`File()`](#file)\n- [`Path()`](#path)\n- [`Secret`](#secret)\n- [`Optional`](#optional)\n- [`List`](#list)\n\n## `BasePredictor`\n\nYou define how Cog runs predictions on your model by defining a class that inherits from `BasePredictor`. It looks something like this:\n\n```python\nfrom cog import BasePredictor, Path, Input\nimport torch\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n self.model = torch.load(\"weights.pth\")\n\n def predict(self,\n image: Path = Input(description=\"Image to enlarge\"),\n scale: float = Input(description=\"Factor to scale image by\", default=1.5)\n ) -> Path:\n \"\"\"Run a single prediction on the model\"\"\"\n # ... pre-processing ...\n output = self.model(image)\n # ... post-processing ...\n return output\n```\n\nYour Predictor class should define two methods: `setup()` and `predict()`.\n\n### `Predictor.setup()`\n\nPrepare the model so multiple predictions run efficiently.\n\nUse this _optional_ method to include expensive one-off operations like loading trained models, instantiating data transformations, etc.\n\nMany models use this method to download their weights (e.g. using [`pget`](https://github.com/replicate/pget)). This has some advantages:\n\n- Smaller image sizes\n- Faster build times\n- Faster pushes and inference on [Replicate](https://replicate.com)\n\nHowever, this may also significantly increase your `setup()` time.\n\nAs an alternative, some choose to store their weights directly in the image. You can simply leave your weights in the directory alongside your `cog.yaml` and ensure they are not excluded in your `.dockerignore` file.\n\nWhile this will increase your image size and build time, it offers other advantages:\n\n- Faster `setup()` time\n- Ensures idempotency and reduces your model's reliance on external systems\n- Preserves reproducibility as your model will be self-contained in the image\n\n> When using this method, you should use the `--separate-weights` flag on `cog build` to store weights in a [separate layer](https://github.com/replicate/cog/blob/12ac02091d93beebebed037f38a0c99cd8749806/docs/getting-started.md?plain=1#L219).\n\n### `Predictor.predict(**kwargs)`\n\nRun a single prediction.\n\nThis _required_ method is where you call the model that was loaded during `setup()`, but you may also want to add pre- and post-processing code here.\n\nThe `predict()` method takes an arbitrary list of named arguments, where each argument name must correspond to an [`Input()`](#inputkwargs) annotation.\n\n`predict()` can return strings, numbers, [`cog.Path`](#path) objects representing files on disk, or lists or dicts of those types. You can also define a custom [`Output()`](#outputbasemodel) for more complex return types.\n\n## `async` predictors and concurrency\n\n> Added in cog 0.14.0.\n\nYou may specify your `predict()` method as `async def predict(...)`. In\naddition, if you have an async `predict()` function you may also have an async\n`setup()` function:\n\n```py\nclass Predictor(BasePredictor):\n async def setup(self) -> None:\n print(\"async setup is also supported...\")\n\n async def predict(self) -> str:\n print(\"async predict\");\n return \"hello world\";\n```\n\nModels that have an async `predict()` function can run predictions concurrently, up to the limit specified by [`concurrency.max`](yaml.md#max) in cog.yaml. Attempting to exceed this limit will return a 409 Conflict response.\n\n## `Input(**kwargs)`\n\nUse cog's `Input()` function to define each of the parameters in your `predict()` method:\n\n```py\nclass Predictor(BasePredictor):\n def predict(self,\n image: Path = Input(description=\"Image to enlarge\"),\n scale: float = Input(description=\"Factor to scale image by\", default=1.5, ge=1.0, le=10.0)\n ) -> Path:\n```\n\nThe `Input()` function takes these keyword arguments:\n\n- `description`: A description of what to pass to this input for users of the model.\n- `default`: A default value to set the input to. If this argument is not passed, the input is required. If it is explicitly set to `None`, the input is optional.\n- `ge`: For `int` or `float` types, the value must be greater than or equal to this number.\n- `le`: For `int` or `float` types, the value must be less than or equal to this number.\n- `min_length`: For `str` types, the minimum length of the string.\n- `max_length`: For `str` types, the maximum length of the string.\n- `regex`: For `str` types, the string must match this regular expression.\n- `choices`: For `str` or `int` types, a list of possible values for this input.\n- `deprecated`: (optional) If set to `True`, marks this input as deprecated. Deprecated inputs will still be accepted, but tools and UIs may warn users that the input is deprecated and may be removed in the future. See [Deprecating inputs](#deprecating-inputs).\n\nEach parameter of the `predict()` method must be annotated with a type like `str`, `int`, `float`, `bool`, etc. See [Input and output types](#input-and-output-types) for the full list of supported types.\n\nUsing the `Input` function provides better documentation and validation constraints to the users of your model, but it is not strictly required. You can also specify default values for your parameters using plain Python, or omit default assignment entirely:\n\n```py\nclass Predictor(BasePredictor):\n def predict(self,\n prompt: str = \"default prompt\", # this is valid\n iterations: int # also valid\n ) -> str:\n # ...\n```\n\n## Deprecating inputs\n\nYou can mark an input as deprecated by passing `deprecated=True` to the `Input()` function. Deprecated inputs will still be accepted, but tools and UIs may warn users that the input is deprecated and may be removed in the future.\n\nThis is useful when you want to phase out an input without breaking existing clients immediately:\n\n```py\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n def predict(self,\n text: str = Input(description=\"Some deprecated text\", deprecated=True),\n prompt: str = Input(description=\"Prompt for the model\")\n ) -> str:\n # ...\n return prompt\n```\n\n## Output\n\nCog predictors can return a simple data type like a string, number, float, or boolean. Use Python's `-> ` syntax to annotate the return type.\n\nHere's an example of a predictor that returns a string:\n\n```py\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def predict(self) -> str:\n return \"hello\"\n```\n\n### Returning an object\n\nTo return a complex object with multiple values, define an `Output` object with multiple fields to return from your `predict()` method:\n\n```py\nfrom cog import BasePredictor, BaseModel, File\n\nclass Output(BaseModel):\n file: File\n text: str\n\nclass Predictor(BasePredictor):\n def predict(self) -> Output:\n return Output(text=\"hello\", file=io.StringIO(\"hello\"))\n```\n\nEach of the output object's properties must be one of the supported output types. For the full list, see [Input and output types](#input-and-output-types). Also, make sure to name the output class as `Output` and nothing else.\n\n### Returning a list\n\nThe `predict()` method can return a list of any of the supported output types. Here's an example that outputs multiple files:\n\n```py\nfrom cog import BasePredictor, Path\n\nclass Predictor(BasePredictor):\n def predict(self) -> list[Path]:\n predictions = [\"foo\", \"bar\", \"baz\"]\n output = []\n for i, prediction in enumerate(predictions):\n out_path = Path(f\"/tmp/out-{i}.txt\")\n with out_path.open(\"w\") as f:\n f.write(prediction)\n output.append(out_path)\n return output\n```\n\nFiles are named in the format `output..`, e.g. `output.0.txt`, `output.1.txt`, and `output.2.txt` from the example above.\n\n### Optional properties\n\nTo conditionally omit properties from the Output object, define them using `typing.Optional`:\n\n```py\nfrom cog import BaseModel, BasePredictor, Path\nfrom typing import Optional\n\nclass Output(BaseModel):\n score: Optional[float]\n file: Optional[Path]\n\nclass Predictor(BasePredictor):\n def predict(self) -> Output:\n if condition:\n return Output(score=1.5)\n else:\n return Output(file=io.StringIO(\"hello\"))\n```\n\n### Streaming output\n\nCog models can stream output as the `predict()` method is running. For example, a language model can output tokens as they're being generated and an image generation model can output images as they are being generated.\n\nTo support streaming output in your Cog model, add `from typing import Iterator` to your predict.py file. The `typing` package is a part of Python's standard library so it doesn't need to be installed. Then add a return type annotation to the `predict()` method in the form `-> Iterator[]` where `` can be one of `str`, `int`, `float`, `bool`, or `cog.Path`.\n\n```py\nfrom cog import BasePredictor, Path\nfrom typing import Iterator\n\nclass Predictor(BasePredictor):\n def predict(self) -> Iterator[Path]:\n done = False\n while not done:\n output_path, done = do_stuff()\n yield Path(output_path)\n```\n\nIf you have an [async `predict()` method](#async-predictors-and-concurrency), you must use `cog.AsyncIterator` instead:\n\n```py\nfrom cog import AsyncIterator, BasePredictor, Path\n\nclass Predictor(BasePredictor):\n async def predict(self) -> AsyncIterator[Path]:\n done = False\n while not done:\n output_path, done = do_stuff()\n yield Path(output_path)\n```\n\nIf you're streaming text output, you can use `ConcatenateIterator` to hint that the output should be concatenated together into a single string. This is useful on Replicate to display the output as a string instead of a list of strings.\n\n```py\nfrom cog import BasePredictor, Path, ConcatenateIterator\n\nclass Predictor(BasePredictor):\n def predict(self) -> ConcatenateIterator[str]:\n tokens = [\"The\", \"quick\", \"brown\", \"fox\", \"jumps\", \"over\", \"the\", \"lazy\", \"dog\"]\n for token in tokens:\n yield token + \" \"\n```\n\nOr for async `predict()` methods, use `AsyncConcatenateIterator`:\n\n```py\nfrom cog import BasePredictor, Path, AsyncConcatenateIterator\n\nclass Predictor(BasePredictor):\n async def predict(self) -> AsyncConcatenateIterator[str]:\n tokens = [\"The\", \"quick\", \"brown\", \"fox\", \"jumps\", \"over\", \"the\", \"lazy\", \"dog\"]\n for token in tokens:\n yield token + \" \"\n```\n\n## Metrics\n\nYou can record custom metrics from your `predict()` function to track model-specific data like token counts, timing breakdowns, or confidence scores. Metrics are included in the prediction response alongside the output.\n\n### Recording metrics\n\nUse `self.record_metric()` inside your `predict()` method:\n\n```python\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def predict(self, prompt: str) -> str:\n self.record_metric(\"temperature\", 0.7)\n self.record_metric(\"token_count\", 42)\n\n result = self.model.generate(prompt)\n return result\n```\n\nFor advanced use (dict-style access, deleting metrics), use `self.scope`:\n\n```python\nself.scope.metrics[\"token_count\"] = 42\ndel self.scope.metrics[\"token_count\"]\n```\n\nMetrics appear in the prediction response `metrics` field:\n\n```json\n{\n \"status\": \"succeeded\",\n \"output\": \"...\",\n \"metrics\": {\n \"temperature\": 0.7,\n \"token_count\": 42,\n \"predict_time\": 1.23\n }\n}\n```\n\nThe `predict_time` metric is always added automatically by the runtime. If you set `predict_time` yourself, the runtime value takes precedence.\n\nSupported value types are `bool`, `int`, `float`, `str`, `list`, and `dict`. Setting a metric to `None` deletes it.\n\n### Accumulation modes\n\nBy default, recording a metric replaces any previous value for that key. You can use accumulation modes to build up values across multiple calls:\n\n```python\n# Increment a counter (adds to the existing numeric value)\nself.record_metric(\"token_count\", 1, mode=\"incr\")\nself.record_metric(\"token_count\", 1, mode=\"incr\")\n# Result: {\"token_count\": 2}\n\n# Append to an array\nself.record_metric(\"steps\", \"preprocessing\", mode=\"append\")\nself.record_metric(\"steps\", \"inference\", mode=\"append\")\n# Result: {\"steps\": [\"preprocessing\", \"inference\"]}\n\n# Replace (default behavior)\nself.record_metric(\"status\", \"running\", mode=\"replace\")\nself.record_metric(\"status\", \"done\", mode=\"replace\")\n# Result: {\"status\": \"done\"}\n```\n\nThe `mode` parameter accepts `\"replace\"` (default), `\"incr\"`, or `\"append\"`.\n\n### Dot-path keys\n\nUse dot-separated keys to create nested objects in the metrics output:\n\n```python\nself.record_metric(\"timing.preprocess\", 0.12)\nself.record_metric(\"timing.inference\", 0.85)\n```\n\nThis produces nested JSON:\n\n```json\n{\n \"metrics\": {\n \"timing\": {\n \"preprocess\": 0.12,\n \"inference\": 0.85\n },\n \"predict_time\": 1.23\n }\n}\n```\n\n### Type safety\n\nOnce a metric key has been assigned a value of a certain type, it cannot be changed to a different type without deleting it first. This prevents accidental type mismatches when using accumulation modes:\n\n```python\nself.record_metric(\"count\", 1)\n\n# This would raise an error — \"count\" is an int, not a string:\n# self.record_metric(\"count\", \"oops\")\n\n# Delete first, then set with new type:\ndel self.scope.metrics[\"count\"]\nself.record_metric(\"count\", \"now a string\")\n```\n\nOutside an active prediction, `self.record_metric()` and `self.scope` are silent no-ops — no need for `None` checks.\n\n## Cancellation\n\nWhen a prediction is canceled (via the [cancel HTTP endpoint](http.md#post-predictionsprediction_idcancel) or a dropped connection), the Cog runtime interrupts the running `predict()` function. The exception raised depends on whether the predictor is sync or async:\n\n| Predictor type | Exception raised |\n| --------------------------- | ------------------------ |\n| Sync (`def predict`) | `CancelationException` |\n| Async (`async def predict`) | `asyncio.CancelledError` |\n\n### `CancelationException`\n\n```python\nfrom cog import CancelationException\n```\n\n`CancelationException` is raised in **sync** predictors when a prediction is cancelled. It is a `BaseException` subclass — **not** an `Exception` subclass. This means bare `except Exception` blocks in your predict code will not accidentally catch it, matching the behavior of `KeyboardInterrupt` and `asyncio.CancelledError`.\n\nYou do **not** need to handle this exception in normal predictor code — the runtime manages cancellation automatically. However, if you need to run cleanup logic when a prediction is cancelled, you can catch it explicitly:\n\n```python\nfrom cog import BasePredictor, CancelationException, Path\n\nclass Predictor(BasePredictor):\n def predict(self, image: Path) -> Path:\n try:\n return self.process(image)\n except CancelationException:\n self.cleanup()\n raise # always re-raise\n```\n\n> [!WARNING]\n> You **must** re-raise `CancelationException` after cleanup. Swallowing it will prevent the runtime from marking the prediction as canceled, and may result in the termination of the container.\n\n`CancelationException` is available as:\n\n- `cog.CancelationException` (recommended)\n- `cog.exceptions.CancelationException`\n\nFor **async** predictors, cancellation follows standard Python async conventions and raises `asyncio.CancelledError` instead.\n\n## Input and output types\n\nEach parameter of the `predict()` method must be annotated with a type. The method's return type must also be annotated. The supported types are:\n\n- `str`: a string\n- `int`: an integer\n- `float`: a floating point number\n- `bool`: a boolean\n- [`cog.File`](#file): a file-like object representing a file\n- [`cog.Path`](#path): a path to a file on disk\n- [`cog.Secret`](#secret): a string containing sensitive information\n\n## `File()`\n\n> [!WARNING] \n> `cog.File` is deprecated and will be removed in a future version of Cog. Use [`cog.Path`](#path) instead.\n\nThe `cog.File` object is used to get files in and out of models. It represents a _file handle_.\n\nFor models that return a `cog.File` object, the prediction output returned by Cog's built-in HTTP server will be a URL.\n\n```python\nfrom cog import BasePredictor, File, Input, Path\nfrom PIL import Image\n\nclass Predictor(BasePredictor):\n def predict(self, source_image: File = Input(description=\"Image to enlarge\")) -> File:\n pillow_img = Image.open(source_image)\n upscaled_image = do_some_processing(pillow_img)\n return File(upscaled_image)\n```\n\n## `Path()`\n\nThe `cog.Path` object is used to get files in and out of models. It represents a _path to a file on disk_.\n\n`cog.Path` is a subclass of Python's [`pathlib.Path`](https://docs.python.org/3/library/pathlib.html#basic-use) and can be used as a drop-in replacement.\n\nFor models that return a `cog.Path` object, the prediction output returned by Cog's built-in HTTP server will be a URL.\n\nThis example takes an input file, resizes it, and returns the resized image:\n\n```python\nimport tempfile\nfrom cog import BasePredictor, Input, Path\n\nclass Predictor(BasePredictor):\n def predict(self, image: Path = Input(description=\"Image to enlarge\")) -> Path:\n upscaled_image = do_some_processing(image)\n\n # To output `cog.Path` objects the file needs to exist, so create a temporary file first.\n # This file will automatically be deleted by Cog after it has been returned.\n output_path = Path(tempfile.mkdtemp()) / \"upscaled.png\"\n upscaled_image.save(output_path)\n return Path(output_path)\n```\n\n## `Secret`\n\nThe `cog.Secret` type is used to signify that an input holds sensitive information,\nlike a password or API token.\n\n`cog.Secret` is a type that redacts its contents in string representations to prevent accidental disclosure.\nYou can access its contents with the `get_secret_value()` method.\n\n```python\nfrom cog import BasePredictor, Secret\n\n\nclass Predictor(BasePredictor):\n def predict(self, api_token: Secret) -> None:\n # Prints '**********'\n print(api_token)\n\n # Use get_secret_value method to see the secret's content.\n print(api_token.get_secret_value())\n```\n\nA predictor's `Secret` inputs are represented in OpenAPI with the following schema:\n\n```json\n{\n \"type\": \"string\",\n \"format\": \"password\",\n \"x-cog-secret\": true\n}\n```\n\nModels uploaded to Replicate treat secret inputs differently throughout its system.\nWhen you create a prediction on Replicate,\nany value passed to a `Secret` input is redacted after being sent to the model.\n\n> [!WARNING] \n> Passing secret values to untrusted models can result in\n> unintended disclosure, exfiltration, or misuse of sensitive data.\n\n## `Optional`\n\nOptional inputs should be explicitly defined as `Optional[T]` so that type checker can warn us about error-prone `None` values.\n\nFor example, the following code might fail if `prompt` is not specified in the inputs:\n\n```python\nclass Predictor(BasePredictor):\n def predict(self, prompt: str=Input(description=\"prompt\", default=None)) -> str:\n return \"hello\" + prompt # TypeError: can only concatenate str (not \"NoneType\") to str\n```\n\nWe can improve it by making `prompt` an `Optional[str]`. Note that `default=None` is now redundant as `Optional` implies it.\n\n```python\nclass Predictor(BasePredictor):\n def predict(self, prompt: Optional[str]=Input(description=\"prompt\")) -> str:\n if prompt is None: # type check can warn us if we forget this\n return \"hello\"\n else:\n return \"hello\" + prompt\n```\n\nNote that the error prone usage of `prompt: str=Input(default=None)` might throw an error in a future release of Cog.\n\n## `List`\n\nThe List type is also supported in inputs. It can hold any supported type.\n\nExample for **List[Path]**:\n\n```py\nclass Predictor(BasePredictor):\n def predict(self, paths: list[Path]) -> str:\n output_parts = [] # Use a list to collect file contents\n for path in paths:\n with open(path) as f:\n output_parts.append(f.read())\n return \"\".join(output_parts)\n```\n\nThe corresponding cog command:\n\n```bash\n$ echo test1 > 1.txt\n$ echo test2 > 2.txt\n$ cog predict -i paths=@1.txt -i paths=@2.txt\nRunning prediction...\ntest1\n\ntest2\n```\n\n- Note the repeated inputs with the same name \"paths\" which constitute the list\n\n\n---\n\n# Training interface reference\n\n> [!WARNING] \n> The `cog train` command is deprecated and will be removed in the next version of Cog. The training API described below may still be used with the HTTP API's `/trainings` endpoint, but the CLI command is no longer recommended for new projects.\n\nCog's training API allows you to define a fine-tuning interface for an existing Cog model, so users of the model can bring their own training data to create derivative fine-tuned models. Real-world examples of this API in use include [fine-tuning SDXL with images](https://replicate.com/blog/fine-tune-sdxl) or [fine-tuning Llama 2 with structured text](https://replicate.com/blog/fine-tune-llama-2).\n\n## How it works\n\nIf you've used Cog before, you've probably seen the [Predictor](./python.md) class, which defines the interface for creating predictions against your model. Cog's training API works similarly: You define a Python function that describes the inputs and outputs of the training process. The inputs are things like training data, epochs, batch size, seed, etc. The output is typically a file with the fine-tuned weights.\n\n`cog.yaml`:\n\n```yaml\nbuild:\n python_version: \"3.13\"\ntrain: \"train.py:train\"\n```\n\n`train.py`:\n\n```python\nfrom cog import BasePredictor, File\nimport io\n\ndef train(param: str) -> File:\n return io.StringIO(\"hello \" + param)\n```\n\nThen you can run it like this:\n\n```\n$ cog train -i param=train\n...\n\n$ cat weights\nhello train\n```\n\nYou can also use classes if you want to run many model trainings and save on setup time. This works the same way as the [Predictor](./python.md) class with the only difference being the `train` method.\n\n`cog.yaml`:\n\n```yaml\nbuild:\n python_version: \"3.13\"\ntrain: \"train.py:Trainer\"\n```\n\n`train.py`:\n\n```python\nfrom cog import BasePredictor, File\nimport io\n\nclass Trainer:\n def setup(self) -> None:\n self.base_model = ... # Load a big base model\n\n def train(self, param: str) -> File:\n return self.base_model.train(param) # Train on top of a base model\n```\n\n## `Input(**kwargs)`\n\nUse Cog's `Input()` function to define each of the parameters in your `train()` function:\n\n```py\nfrom cog import Input, Path\n\ndef train(\n train_data: Path = Input(description=\"HTTPS URL of a file containing training data\"),\n learning_rate: float = Input(description=\"learning rate, for learning!\", default=1e-4, ge=0),\n seed: int = Input(description=\"random seed to use for training\", default=None)\n) -> str:\n return \"hello, weights\"\n```\n\nThe `Input()` function takes these keyword arguments:\n\n- `description`: A description of what to pass to this input for users of the model.\n- `default`: A default value to set the input to. If this argument is not passed, the input is required. If it is explicitly set to `None`, the input is optional.\n- `ge`: For `int` or `float` types, the value must be greater than or equal to this number.\n- `le`: For `int` or `float` types, the value must be less than or equal to this number.\n- `min_length`: For `str` types, the minimum length of the string.\n- `max_length`: For `str` types, the maximum length of the string.\n- `regex`: For `str` types, the string must match this regular expression.\n- `choices`: For `str` or `int` types, a list of possible values for this input.\n\nEach parameter of the `train()` function must be annotated with a type like `str`, `int`, `float`, `bool`, etc. See [Input and output types](./python.md#input-and-output-types) for the full list of supported types.\n\nUsing the `Input` function provides better documentation and validation constraints to the users of your model, but it is not strictly required. You can also specify default values for your parameters using plain Python, or omit default assignment entirely:\n\n```py\ndef train(self,\n training_data: str = \"foo bar\", # this is valid\n iterations: int # also valid\n) -> str:\n # ...\n```\n\n## Training Output\n\nTraining output is typically a binary weights file. To return a custom output object or a complex object with multiple values, define a `TrainingOutput` object with multiple fields to return from your `train()` function, and specify it as the return type for the train function using Python's `->` return type annotation:\n\n```python\nfrom cog import BaseModel, Input, Path\n\nclass TrainingOutput(BaseModel):\n weights: Path\n\ndef train(\n train_data: Path = Input(description=\"HTTPS URL of a file containing training data\"),\n learning_rate: float = Input(description=\"learning rate, for learning!\", default=1e-4, ge=0),\n seed: int = Input(description=\"random seed to use for training\", default=42)\n) -> TrainingOutput:\n weights_file = generate_weights(\"...\")\n return TrainingOutput(weights=Path(weights_file))\n```\n\n## Testing\n\nIf you are doing development of a Cog model like Llama or SDXL, you can test that the fine-tuned code path works before pushing by specifying a `COG_WEIGHTS` environment variable when running `predict`:\n\n```console\ncog predict -e COG_WEIGHTS=https://replicate.delivery/pbxt/xyz/weights.tar -i prompt=\"a photo of TOK\"\n```\n\n\n---\n\n# Using `cog` on Windows 11 with WSL 2\n\n- [0. Prerequisites](#0-prerequisites)\n- [1. Install the GPU driver](#1-install-the-gpu-driver)\n- [2. Unlocking features](#2-unlocking-features)\n - [2.1. Unlock WSL2](#21-unlock-wsl2)\n - [2.2. Unlock virtualization](#22-unlock-virtualization)\n - [2.3. Reboot](#23-reboot)\n- [3. Update MS Linux kernel](#3-update-ms-linux-kernel)\n- [4. Configure WSL 2](#4-configure-wsl-2)\n- [5. Configure CUDA WSL-Ubuntu Toolkit](#5-configure-cuda-wsl-ubuntu-toolkit)\n- [6. Install Docker](#6-install-docker)\n- [7. Install `cog` and pull an image](#7-install-cog-and-pull-an-image)\n- [8. Run a model in WSL 2](#8-run-a-model-in-wsl-2)\n- [9. References](#9-references)\n\nRunning cog on Windows is now possible thanks to WSL 2. Follow this guide to enable WSL 2 and GPU passthrough on Windows 11.\n\n**Windows 10 is not officially supported, as you need to be on an insider build in order to use GPU passthrough.**\n\n## 0. Prerequisites\n\nBefore beginning installation, make sure you have:\n\n- Windows 11.\n- NVIDIA GPU.\n - RTX 2000/3000 series\n - Kesler/Tesla/Volta/Ampere series\n - Other configurations are not guaranteed to work.\n\n## 1. Install the GPU driver\n\nPer NVIDIA, the first order of business is to install the latest Game Ready drivers for your NVIDIA GPU.\n\n\n\nI have an NVIDIA RTX 2070 Super, so filled out the form as such:\n\n![a form showing the correct model number selected for an RTX 2070 Super](images/nvidia_driver_select.png)\n\nClick \"search\", and follow the dialogue to download and install the driver.\n\nRestart your computer once the driver has finished installation.\n\n## 2. Unlocking features\n\nOpen Windows Terminal as an administrator.\n\n- Use start to search for \"Terminal\"\n- Right click -> Run as administrator...\n\nRun the following powershell command to enable the Windows Subsystem for Linux and Virtual Machine Platform capabilities.\n\n### 2.1. Unlock WSL2\n\n```powershell\ndism.exe /online /enable-feature /featurename:Microsoft-Windows-Subsystem-Linux /all /norestart\n```\n\nIf you see an error about permissions, make sure the terminal you are using is run as an administrator and that you have an account with administrator-level privileges.\n\n### 2.2. Unlock virtualization\n\n```powershell\ndism.exe /online /enable-feature /featurename:VirtualMachinePlatform /all /norestart\n```\n\nIf this command fails, make sure to [enable virtualization capabilities](https://docs.microsoft.com/en-us/windows/wsl/troubleshooting#error-0x80370102-the-virtual-machine-could-not-be-started-because-a-required-feature-is-not-installed) in your computer's BIOS/UEFI. A successful output will print `The operation completed successfully.`\n\n![Output from running the above commands successfully. Should read \"The operation completed successfully\".](images/enable_feature_success.png)\n\n### 2.3. Reboot\n\nBefore moving forward, make sure you reboot your computer so that Windows 11 will have WSL2 and virtualization available to it.\n\n## 3. Update MS Linux kernel\n\nDownload and run the [WSL2 Linux kernel update package for x64 machines](https://wslstorestorage.blob.core.windows.net/wslblob/wsl_update_x64.msi) msi installer. When prompted for elevated permissions, click 'yes' to approve the installation.\n\nTo ensure you are using the correct WSL kernel, `open Windows Terminal as an administrator` and enter:\n\n```powershell\nwsl cat /proc/version\n```\n\nThis will return a complicated string such as:\n\n```sh\nLinux version 5.10.102.1-microsoft-standard-WSL2 (oe-user@oe-host) (x86_64-msft-linux-gcc (GCC) 9.3.0, GNU ld (GNU Binutils) 2.34.0.20200220)\n```\n\nThe version we are interested in is `Linux version 5.10.102.1`. At this point, you should have updated your kernel to be at least `Linux version 5.10.43.3`.\n\nIf you can't get the correct kernel version to show:\n\nOpen `Settings` → `Windows Update` → `Advanced options` and ensure `Receive updates for other Microsoft products` is enabled. Then go to `Windows Update` again and click `Check for updates`.\n\n## 4. Configure WSL 2\n\nFirst, configure Windows to use the virtualization-based version of WSL (version 2) by default. In a Windows Terminal with administrator privileges, type the following:\n\n```powershell\nwsl --set-default-version 2\n```\n\nNow, you will need to go to the Microsoft Store and [Download Ubuntu 18.04](https://www.microsoft.com/store/apps/9N9TNGVNDL3Q)\n\n![Screenshot showing the \"Ubuntu\" store page](https://docs.microsoft.com/en-us/windows/wsl/media/ubuntustore.png)\n\nLaunch the \"Ubuntu\" app available in your Start Menu. Linux will require its own user account and password, which you will need to enter now:\n\n![a terminal showing input for user account info on WSL 2](https://docs.microsoft.com/en-us/windows/wsl/media/ubuntuinstall.png)\n\n## 5. Configure CUDA WSL-Ubuntu Toolkit\n\nBy default, a shimmed version of the CUDA tooling is provided by your Windows GPU drivers.\n\nImportant: you should _never_ use instructions for installing CUDA-toolkit in a generic linux fashion. in WSL 2, you _always_ want to use the provided `CUDA Toolkit using WSL-Ubuntu Package`.\n\nFirst, open PowerShell or Windows Command Prompt in administrator mode\nby right-clicking and selecting \"Run as administrator\".\nThen enter the following command:\n\n```powershell\nwsl.exe\n```\n\nThis should drop you into your running linux VM. Now you can run the following bash commands to install the correct version of cuda-toolkit for WSL-Ubuntu. Note that the version of CUDA used below may not be the version of CUDA your GPU supports.\n\n```sh\nsudo apt-key del 7fa2af80 # if this line fails, you may remove it.\nwget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin\nsudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600\nwget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb\nsudo dpkg -i cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb\nsudo cp /var/cuda-repo-wsl-ubuntu-11-7-local/cuda-B81839D3-keyring.gpg /usr/share/keyrings/\nsudo apt-get update\nsudo apt-get -y install cuda-toolkit-11-7\n```\n\n## 6. Install Docker\n\nDownload and install [Docker Desktop for Windows](https://desktop.docker.com/win/main/amd64/Docker%20Desktop%20Installer.exe). It has WSL 2 support built in by default.\n\nOnce installed, run `Docker Desktop`, you can ignore the first-run tutorial. Go to **Settings → General** and ensure **Use the WSL 2 based engine** has a checkmark next to it. Click **Apply & Restart**.\n\n![\"Use the WSL 2 based engine\" is checked in this interface](images/wsl2-enable.png)\n\nReboot your computer one more time.\n\n## 7. Install `cog` and pull an image\n\nOpen Windows Terminal and enter your WSL 2 VM:\n\n```powershell\nwsl.exe\n```\n\nDownload and install `cog` inside the VM:\n\n```bash\nsudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`\nsudo chmod +x /usr/local/bin/cog\n```\n\nMake sure it's available by typing:\n\n```bash\nwhich cog # should output /usr/local/bin/cog\ncog --version # should output the cog version number.\n```\n\n## 8. Run a model in WSL 2\n\nFinally, make sure it works. Let's try running `afiaka87/glid-3-xl` locally:\n\n```bash\ncog predict 'r8.im/afiaka87/glid-3-xl' -i prompt=\"a fresh avocado floating in the water\" -o prediction.json\n```\n\n![Output from a running cog prediction in Windows Terminal](images/cog_model_output.png)\n\nWhile your prediction is running, you can use `Task Manager` to keep an eye on GPU memory consumption:\n\n![Windows task manager will show the shared host/guest GPU memory](images/memory-usage.png)\n\nThis model just barely manages to fit under 8 GB of VRAM.\n\nNotice that output is returned as JSON for this model as it has a complex return type. You will want to convert the base64 string in the json array to an image.\n\n`jq` can help with this:\n\n```sh\nsudo apt install jq\n```\n\nThe following bash uses `jq` to grab the first element in our prediction array and converts it from a base64 string to a `png` file.\n\n```bash\njq -cs '.[0][0][0]' prediction.json | cut --delimiter \",\" --field 2 | base64 --ignore-garbage --decode > prediction.png\n```\n\nWhen using WSL 2, you can access Windows binaries with the `.exe` extension. This lets you open photos easily within linux.\n\n```bash\nexplorer.exe prediction.png\n```\n\n![a square image of an avocado, generated by the model](images/glide_out.png)\n\n## 9. References\n\n- \n- \n- \n- \n- \n\n\n---\n\n# `cog.yaml` reference\n\n`cog.yaml` defines how to build a Docker image and how to run predictions on your model inside that image.\n\nIt has three keys: [`build`](#build), [`image`](#image), and [`predict`](#predict). It looks a bit like this:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n python_requirements: requirements.txt\n system_packages:\n - \"ffmpeg\"\n - \"git\"\npredict: \"predict.py:Predictor\"\n```\n\nTip: Run [`cog init`](getting-started-own-model.md#initialization) to generate an annotated `cog.yaml` file that can be used as a starting point for setting up your model.\n\n## `build`\n\nThis stanza describes how to build the Docker image your model runs in. It contains various options within it:\n\n\n\n### `cuda`\n\nCog automatically picks the correct version of CUDA to install, but this lets you override it for whatever reason by specifying the minor (`11.8`) or patch (`11.8.0`) version of CUDA to use.\n\nFor example:\n\n```yaml\nbuild:\n cuda: \"11.8\"\n```\n\n### `gpu`\n\nEnable GPUs for this model. When enabled, the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image will be used, and Cog will automatically figure out what versions of CUDA and cuDNN to use based on the version of Python, PyTorch, and Tensorflow that you are using.\n\nFor example:\n\n```yaml\nbuild:\n gpu: true\n```\n\nWhen you use `cog run` or `cog predict`, Cog will automatically pass the `--gpus=all` flag to Docker. When you run a Docker image built with Cog, you'll need to pass this option to `docker run`.\n\n### `python_requirements`\n\nA pip requirements file specifying the Python packages to install. For example:\n\n```yaml\nbuild:\n python_requirements: requirements.txt\n```\n\nYour `cog.yaml` file can set either `python_packages` or `python_requirements`, but not both. Use `python_requirements` when you need to configure options like `--extra-index-url` or `--trusted-host` to fetch Python package dependencies.\n\nThis follows the standard [requirements.txt](https://pip.pypa.io/en/stable/reference/requirements-file-format/) format.\n\nTo install Git-hosted Python packages, add `git` to the `system_packages` list, then use the `git+https://` syntax to specify the package name. For example:\n\n`cog.yaml`:\n\n```yaml\nbuild:\n system_packages:\n - \"git\"\n python_requirements: requirements.txt\n```\n\n`requirements.txt`:\n\n```\ngit+https://github.com/huggingface/transformers\n```\n\nYou can also pin Python package installations to a specific git commit:\n\n`cog.yaml`:\n\n```yaml\nbuild:\n system_packages:\n - \"git\"\n python_requirements: requirements.txt\n```\n\n`requirements.txt`:\n\n```\ngit+https://github.com/huggingface/transformers@2d1602a\n```\n\nNote that you can use a shortened prefix of the 40-character git commit SHA, but you must use at least six characters, like `2d1602a` above.\n\n### `python_packages`\n\n**DEPRECATED**: This will be removed in future versions, please use [python_requirements](#python_requirements) instead.\n\nA list of Python packages to install from the PyPi package index, in the format `package==version`. For example:\n\n```yaml\nbuild:\n python_packages:\n - pillow==8.3.1\n - tensorflow==2.5.0\n```\n\nYour `cog.yaml` file can set either `python_packages` or `python_requirements`, but not both.\n\n### `python_version`\n\nThe minor (`3.13`) or patch (`3.13.1`) version of Python to use. For example:\n\n```yaml\nbuild:\n python_version: \"3.13.1\"\n```\n\nCog supports Python 3.10, 3.11, 3.12, and 3.13. If you don't define a version, Cog will use the latest version of Python 3.13 or a version of Python that is compatible with the versions of PyTorch or TensorFlow you specify.\n\nNote that these are the versions supported **in the Docker container**, not your host machine. You can run any version(s) of Python you wish on your host machine.\n\n### `run`\n\nA list of setup commands to run in the environment after your system packages and Python packages have been installed. If you're familiar with Docker, it's like a `RUN` instruction in your `Dockerfile`.\n\nFor example:\n\n```yaml\nbuild:\n run:\n - curl -L https://github.com/cowsay-org/cowsay/archive/refs/tags/v3.7.0.tar.gz | tar -xzf -\n - cd cowsay-3.7.0 && make install\n```\n\nYour code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally.\n\nEach command in `run` can be either a string or a dictionary in the following format:\n\n```yaml\nbuild:\n run:\n - command: pip install\n mounts:\n - type: secret\n id: pip\n target: /etc/pip.conf\n```\n\nYou can use secret mounts to securely pass credentials to setup commands, without baking them into the image. For more information, see [Dockerfile reference](https://docs.docker.com/engine/reference/builder/#run---mounttypesecret).\n\n### `sdk_version`\n\nPin the version of the cog Python SDK installed in the container. Accepts a [PEP 440](https://peps.python.org/pep-0440/) version string. When omitted, the latest release is installed.\n\n```yaml\nbuild:\n python_version: \"3.13\"\n sdk_version: \"0.18.0\"\n```\n\nPre-release versions are also supported:\n\n```yaml\nbuild:\n sdk_version: \"0.18.0a1\"\n```\n\nWhen a pre-release `sdk_version` is set, `--pre` is automatically passed to the pip install commands for both `cog` and `coglet`, so pip will resolve matching pre-release packages.\n\nThe minimum supported version is `0.16.0`. Specifying an older version will cause `cog build` to fail with an error.\n\nThe `COG_SDK_WHEEL` environment variable takes precedence over `sdk_version`. See [Environment variables](./environment.md) for details.\n\n### `system_packages`\n\nA list of Ubuntu APT packages to install. For example:\n\n```yaml\nbuild:\n system_packages:\n - \"ffmpeg\"\n - \"libavcodec-dev\"\n```\n\n## `concurrency`\n\n> Added in cog 0.14.0.\n\nThis stanza describes the concurrency capabilities of the model. It has one option:\n\n### `max`\n\nThe maximum number of concurrent predictions the model can process. If this is set, the model must specify an [async `predict()` method](python.md#async-predictors-and-concurrency).\n\nFor example:\n\n```yaml\nconcurrency:\n max: 10\n```\n\n## `image`\n\nThe name given to built Docker images. If you want to push to a registry, this should also include the registry name.\n\nFor example:\n\n```yaml\nimage: \"r8.im/your-username/your-model\"\n```\n\nr8.im is Replicate's registry, but this can be any Docker registry.\n\nIf you don't set this, then a name will be generated from the directory name.\n\nIf you set this, then you can run `cog push` without specifying the model name.\n\nIf you specify an image name argument when pushing (like `cog push your-username/custom-model-name`), the argument will be used and the value of `image` in cog.yaml will be ignored.\n\n## `predict`\n\nThe pointer to the `Predictor` object in your code, which defines how predictions are run on your model.\n\nFor example:\n\n```yaml\npredict: \"predict.py:Predictor\"\n```\n\nSee [the Python API documentation for more information](python.md).\n" }, { "path": "docs/notebooks.md", "content": "# Notebooks\n\nCog plays nicely with Jupyter notebooks.\n\n## Install the jupyterlab Python package\n\nFirst, add `jupyterlab` to your `requirements.txt` file and reference it in [`cog.yaml`](yaml.md):\n\n`requirements.txt`:\n\n```\njupyterlab\n```\n\n`cog.yaml`:\n\n```yaml\nbuild:\n python_requirements: requirements.txt\n```\n\n## Run a notebook\n\nCog can run notebooks in the environment you've defined in `cog.yaml` with the following command:\n\n```sh\ncog run -p 8888 jupyter lab --allow-root --ip=0.0.0.0\n```\n\n## Use notebook code in your predictor\n\nYou can also import a notebook into your Cog [Predictor](python.md) file.\n\nFirst, export your notebook to a Python file:\n\n```sh\njupyter nbconvert --to script my_notebook.ipynb # creates my_notebook.py\n```\n\nThen import the exported Python script into your `predict.py` file. Any functions or variables defined in your notebook will be available to your predictor:\n\n```python\nfrom cog import BasePredictor, Input\n\nimport my_notebook\n\nclass Predictor(BasePredictor):\n def predict(self, prompt: str = Input(description=\"string prompt\")) -> str:\n output = my_notebook.do_stuff(prompt)\n return output\n```\n" }, { "path": "docs/private-package-registry.md", "content": "# Private package registry\n\nThis guide describes how to build a Docker image with Cog that fetches Python packages from a private registry during setup.\n\n## `pip.conf`\n\nIn a directory outside your Cog project, create a `pip.conf` file with an `index-url` set to the registry's URL with embedded credentials.\n\n```conf\n[global]\nindex-url = https://username:password@my-private-registry.com\n```\n\n> **Warning**\n> Be careful not to commit secrets in Git or include them in Docker images. If your Cog project contains any sensitive files, make sure they're listed in `.gitignore` and `.dockerignore`.\n\n## `cog.yaml`\n\nIn your project's [`cog.yaml`](yaml.md) file, add a setup command to run `pip install` with a secret configuration file mounted to `/etc/pip.conf`.\n\n```yaml\nbuild:\n run:\n - command: pip install\n mounts:\n - type: secret\n id: pip\n target: /etc/pip.conf\n```\n\n## Build\n\nWhen building or pushing your model with Cog, pass the `--secret` option with an `id` matching the one specified in `cog.yaml`, along with a path to your local `pip.conf` file.\n\n```console\n$ cog build --secret id=pip,source=/path/to/pip.conf\n```\n\nUsing a secret mount allows the private registry credentials to be securely passed to the `pip install` setup command, without baking them into the Docker image.\n\n> **Warning**\n> If you run `cog build` or `cog push` and then change the contents of a secret source file, the cached version of the file will be used on subsequent builds, ignoring any changes you made. To update the contents of the target secret file, either change the `id` value in `cog.yaml` and the `--secret` option, or pass the `--no-cache` option to bypass the cache entirely.\n" }, { "path": "docs/python.md", "content": "# Prediction interface reference\n\nThis document defines the API of the `cog` Python module, which is used to define the interface for running predictions on your model.\n\n> [!TIP]\n> Run [`cog init`](getting-started-own-model.md#initialization) to generate an annotated `predict.py` file that can be used as a starting point for setting up your model.\n\n> [!TIP]\n> Using a language model to help you write the code for your new Cog model?\n>\n> Feed it [https://cog.run/llms.txt](https://cog.run/llms.txt), which has all of Cog's documentation bundled into a single file. To learn more about this format, check out [llmstxt.org](https://llmstxt.org).\n\n## Contents\n\n- [Contents](#contents)\n- [`BasePredictor`](#basepredictor)\n - [`Predictor.setup()`](#predictorsetup)\n - [`Predictor.predict(**kwargs)`](#predictorpredictkwargs)\n- [`async` predictors and concurrency](#async-predictors-and-concurrency)\n- [`Input(**kwargs)`](#inputkwargs)\n - [Deprecating inputs](#deprecating-inputs)\n- [Output](#output)\n - [Returning an object](#returning-an-object)\n - [Returning a list](#returning-a-list)\n - [Optional properties](#optional-properties)\n - [Streaming output](#streaming-output)\n- [Metrics](#metrics)\n - [Recording metrics](#recording-metrics)\n - [Accumulation modes](#accumulation-modes)\n - [Dot-path keys](#dot-path-keys)\n - [Type safety](#type-safety)\n- [Cancellation](#cancellation)\n - [`CancelationException`](#cancelationexception)\n- [Input and output types](#input-and-output-types)\n- [`File()`](#file)\n- [`Path()`](#path)\n- [`Secret`](#secret)\n- [`Optional`](#optional)\n- [`List`](#list)\n\n## `BasePredictor`\n\nYou define how Cog runs predictions on your model by defining a class that inherits from `BasePredictor`. It looks something like this:\n\n```python\nfrom cog import BasePredictor, Path, Input\nimport torch\n\nclass Predictor(BasePredictor):\n def setup(self):\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n self.model = torch.load(\"weights.pth\")\n\n def predict(self,\n image: Path = Input(description=\"Image to enlarge\"),\n scale: float = Input(description=\"Factor to scale image by\", default=1.5)\n ) -> Path:\n \"\"\"Run a single prediction on the model\"\"\"\n # ... pre-processing ...\n output = self.model(image)\n # ... post-processing ...\n return output\n```\n\nYour Predictor class should define two methods: `setup()` and `predict()`.\n\n### `Predictor.setup()`\n\nPrepare the model so multiple predictions run efficiently.\n\nUse this _optional_ method to include expensive one-off operations like loading trained models, instantiating data transformations, etc.\n\nMany models use this method to download their weights (e.g. using [`pget`](https://github.com/replicate/pget)). This has some advantages:\n\n- Smaller image sizes\n- Faster build times\n- Faster pushes and inference on [Replicate](https://replicate.com)\n\nHowever, this may also significantly increase your `setup()` time.\n\nAs an alternative, some choose to store their weights directly in the image. You can simply leave your weights in the directory alongside your `cog.yaml` and ensure they are not excluded in your `.dockerignore` file.\n\nWhile this will increase your image size and build time, it offers other advantages:\n\n- Faster `setup()` time\n- Ensures idempotency and reduces your model's reliance on external systems\n- Preserves reproducibility as your model will be self-contained in the image\n\n> When using this method, you should use the `--separate-weights` flag on `cog build` to store weights in a [separate layer](https://github.com/replicate/cog/blob/12ac02091d93beebebed037f38a0c99cd8749806/docs/getting-started.md?plain=1#L219).\n\n### `Predictor.predict(**kwargs)`\n\nRun a single prediction.\n\nThis _required_ method is where you call the model that was loaded during `setup()`, but you may also want to add pre- and post-processing code here.\n\nThe `predict()` method takes an arbitrary list of named arguments, where each argument name must correspond to an [`Input()`](#inputkwargs) annotation.\n\n`predict()` can return strings, numbers, [`cog.Path`](#path) objects representing files on disk, or lists or dicts of those types. You can also define a custom [`Output()`](#outputbasemodel) for more complex return types.\n\n## `async` predictors and concurrency\n\n> Added in cog 0.14.0.\n\nYou may specify your `predict()` method as `async def predict(...)`. In\naddition, if you have an async `predict()` function you may also have an async\n`setup()` function:\n\n```py\nclass Predictor(BasePredictor):\n async def setup(self) -> None:\n print(\"async setup is also supported...\")\n\n async def predict(self) -> str:\n print(\"async predict\");\n return \"hello world\";\n```\n\nModels that have an async `predict()` function can run predictions concurrently, up to the limit specified by [`concurrency.max`](yaml.md#max) in cog.yaml. Attempting to exceed this limit will return a 409 Conflict response.\n\n## `Input(**kwargs)`\n\nUse cog's `Input()` function to define each of the parameters in your `predict()` method:\n\n```py\nclass Predictor(BasePredictor):\n def predict(self,\n image: Path = Input(description=\"Image to enlarge\"),\n scale: float = Input(description=\"Factor to scale image by\", default=1.5, ge=1.0, le=10.0)\n ) -> Path:\n```\n\nThe `Input()` function takes these keyword arguments:\n\n- `description`: A description of what to pass to this input for users of the model.\n- `default`: A default value to set the input to. If this argument is not passed, the input is required. If it is explicitly set to `None`, the input is optional.\n- `ge`: For `int` or `float` types, the value must be greater than or equal to this number.\n- `le`: For `int` or `float` types, the value must be less than or equal to this number.\n- `min_length`: For `str` types, the minimum length of the string.\n- `max_length`: For `str` types, the maximum length of the string.\n- `regex`: For `str` types, the string must match this regular expression.\n- `choices`: For `str` or `int` types, a list of possible values for this input.\n- `deprecated`: (optional) If set to `True`, marks this input as deprecated. Deprecated inputs will still be accepted, but tools and UIs may warn users that the input is deprecated and may be removed in the future. See [Deprecating inputs](#deprecating-inputs).\n\nEach parameter of the `predict()` method must be annotated with a type like `str`, `int`, `float`, `bool`, etc. See [Input and output types](#input-and-output-types) for the full list of supported types.\n\nUsing the `Input` function provides better documentation and validation constraints to the users of your model, but it is not strictly required. You can also specify default values for your parameters using plain Python, or omit default assignment entirely:\n\n```py\nclass Predictor(BasePredictor):\n def predict(self,\n prompt: str = \"default prompt\", # this is valid\n iterations: int # also valid\n ) -> str:\n # ...\n```\n\n## Deprecating inputs\n\nYou can mark an input as deprecated by passing `deprecated=True` to the `Input()` function. Deprecated inputs will still be accepted, but tools and UIs may warn users that the input is deprecated and may be removed in the future.\n\nThis is useful when you want to phase out an input without breaking existing clients immediately:\n\n```py\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n def predict(self,\n text: str = Input(description=\"Some deprecated text\", deprecated=True),\n prompt: str = Input(description=\"Prompt for the model\")\n ) -> str:\n # ...\n return prompt\n```\n\n## Output\n\nCog predictors can return a simple data type like a string, number, float, or boolean. Use Python's `-> ` syntax to annotate the return type.\n\nHere's an example of a predictor that returns a string:\n\n```py\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def predict(self) -> str:\n return \"hello\"\n```\n\n### Returning an object\n\nTo return a complex object with multiple values, define an `Output` object with multiple fields to return from your `predict()` method:\n\n```py\nfrom cog import BasePredictor, BaseModel, File\n\nclass Output(BaseModel):\n file: File\n text: str\n\nclass Predictor(BasePredictor):\n def predict(self) -> Output:\n return Output(text=\"hello\", file=io.StringIO(\"hello\"))\n```\n\nEach of the output object's properties must be one of the supported output types. For the full list, see [Input and output types](#input-and-output-types). Also, make sure to name the output class as `Output` and nothing else.\n\n### Returning a list\n\nThe `predict()` method can return a list of any of the supported output types. Here's an example that outputs multiple files:\n\n```py\nfrom cog import BasePredictor, Path\n\nclass Predictor(BasePredictor):\n def predict(self) -> list[Path]:\n predictions = [\"foo\", \"bar\", \"baz\"]\n output = []\n for i, prediction in enumerate(predictions):\n out_path = Path(f\"/tmp/out-{i}.txt\")\n with out_path.open(\"w\") as f:\n f.write(prediction)\n output.append(out_path)\n return output\n```\n\nFiles are named in the format `output..`, e.g. `output.0.txt`, `output.1.txt`, and `output.2.txt` from the example above.\n\n### Optional properties\n\nTo conditionally omit properties from the Output object, define them using `typing.Optional`:\n\n```py\nfrom cog import BaseModel, BasePredictor, Path\nfrom typing import Optional\n\nclass Output(BaseModel):\n score: Optional[float]\n file: Optional[Path]\n\nclass Predictor(BasePredictor):\n def predict(self) -> Output:\n if condition:\n return Output(score=1.5)\n else:\n return Output(file=io.StringIO(\"hello\"))\n```\n\n### Streaming output\n\nCog models can stream output as the `predict()` method is running. For example, a language model can output tokens as they're being generated and an image generation model can output images as they are being generated.\n\nTo support streaming output in your Cog model, add `from typing import Iterator` to your predict.py file. The `typing` package is a part of Python's standard library so it doesn't need to be installed. Then add a return type annotation to the `predict()` method in the form `-> Iterator[]` where `` can be one of `str`, `int`, `float`, `bool`, or `cog.Path`.\n\n```py\nfrom cog import BasePredictor, Path\nfrom typing import Iterator\n\nclass Predictor(BasePredictor):\n def predict(self) -> Iterator[Path]:\n done = False\n while not done:\n output_path, done = do_stuff()\n yield Path(output_path)\n```\n\nIf you have an [async `predict()` method](#async-predictors-and-concurrency), you must use `cog.AsyncIterator` instead:\n\n```py\nfrom cog import AsyncIterator, BasePredictor, Path\n\nclass Predictor(BasePredictor):\n async def predict(self) -> AsyncIterator[Path]:\n done = False\n while not done:\n output_path, done = do_stuff()\n yield Path(output_path)\n```\n\nIf you're streaming text output, you can use `ConcatenateIterator` to hint that the output should be concatenated together into a single string. This is useful on Replicate to display the output as a string instead of a list of strings.\n\n```py\nfrom cog import BasePredictor, Path, ConcatenateIterator\n\nclass Predictor(BasePredictor):\n def predict(self) -> ConcatenateIterator[str]:\n tokens = [\"The\", \"quick\", \"brown\", \"fox\", \"jumps\", \"over\", \"the\", \"lazy\", \"dog\"]\n for token in tokens:\n yield token + \" \"\n```\n\nOr for async `predict()` methods, use `AsyncConcatenateIterator`:\n\n```py\nfrom cog import BasePredictor, Path, AsyncConcatenateIterator\n\nclass Predictor(BasePredictor):\n async def predict(self) -> AsyncConcatenateIterator[str]:\n tokens = [\"The\", \"quick\", \"brown\", \"fox\", \"jumps\", \"over\", \"the\", \"lazy\", \"dog\"]\n for token in tokens:\n yield token + \" \"\n```\n\n## Metrics\n\nYou can record custom metrics from your `predict()` function to track model-specific data like token counts, timing breakdowns, or confidence scores. Metrics are included in the prediction response alongside the output.\n\n### Recording metrics\n\nUse `self.record_metric()` inside your `predict()` method:\n\n```python\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def predict(self, prompt: str) -> str:\n self.record_metric(\"temperature\", 0.7)\n self.record_metric(\"token_count\", 42)\n\n result = self.model.generate(prompt)\n return result\n```\n\nFor advanced use (dict-style access, deleting metrics), use `self.scope`:\n\n```python\nself.scope.metrics[\"token_count\"] = 42\ndel self.scope.metrics[\"token_count\"]\n```\n\nMetrics appear in the prediction response `metrics` field:\n\n```json\n{\n \"status\": \"succeeded\",\n \"output\": \"...\",\n \"metrics\": {\n \"temperature\": 0.7,\n \"token_count\": 42,\n \"predict_time\": 1.23\n }\n}\n```\n\nThe `predict_time` metric is always added automatically by the runtime. If you set `predict_time` yourself, the runtime value takes precedence.\n\nSupported value types are `bool`, `int`, `float`, `str`, `list`, and `dict`. Setting a metric to `None` deletes it.\n\n### Accumulation modes\n\nBy default, recording a metric replaces any previous value for that key. You can use accumulation modes to build up values across multiple calls:\n\n```python\n# Increment a counter (adds to the existing numeric value)\nself.record_metric(\"token_count\", 1, mode=\"incr\")\nself.record_metric(\"token_count\", 1, mode=\"incr\")\n# Result: {\"token_count\": 2}\n\n# Append to an array\nself.record_metric(\"steps\", \"preprocessing\", mode=\"append\")\nself.record_metric(\"steps\", \"inference\", mode=\"append\")\n# Result: {\"steps\": [\"preprocessing\", \"inference\"]}\n\n# Replace (default behavior)\nself.record_metric(\"status\", \"running\", mode=\"replace\")\nself.record_metric(\"status\", \"done\", mode=\"replace\")\n# Result: {\"status\": \"done\"}\n```\n\nThe `mode` parameter accepts `\"replace\"` (default), `\"incr\"`, or `\"append\"`.\n\n### Dot-path keys\n\nUse dot-separated keys to create nested objects in the metrics output:\n\n```python\nself.record_metric(\"timing.preprocess\", 0.12)\nself.record_metric(\"timing.inference\", 0.85)\n```\n\nThis produces nested JSON:\n\n```json\n{\n \"metrics\": {\n \"timing\": {\n \"preprocess\": 0.12,\n \"inference\": 0.85\n },\n \"predict_time\": 1.23\n }\n}\n```\n\n### Type safety\n\nOnce a metric key has been assigned a value of a certain type, it cannot be changed to a different type without deleting it first. This prevents accidental type mismatches when using accumulation modes:\n\n```python\nself.record_metric(\"count\", 1)\n\n# This would raise an error — \"count\" is an int, not a string:\n# self.record_metric(\"count\", \"oops\")\n\n# Delete first, then set with new type:\ndel self.scope.metrics[\"count\"]\nself.record_metric(\"count\", \"now a string\")\n```\n\nOutside an active prediction, `self.record_metric()` and `self.scope` are silent no-ops — no need for `None` checks.\n\n## Cancellation\n\nWhen a prediction is canceled (via the [cancel HTTP endpoint](http.md#post-predictionsprediction_idcancel) or a dropped connection), the Cog runtime interrupts the running `predict()` function. The exception raised depends on whether the predictor is sync or async:\n\n| Predictor type | Exception raised |\n| --------------------------- | ------------------------ |\n| Sync (`def predict`) | `CancelationException` |\n| Async (`async def predict`) | `asyncio.CancelledError` |\n\n### `CancelationException`\n\n```python\nfrom cog import CancelationException\n```\n\n`CancelationException` is raised in **sync** predictors when a prediction is cancelled. It is a `BaseException` subclass — **not** an `Exception` subclass. This means bare `except Exception` blocks in your predict code will not accidentally catch it, matching the behavior of `KeyboardInterrupt` and `asyncio.CancelledError`.\n\nYou do **not** need to handle this exception in normal predictor code — the runtime manages cancellation automatically. However, if you need to run cleanup logic when a prediction is cancelled, you can catch it explicitly:\n\n```python\nfrom cog import BasePredictor, CancelationException, Path\n\nclass Predictor(BasePredictor):\n def predict(self, image: Path) -> Path:\n try:\n return self.process(image)\n except CancelationException:\n self.cleanup()\n raise # always re-raise\n```\n\n> [!WARNING]\n> You **must** re-raise `CancelationException` after cleanup. Swallowing it will prevent the runtime from marking the prediction as canceled, and may result in the termination of the container.\n\n`CancelationException` is available as:\n\n- `cog.CancelationException` (recommended)\n- `cog.exceptions.CancelationException`\n\nFor **async** predictors, cancellation follows standard Python async conventions and raises `asyncio.CancelledError` instead.\n\n## Input and output types\n\nEach parameter of the `predict()` method must be annotated with a type. The method's return type must also be annotated. The supported types are:\n\n- `str`: a string\n- `int`: an integer\n- `float`: a floating point number\n- `bool`: a boolean\n- [`cog.File`](#file): a file-like object representing a file\n- [`cog.Path`](#path): a path to a file on disk\n- [`cog.Secret`](#secret): a string containing sensitive information\n\n## `File()`\n\n> [!WARNING] \n> `cog.File` is deprecated and will be removed in a future version of Cog. Use [`cog.Path`](#path) instead.\n\nThe `cog.File` object is used to get files in and out of models. It represents a _file handle_.\n\nFor models that return a `cog.File` object, the prediction output returned by Cog's built-in HTTP server will be a URL.\n\n```python\nfrom cog import BasePredictor, File, Input, Path\nfrom PIL import Image\n\nclass Predictor(BasePredictor):\n def predict(self, source_image: File = Input(description=\"Image to enlarge\")) -> File:\n pillow_img = Image.open(source_image)\n upscaled_image = do_some_processing(pillow_img)\n return File(upscaled_image)\n```\n\n## `Path()`\n\nThe `cog.Path` object is used to get files in and out of models. It represents a _path to a file on disk_.\n\n`cog.Path` is a subclass of Python's [`pathlib.Path`](https://docs.python.org/3/library/pathlib.html#basic-use) and can be used as a drop-in replacement.\n\nFor models that return a `cog.Path` object, the prediction output returned by Cog's built-in HTTP server will be a URL.\n\nThis example takes an input file, resizes it, and returns the resized image:\n\n```python\nimport tempfile\nfrom cog import BasePredictor, Input, Path\n\nclass Predictor(BasePredictor):\n def predict(self, image: Path = Input(description=\"Image to enlarge\")) -> Path:\n upscaled_image = do_some_processing(image)\n\n # To output `cog.Path` objects the file needs to exist, so create a temporary file first.\n # This file will automatically be deleted by Cog after it has been returned.\n output_path = Path(tempfile.mkdtemp()) / \"upscaled.png\"\n upscaled_image.save(output_path)\n return Path(output_path)\n```\n\n## `Secret`\n\nThe `cog.Secret` type is used to signify that an input holds sensitive information,\nlike a password or API token.\n\n`cog.Secret` is a type that redacts its contents in string representations to prevent accidental disclosure.\nYou can access its contents with the `get_secret_value()` method.\n\n```python\nfrom cog import BasePredictor, Secret\n\n\nclass Predictor(BasePredictor):\n def predict(self, api_token: Secret) -> None:\n # Prints '**********'\n print(api_token)\n\n # Use get_secret_value method to see the secret's content.\n print(api_token.get_secret_value())\n```\n\nA predictor's `Secret` inputs are represented in OpenAPI with the following schema:\n\n```json\n{\n \"type\": \"string\",\n \"format\": \"password\",\n \"x-cog-secret\": true\n}\n```\n\nModels uploaded to Replicate treat secret inputs differently throughout its system.\nWhen you create a prediction on Replicate,\nany value passed to a `Secret` input is redacted after being sent to the model.\n\n> [!WARNING] \n> Passing secret values to untrusted models can result in\n> unintended disclosure, exfiltration, or misuse of sensitive data.\n\n## `Optional`\n\nOptional inputs should be explicitly defined as `Optional[T]` so that type checker can warn us about error-prone `None` values.\n\nFor example, the following code might fail if `prompt` is not specified in the inputs:\n\n```python\nclass Predictor(BasePredictor):\n def predict(self, prompt: str=Input(description=\"prompt\", default=None)) -> str:\n return \"hello\" + prompt # TypeError: can only concatenate str (not \"NoneType\") to str\n```\n\nWe can improve it by making `prompt` an `Optional[str]`. Note that `default=None` is now redundant as `Optional` implies it.\n\n```python\nclass Predictor(BasePredictor):\n def predict(self, prompt: Optional[str]=Input(description=\"prompt\")) -> str:\n if prompt is None: # type check can warn us if we forget this\n return \"hello\"\n else:\n return \"hello\" + prompt\n```\n\nNote that the error prone usage of `prompt: str=Input(default=None)` might throw an error in a future release of Cog.\n\n## `List`\n\nThe List type is also supported in inputs. It can hold any supported type.\n\nExample for **List[Path]**:\n\n```py\nclass Predictor(BasePredictor):\n def predict(self, paths: list[Path]) -> str:\n output_parts = [] # Use a list to collect file contents\n for path in paths:\n with open(path) as f:\n output_parts.append(f.read())\n return \"\".join(output_parts)\n```\n\nThe corresponding cog command:\n\n```bash\n$ echo test1 > 1.txt\n$ echo test2 > 2.txt\n$ cog predict -i paths=@1.txt -i paths=@2.txt\nRunning prediction...\ntest1\n\ntest2\n```\n\n- Note the repeated inputs with the same name \"paths\" which constitute the list\n" }, { "path": "docs/stylesheets/extra.css", "content": ".md-typeset h1,\n.md-typeset h2,\n.md-typeset h3 {\n font-weight: 600;\n}\n\n/* move the \"Cog\" header to the left */\n[dir=\"ltr\"] .md-header__title {\n margin-left: 0;\n}\n\n/* Remove the superfluous \"Cog\" label above the TOC */\n.md-nav__title {\n display: none;\n}\n" }, { "path": "docs/training.md", "content": "# Training interface reference\n\n> [!WARNING] \n> The `cog train` command is deprecated and will be removed in the next version of Cog. The training API described below may still be used with the HTTP API's `/trainings` endpoint, but the CLI command is no longer recommended for new projects.\n\nCog's training API allows you to define a fine-tuning interface for an existing Cog model, so users of the model can bring their own training data to create derivative fine-tuned models. Real-world examples of this API in use include [fine-tuning SDXL with images](https://replicate.com/blog/fine-tune-sdxl) or [fine-tuning Llama 2 with structured text](https://replicate.com/blog/fine-tune-llama-2).\n\n## How it works\n\nIf you've used Cog before, you've probably seen the [Predictor](./python.md) class, which defines the interface for creating predictions against your model. Cog's training API works similarly: You define a Python function that describes the inputs and outputs of the training process. The inputs are things like training data, epochs, batch size, seed, etc. The output is typically a file with the fine-tuned weights.\n\n`cog.yaml`:\n\n```yaml\nbuild:\n python_version: \"3.13\"\ntrain: \"train.py:train\"\n```\n\n`train.py`:\n\n```python\nfrom cog import BasePredictor, File\nimport io\n\ndef train(param: str) -> File:\n return io.StringIO(\"hello \" + param)\n```\n\nThen you can run it like this:\n\n```\n$ cog train -i param=train\n...\n\n$ cat weights\nhello train\n```\n\nYou can also use classes if you want to run many model trainings and save on setup time. This works the same way as the [Predictor](./python.md) class with the only difference being the `train` method.\n\n`cog.yaml`:\n\n```yaml\nbuild:\n python_version: \"3.13\"\ntrain: \"train.py:Trainer\"\n```\n\n`train.py`:\n\n```python\nfrom cog import BasePredictor, File\nimport io\n\nclass Trainer:\n def setup(self) -> None:\n self.base_model = ... # Load a big base model\n\n def train(self, param: str) -> File:\n return self.base_model.train(param) # Train on top of a base model\n```\n\n## `Input(**kwargs)`\n\nUse Cog's `Input()` function to define each of the parameters in your `train()` function:\n\n```py\nfrom cog import Input, Path\n\ndef train(\n train_data: Path = Input(description=\"HTTPS URL of a file containing training data\"),\n learning_rate: float = Input(description=\"learning rate, for learning!\", default=1e-4, ge=0),\n seed: int = Input(description=\"random seed to use for training\", default=None)\n) -> str:\n return \"hello, weights\"\n```\n\nThe `Input()` function takes these keyword arguments:\n\n- `description`: A description of what to pass to this input for users of the model.\n- `default`: A default value to set the input to. If this argument is not passed, the input is required. If it is explicitly set to `None`, the input is optional.\n- `ge`: For `int` or `float` types, the value must be greater than or equal to this number.\n- `le`: For `int` or `float` types, the value must be less than or equal to this number.\n- `min_length`: For `str` types, the minimum length of the string.\n- `max_length`: For `str` types, the maximum length of the string.\n- `regex`: For `str` types, the string must match this regular expression.\n- `choices`: For `str` or `int` types, a list of possible values for this input.\n\nEach parameter of the `train()` function must be annotated with a type like `str`, `int`, `float`, `bool`, etc. See [Input and output types](./python.md#input-and-output-types) for the full list of supported types.\n\nUsing the `Input` function provides better documentation and validation constraints to the users of your model, but it is not strictly required. You can also specify default values for your parameters using plain Python, or omit default assignment entirely:\n\n```py\ndef train(self,\n training_data: str = \"foo bar\", # this is valid\n iterations: int # also valid\n) -> str:\n # ...\n```\n\n## Training Output\n\nTraining output is typically a binary weights file. To return a custom output object or a complex object with multiple values, define a `TrainingOutput` object with multiple fields to return from your `train()` function, and specify it as the return type for the train function using Python's `->` return type annotation:\n\n```python\nfrom cog import BaseModel, Input, Path\n\nclass TrainingOutput(BaseModel):\n weights: Path\n\ndef train(\n train_data: Path = Input(description=\"HTTPS URL of a file containing training data\"),\n learning_rate: float = Input(description=\"learning rate, for learning!\", default=1e-4, ge=0),\n seed: int = Input(description=\"random seed to use for training\", default=42)\n) -> TrainingOutput:\n weights_file = generate_weights(\"...\")\n return TrainingOutput(weights=Path(weights_file))\n```\n\n## Testing\n\nIf you are doing development of a Cog model like Llama or SDXL, you can test that the fine-tuned code path works before pushing by specifying a `COG_WEIGHTS` environment variable when running `predict`:\n\n```console\ncog predict -e COG_WEIGHTS=https://replicate.delivery/pbxt/xyz/weights.tar -i prompt=\"a photo of TOK\"\n```\n" }, { "path": "docs/wsl2/wsl2.md", "content": "# Using `cog` on Windows 11 with WSL 2\n\n- [0. Prerequisites](#0-prerequisites)\n- [1. Install the GPU driver](#1-install-the-gpu-driver)\n- [2. Unlocking features](#2-unlocking-features)\n - [2.1. Unlock WSL2](#21-unlock-wsl2)\n - [2.2. Unlock virtualization](#22-unlock-virtualization)\n - [2.3. Reboot](#23-reboot)\n- [3. Update MS Linux kernel](#3-update-ms-linux-kernel)\n- [4. Configure WSL 2](#4-configure-wsl-2)\n- [5. Configure CUDA WSL-Ubuntu Toolkit](#5-configure-cuda-wsl-ubuntu-toolkit)\n- [6. Install Docker](#6-install-docker)\n- [7. Install `cog` and pull an image](#7-install-cog-and-pull-an-image)\n- [8. Run a model in WSL 2](#8-run-a-model-in-wsl-2)\n- [9. References](#9-references)\n\nRunning cog on Windows is now possible thanks to WSL 2. Follow this guide to enable WSL 2 and GPU passthrough on Windows 11.\n\n**Windows 10 is not officially supported, as you need to be on an insider build in order to use GPU passthrough.**\n\n## 0. Prerequisites\n\nBefore beginning installation, make sure you have:\n\n- Windows 11.\n- NVIDIA GPU.\n - RTX 2000/3000 series\n - Kesler/Tesla/Volta/Ampere series\n - Other configurations are not guaranteed to work.\n\n## 1. Install the GPU driver\n\nPer NVIDIA, the first order of business is to install the latest Game Ready drivers for your NVIDIA GPU.\n\n\n\nI have an NVIDIA RTX 2070 Super, so filled out the form as such:\n\n![a form showing the correct model number selected for an RTX 2070 Super](images/nvidia_driver_select.png)\n\nClick \"search\", and follow the dialogue to download and install the driver.\n\nRestart your computer once the driver has finished installation.\n\n## 2. Unlocking features\n\nOpen Windows Terminal as an administrator.\n\n- Use start to search for \"Terminal\"\n- Right click -> Run as administrator...\n\nRun the following powershell command to enable the Windows Subsystem for Linux and Virtual Machine Platform capabilities.\n\n### 2.1. Unlock WSL2\n\n```powershell\ndism.exe /online /enable-feature /featurename:Microsoft-Windows-Subsystem-Linux /all /norestart\n```\n\nIf you see an error about permissions, make sure the terminal you are using is run as an administrator and that you have an account with administrator-level privileges.\n\n### 2.2. Unlock virtualization\n\n```powershell\ndism.exe /online /enable-feature /featurename:VirtualMachinePlatform /all /norestart\n```\n\nIf this command fails, make sure to [enable virtualization capabilities](https://docs.microsoft.com/en-us/windows/wsl/troubleshooting#error-0x80370102-the-virtual-machine-could-not-be-started-because-a-required-feature-is-not-installed) in your computer's BIOS/UEFI. A successful output will print `The operation completed successfully.`\n\n![Output from running the above commands successfully. Should read \"The operation completed successfully\".](images/enable_feature_success.png)\n\n### 2.3. Reboot\n\nBefore moving forward, make sure you reboot your computer so that Windows 11 will have WSL2 and virtualization available to it.\n\n## 3. Update MS Linux kernel\n\nDownload and run the [WSL2 Linux kernel update package for x64 machines](https://wslstorestorage.blob.core.windows.net/wslblob/wsl_update_x64.msi) msi installer. When prompted for elevated permissions, click 'yes' to approve the installation.\n\nTo ensure you are using the correct WSL kernel, `open Windows Terminal as an administrator` and enter:\n\n```powershell\nwsl cat /proc/version\n```\n\nThis will return a complicated string such as:\n\n```sh\nLinux version 5.10.102.1-microsoft-standard-WSL2 (oe-user@oe-host) (x86_64-msft-linux-gcc (GCC) 9.3.0, GNU ld (GNU Binutils) 2.34.0.20200220)\n```\n\nThe version we are interested in is `Linux version 5.10.102.1`. At this point, you should have updated your kernel to be at least `Linux version 5.10.43.3`.\n\nIf you can't get the correct kernel version to show:\n\nOpen `Settings` → `Windows Update` → `Advanced options` and ensure `Receive updates for other Microsoft products` is enabled. Then go to `Windows Update` again and click `Check for updates`.\n\n## 4. Configure WSL 2\n\nFirst, configure Windows to use the virtualization-based version of WSL (version 2) by default. In a Windows Terminal with administrator privileges, type the following:\n\n```powershell\nwsl --set-default-version 2\n```\n\nNow, you will need to go to the Microsoft Store and [Download Ubuntu 18.04](https://www.microsoft.com/store/apps/9N9TNGVNDL3Q)\n\n![Screenshot showing the \"Ubuntu\" store page](https://docs.microsoft.com/en-us/windows/wsl/media/ubuntustore.png)\n\nLaunch the \"Ubuntu\" app available in your Start Menu. Linux will require its own user account and password, which you will need to enter now:\n\n![a terminal showing input for user account info on WSL 2](https://docs.microsoft.com/en-us/windows/wsl/media/ubuntuinstall.png)\n\n## 5. Configure CUDA WSL-Ubuntu Toolkit\n\nBy default, a shimmed version of the CUDA tooling is provided by your Windows GPU drivers.\n\nImportant: you should _never_ use instructions for installing CUDA-toolkit in a generic linux fashion. in WSL 2, you _always_ want to use the provided `CUDA Toolkit using WSL-Ubuntu Package`.\n\nFirst, open PowerShell or Windows Command Prompt in administrator mode\nby right-clicking and selecting \"Run as administrator\".\nThen enter the following command:\n\n```powershell\nwsl.exe\n```\n\nThis should drop you into your running linux VM. Now you can run the following bash commands to install the correct version of cuda-toolkit for WSL-Ubuntu. Note that the version of CUDA used below may not be the version of CUDA your GPU supports.\n\n```sh\nsudo apt-key del 7fa2af80 # if this line fails, you may remove it.\nwget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin\nsudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600\nwget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb\nsudo dpkg -i cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb\nsudo cp /var/cuda-repo-wsl-ubuntu-11-7-local/cuda-B81839D3-keyring.gpg /usr/share/keyrings/\nsudo apt-get update\nsudo apt-get -y install cuda-toolkit-11-7\n```\n\n## 6. Install Docker\n\nDownload and install [Docker Desktop for Windows](https://desktop.docker.com/win/main/amd64/Docker%20Desktop%20Installer.exe). It has WSL 2 support built in by default.\n\nOnce installed, run `Docker Desktop`, you can ignore the first-run tutorial. Go to **Settings → General** and ensure **Use the WSL 2 based engine** has a checkmark next to it. Click **Apply & Restart**.\n\n![\"Use the WSL 2 based engine\" is checked in this interface](images/wsl2-enable.png)\n\nReboot your computer one more time.\n\n## 7. Install `cog` and pull an image\n\nOpen Windows Terminal and enter your WSL 2 VM:\n\n```powershell\nwsl.exe\n```\n\nDownload and install `cog` inside the VM:\n\n```bash\nsudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`\nsudo chmod +x /usr/local/bin/cog\n```\n\nMake sure it's available by typing:\n\n```bash\nwhich cog # should output /usr/local/bin/cog\ncog --version # should output the cog version number.\n```\n\n## 8. Run a model in WSL 2\n\nFinally, make sure it works. Let's try running `afiaka87/glid-3-xl` locally:\n\n```bash\ncog predict 'r8.im/afiaka87/glid-3-xl' -i prompt=\"a fresh avocado floating in the water\" -o prediction.json\n```\n\n![Output from a running cog prediction in Windows Terminal](images/cog_model_output.png)\n\nWhile your prediction is running, you can use `Task Manager` to keep an eye on GPU memory consumption:\n\n![Windows task manager will show the shared host/guest GPU memory](images/memory-usage.png)\n\nThis model just barely manages to fit under 8 GB of VRAM.\n\nNotice that output is returned as JSON for this model as it has a complex return type. You will want to convert the base64 string in the json array to an image.\n\n`jq` can help with this:\n\n```sh\nsudo apt install jq\n```\n\nThe following bash uses `jq` to grab the first element in our prediction array and converts it from a base64 string to a `png` file.\n\n```bash\njq -cs '.[0][0][0]' prediction.json | cut --delimiter \",\" --field 2 | base64 --ignore-garbage --decode > prediction.png\n```\n\nWhen using WSL 2, you can access Windows binaries with the `.exe` extension. This lets you open photos easily within linux.\n\n```bash\nexplorer.exe prediction.png\n```\n\n![a square image of an avocado, generated by the model](images/glide_out.png)\n\n## 9. References\n\n- \n- \n- \n- \n- \n" }, { "path": "docs/yaml.md", "content": "# `cog.yaml` reference\n\n`cog.yaml` defines how to build a Docker image and how to run predictions on your model inside that image.\n\nIt has three keys: [`build`](#build), [`image`](#image), and [`predict`](#predict). It looks a bit like this:\n\n```yaml\nbuild:\n python_version: \"3.13\"\n python_requirements: requirements.txt\n system_packages:\n - \"ffmpeg\"\n - \"git\"\npredict: \"predict.py:Predictor\"\n```\n\nTip: Run [`cog init`](getting-started-own-model.md#initialization) to generate an annotated `cog.yaml` file that can be used as a starting point for setting up your model.\n\n## `build`\n\nThis stanza describes how to build the Docker image your model runs in. It contains various options within it:\n\n\n\n### `cuda`\n\nCog automatically picks the correct version of CUDA to install, but this lets you override it for whatever reason by specifying the minor (`11.8`) or patch (`11.8.0`) version of CUDA to use.\n\nFor example:\n\n```yaml\nbuild:\n cuda: \"11.8\"\n```\n\n### `gpu`\n\nEnable GPUs for this model. When enabled, the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image will be used, and Cog will automatically figure out what versions of CUDA and cuDNN to use based on the version of Python, PyTorch, and Tensorflow that you are using.\n\nFor example:\n\n```yaml\nbuild:\n gpu: true\n```\n\nWhen you use `cog run` or `cog predict`, Cog will automatically pass the `--gpus=all` flag to Docker. When you run a Docker image built with Cog, you'll need to pass this option to `docker run`.\n\n### `python_requirements`\n\nA pip requirements file specifying the Python packages to install. For example:\n\n```yaml\nbuild:\n python_requirements: requirements.txt\n```\n\nYour `cog.yaml` file can set either `python_packages` or `python_requirements`, but not both. Use `python_requirements` when you need to configure options like `--extra-index-url` or `--trusted-host` to fetch Python package dependencies.\n\nThis follows the standard [requirements.txt](https://pip.pypa.io/en/stable/reference/requirements-file-format/) format.\n\nTo install Git-hosted Python packages, add `git` to the `system_packages` list, then use the `git+https://` syntax to specify the package name. For example:\n\n`cog.yaml`:\n\n```yaml\nbuild:\n system_packages:\n - \"git\"\n python_requirements: requirements.txt\n```\n\n`requirements.txt`:\n\n```\ngit+https://github.com/huggingface/transformers\n```\n\nYou can also pin Python package installations to a specific git commit:\n\n`cog.yaml`:\n\n```yaml\nbuild:\n system_packages:\n - \"git\"\n python_requirements: requirements.txt\n```\n\n`requirements.txt`:\n\n```\ngit+https://github.com/huggingface/transformers@2d1602a\n```\n\nNote that you can use a shortened prefix of the 40-character git commit SHA, but you must use at least six characters, like `2d1602a` above.\n\n### `python_packages`\n\n**DEPRECATED**: This will be removed in future versions, please use [python_requirements](#python_requirements) instead.\n\nA list of Python packages to install from the PyPi package index, in the format `package==version`. For example:\n\n```yaml\nbuild:\n python_packages:\n - pillow==8.3.1\n - tensorflow==2.5.0\n```\n\nYour `cog.yaml` file can set either `python_packages` or `python_requirements`, but not both.\n\n### `python_version`\n\nThe minor (`3.13`) or patch (`3.13.1`) version of Python to use. For example:\n\n```yaml\nbuild:\n python_version: \"3.13.1\"\n```\n\nCog supports Python 3.10, 3.11, 3.12, and 3.13. If you don't define a version, Cog will use the latest version of Python 3.13 or a version of Python that is compatible with the versions of PyTorch or TensorFlow you specify.\n\nNote that these are the versions supported **in the Docker container**, not your host machine. You can run any version(s) of Python you wish on your host machine.\n\n### `run`\n\nA list of setup commands to run in the environment after your system packages and Python packages have been installed. If you're familiar with Docker, it's like a `RUN` instruction in your `Dockerfile`.\n\nFor example:\n\n```yaml\nbuild:\n run:\n - curl -L https://github.com/cowsay-org/cowsay/archive/refs/tags/v3.7.0.tar.gz | tar -xzf -\n - cd cowsay-3.7.0 && make install\n```\n\nYour code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally.\n\nEach command in `run` can be either a string or a dictionary in the following format:\n\n```yaml\nbuild:\n run:\n - command: pip install\n mounts:\n - type: secret\n id: pip\n target: /etc/pip.conf\n```\n\nYou can use secret mounts to securely pass credentials to setup commands, without baking them into the image. For more information, see [Dockerfile reference](https://docs.docker.com/engine/reference/builder/#run---mounttypesecret).\n\n### `sdk_version`\n\nPin the version of the cog Python SDK installed in the container. Accepts a [PEP 440](https://peps.python.org/pep-0440/) version string. When omitted, the latest release is installed.\n\n```yaml\nbuild:\n python_version: \"3.13\"\n sdk_version: \"0.18.0\"\n```\n\nPre-release versions are also supported:\n\n```yaml\nbuild:\n sdk_version: \"0.18.0a1\"\n```\n\nWhen a pre-release `sdk_version` is set, `--pre` is automatically passed to the pip install commands for both `cog` and `coglet`, so pip will resolve matching pre-release packages.\n\nThe minimum supported version is `0.16.0`. Specifying an older version will cause `cog build` to fail with an error.\n\nThe `COG_SDK_WHEEL` environment variable takes precedence over `sdk_version`. See [Environment variables](./environment.md) for details.\n\n### `system_packages`\n\nA list of Ubuntu APT packages to install. For example:\n\n```yaml\nbuild:\n system_packages:\n - \"ffmpeg\"\n - \"libavcodec-dev\"\n```\n\n## `concurrency`\n\n> Added in cog 0.14.0.\n\nThis stanza describes the concurrency capabilities of the model. It has one option:\n\n### `max`\n\nThe maximum number of concurrent predictions the model can process. If this is set, the model must specify an [async `predict()` method](python.md#async-predictors-and-concurrency).\n\nFor example:\n\n```yaml\nconcurrency:\n max: 10\n```\n\n## `image`\n\nThe name given to built Docker images. If you want to push to a registry, this should also include the registry name.\n\nFor example:\n\n```yaml\nimage: \"r8.im/your-username/your-model\"\n```\n\nr8.im is Replicate's registry, but this can be any Docker registry.\n\nIf you don't set this, then a name will be generated from the directory name.\n\nIf you set this, then you can run `cog push` without specifying the model name.\n\nIf you specify an image name argument when pushing (like `cog push your-username/custom-model-name`), the argument will be used and the value of `image` in cog.yaml will be ignored.\n\n## `predict`\n\nThe pointer to the `Predictor` object in your code, which defines how predictions are run on your model.\n\nFor example:\n\n```yaml\npredict: \"predict.py:Predictor\"\n```\n\nSee [the Python API documentation for more information](python.md).\n" }, { "path": "go.mod", "content": "module github.com/replicate/cog\n\ngo 1.26\n\nrequire (\n\tgithub.com/anaskhan96/soup v1.2.5\n\tgithub.com/containerd/errdefs v1.0.0\n\tgithub.com/creack/pty v1.1.24\n\tgithub.com/docker/cli v29.2.1+incompatible\n\tgithub.com/docker/docker v28.5.2+incompatible\n\tgithub.com/docker/go-connections v0.6.0\n\tgithub.com/getkin/kin-openapi v0.133.0\n\tgithub.com/google/go-containerregistry v0.21.1\n\tgithub.com/hashicorp/go-version v1.7.0\n\tgithub.com/logrusorgru/aurora v2.0.3+incompatible\n\tgithub.com/mattn/go-isatty v0.0.20\n\tgithub.com/mitchellh/go-homedir v1.1.0\n\tgithub.com/moby/buildkit v0.28.0\n\tgithub.com/moby/docker-image-spec v1.3.1\n\tgithub.com/moby/term v0.5.2\n\tgithub.com/opencontainers/image-spec v1.1.1\n\tgithub.com/pkg/errors v0.9.1\n\tgithub.com/replicate/go v0.0.0-20250205165008-b772d7cd506b\n\tgithub.com/rogpeppe/go-internal v1.14.1\n\tgithub.com/sabhiram/go-gitignore v0.0.0-20210923224102-525f6e181f06\n\tgithub.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82\n\tgithub.com/spf13/cobra v1.10.2\n\tgithub.com/spf13/pflag v1.0.10\n\tgithub.com/stretchr/testify v1.11.1\n\tgithub.com/testcontainers/testcontainers-go v0.40.0\n\tgithub.com/testcontainers/testcontainers-go/modules/registry v0.40.0\n\tgithub.com/tonistiigi/go-csvvalue v0.0.0-20240814133006-030d3b2625d0\n\tgithub.com/vincent-petithory/dataurl v1.0.0\n\tgithub.com/xeipuuv/gojsonschema v1.2.0\n\tgithub.com/xeonx/timeago v1.0.0-rc5\n\tgo.yaml.in/yaml/v4 v4.0.0-rc.4\n\tgolang.org/x/crypto v0.48.0\n\tgolang.org/x/exp v0.0.0-20250911091902-df9299821621\n\tgolang.org/x/sync v0.19.0\n\tgolang.org/x/sys v0.42.0\n\tgolang.org/x/term v0.40.0\n\tgoogle.golang.org/grpc v1.79.1\n)\n\nrequire (\n\tdario.cat/mergo v1.0.2 // indirect\n\tgithub.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c // indirect\n\tgithub.com/Microsoft/go-winio v0.6.2 // indirect\n\tgithub.com/bitfield/gotestdox v0.2.2 // indirect\n\tgithub.com/cenkalti/backoff/v4 v4.3.0 // indirect\n\tgithub.com/cespare/xxhash/v2 v2.3.0 // indirect\n\tgithub.com/containerd/console v1.0.5 // indirect\n\tgithub.com/containerd/containerd/api v1.10.0 // indirect\n\tgithub.com/containerd/containerd/v2 v2.2.1 // indirect\n\tgithub.com/containerd/continuity v0.4.5 // indirect\n\tgithub.com/containerd/errdefs/pkg v0.3.0 // indirect\n\tgithub.com/containerd/log v0.1.0 // indirect\n\tgithub.com/containerd/platforms v1.0.0-rc.2 // indirect\n\tgithub.com/containerd/stargz-snapshotter/estargz v0.18.2 // indirect\n\tgithub.com/containerd/ttrpc v1.2.7 // indirect\n\tgithub.com/containerd/typeurl/v2 v2.2.3 // indirect\n\tgithub.com/cpuguy83/dockercfg v0.3.2 // indirect\n\tgithub.com/cpuguy83/go-md2man/v2 v2.0.7 // indirect\n\tgithub.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect\n\tgithub.com/distribution/reference v0.6.0 // indirect\n\tgithub.com/dnephin/pflag v1.0.7 // indirect\n\tgithub.com/docker/distribution v2.8.3+incompatible // indirect\n\tgithub.com/docker/docker-credential-helpers v0.9.5 // indirect\n\tgithub.com/docker/go-units v0.5.0 // indirect\n\tgithub.com/ebitengine/purego v0.8.4 // indirect\n\tgithub.com/fatih/color v1.18.0 // indirect\n\tgithub.com/felixge/httpsnoop v1.0.4 // indirect\n\tgithub.com/fsnotify/fsnotify v1.9.0 // indirect\n\tgithub.com/fvbommel/sortorder v1.1.0 // indirect\n\tgithub.com/go-logr/logr v1.4.3 // indirect\n\tgithub.com/go-logr/stdr v1.2.2 // indirect\n\tgithub.com/go-ole/go-ole v1.2.6 // indirect\n\tgithub.com/go-openapi/jsonpointer v0.22.4 // indirect\n\tgithub.com/go-openapi/swag/jsonname v0.25.4 // indirect\n\tgithub.com/go-test/deep v1.1.1 // indirect\n\tgithub.com/gofrs/flock v0.13.0 // indirect\n\tgithub.com/gogo/protobuf v1.3.2 // indirect\n\tgithub.com/golang/protobuf v1.5.4 // indirect\n\tgithub.com/google/go-cmp v0.7.0 // indirect\n\tgithub.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect\n\tgithub.com/google/uuid v1.6.0 // indirect\n\tgithub.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3 // indirect\n\tgithub.com/in-toto/attestation v1.1.2 // indirect\n\tgithub.com/in-toto/in-toto-golang v0.10.0 // indirect\n\tgithub.com/inconshreveable/mousetrap v1.1.0 // indirect\n\tgithub.com/josharian/intern v1.0.0 // indirect\n\tgithub.com/klauspost/compress v1.18.4 // indirect\n\tgithub.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect\n\tgithub.com/magiconair/properties v1.8.10 // indirect\n\tgithub.com/mailru/easyjson v0.9.0 // indirect\n\tgithub.com/mattn/go-colorable v0.1.14 // indirect\n\tgithub.com/moby/go-archive v0.2.0 // indirect\n\tgithub.com/moby/locker v1.0.1 // indirect\n\tgithub.com/moby/moby/api v1.53.0 // indirect\n\tgithub.com/moby/moby/client v0.2.2 // indirect\n\tgithub.com/moby/patternmatcher v0.6.0 // indirect\n\tgithub.com/moby/sys/atomicwriter v0.1.0 // indirect\n\tgithub.com/moby/sys/sequential v0.6.0 // indirect\n\tgithub.com/moby/sys/signal v0.7.1 // indirect\n\tgithub.com/moby/sys/user v0.4.0 // indirect\n\tgithub.com/moby/sys/userns v0.1.0 // indirect\n\tgithub.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect\n\tgithub.com/morikuni/aec v1.1.0 // indirect\n\tgithub.com/oasdiff/yaml v0.0.0-20250309154309-f31be36b4037 // indirect\n\tgithub.com/oasdiff/yaml3 v0.0.0-20250309153720-d2182401db90 // indirect\n\tgithub.com/opencontainers/go-digest v1.0.0 // indirect\n\tgithub.com/perimeterx/marshmallow v1.1.5 // indirect\n\tgithub.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect\n\tgithub.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect\n\tgithub.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect\n\tgithub.com/russross/blackfriday/v2 v2.1.0 // indirect\n\tgithub.com/secure-systems-lab/go-securesystemslib v0.10.0 // indirect\n\tgithub.com/shibumi/go-pathspec v1.3.0 // indirect\n\tgithub.com/shirou/gopsutil/v4 v4.25.6 // indirect\n\tgithub.com/sirupsen/logrus v1.9.4 // indirect\n\tgithub.com/stretchr/objx v0.5.2 // indirect\n\tgithub.com/tklauser/go-sysconf v0.3.12 // indirect\n\tgithub.com/tklauser/numcpus v0.6.1 // indirect\n\tgithub.com/tonistiigi/fsutil v0.0.0-20251211185533-a2aa163d723f // indirect\n\tgithub.com/tonistiigi/units v0.0.0-20180711220420-6950e57a87ea // indirect\n\tgithub.com/tonistiigi/vt100 v0.0.0-20240514184818-90bafcd6abab // indirect\n\tgithub.com/vbatts/tar-split v0.12.2 // indirect\n\tgithub.com/woodsbury/decimal128 v1.3.0 // indirect\n\tgithub.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect\n\tgithub.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect\n\tgithub.com/yusufpapurcu/wmi v1.2.4 // indirect\n\tgo.opentelemetry.io/auto/sdk v1.2.1 // indirect\n\tgo.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0 // indirect\n\tgo.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace v0.63.0 // indirect\n\tgo.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 // indirect\n\tgo.opentelemetry.io/otel v1.40.0 // indirect\n\tgo.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0 // indirect\n\tgo.opentelemetry.io/otel/metric v1.40.0 // indirect\n\tgo.opentelemetry.io/otel/sdk v1.40.0 // indirect\n\tgo.opentelemetry.io/otel/trace v1.40.0 // indirect\n\tgo.opentelemetry.io/proto/otlp v1.7.1 // indirect\n\tgo.yaml.in/yaml/v3 v3.0.4 // indirect\n\tgolang.org/x/mod v0.33.0 // indirect\n\tgolang.org/x/net v0.51.0 // indirect\n\tgolang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4 // indirect\n\tgolang.org/x/text v0.34.0 // indirect\n\tgolang.org/x/time v0.14.0 // indirect\n\tgolang.org/x/tools v0.42.0 // indirect\n\tgoogle.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect\n\tgoogle.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b // indirect\n\tgoogle.golang.org/protobuf v1.36.11 // indirect\n\tgopkg.in/yaml.v3 v3.0.1 // indirect\n\tgotest.tools/gotestsum v1.12.2 // indirect\n)\n\nreplace (\n\tgithub.com/mholt/archiver/v3 => github.com/bfirsh/archiver/v3 v3.5.1-0.20210316180101-755470a1a69b\n\tgopkg.in/fsnotify.v1 => github.com/kolaente/fsnotify v1.4.10-0.20200411160148-1bc3c8ff4048\n)\n\ntool (\n\tgolang.org/x/tools/cmd/goimports\n\tgotest.tools/gotestsum\n)\n" }, { "path": "go.sum", "content": "cyphar.com/go-pathrs v0.2.1 h1:9nx1vOgwVvX1mNBWDu93+vaceedpbsDqo+XuBGL40b8=\ncyphar.com/go-pathrs v0.2.1/go.mod h1:y8f1EMG7r+hCuFf/rXsKqMJrJAUoADZGNh5/vZPKcGc=\ndario.cat/mergo v1.0.2 h1:85+piFYR1tMbRrLcDwR18y4UKJ3aH1Tbzi24VRW1TK8=\ndario.cat/mergo v1.0.2/go.mod h1:E/hbnu0NxMFBjpMIE34DRGLWqDy0g5FuKDhCb31ngxA=\ngithub.com/AdaLogics/go-fuzz-headers v0.0.0-20240806141605-e8a1dd7889d6 h1:He8afgbRMd7mFxO99hRNu+6tazq8nFF9lIwo9JFroBk=\ngithub.com/AdaLogics/go-fuzz-headers v0.0.0-20240806141605-e8a1dd7889d6/go.mod h1:8o94RPi1/7XTJvwPpRSzSUedZrtlirdB3r9Z20bi2f8=\ngithub.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c h1:udKWzYgxTojEKWjV8V+WSxDXJ4NFATAsZjh8iIbsQIg=\ngithub.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=\ngithub.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=\ngithub.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=\ngithub.com/Microsoft/hcsshim v0.14.0-rc.1 h1:qAPXKwGOkVn8LlqgBN8GS0bxZ83hOJpcjxzmlQKxKsQ=\ngithub.com/Microsoft/hcsshim v0.14.0-rc.1/go.mod h1:hTKFGbnDtQb1wHiOWv4v0eN+7boSWAHyK/tNAaYZL0c=\ngithub.com/ProtonMail/go-crypto v1.3.0 h1:ILq8+Sf5If5DCpHQp4PbZdS1J7HDFRXz/+xKBiRGFrw=\ngithub.com/ProtonMail/go-crypto v1.3.0/go.mod h1:9whxjD8Rbs29b4XWbB8irEcE8KHMqaR2e7GWU1R+/PE=\ngithub.com/anaskhan96/soup v1.2.5 h1:V/FHiusdTrPrdF4iA1YkVxsOpdNcgvqT1hG+YtcZ5hM=\ngithub.com/anaskhan96/soup v1.2.5/go.mod h1:6YnEp9A2yywlYdM4EgDz9NEHclocMepEtku7wg6Cq3s=\ngithub.com/anchore/go-struct-converter v0.1.0 h1:2rDRssAl6mgKBSLNiVCMADgZRhoqtw9dedlWa0OhD30=\ngithub.com/anchore/go-struct-converter v0.1.0/go.mod h1:rYqSE9HbjzpHTI74vwPvae4ZVYZd1lue2ta6xHPdblA=\ngithub.com/bitfield/gotestdox v0.2.2 h1:x6RcPAbBbErKLnapz1QeAlf3ospg8efBsedU93CDsnE=\ngithub.com/bitfield/gotestdox v0.2.2/go.mod h1:D+gwtS0urjBrzguAkTM2wodsTQYFHdpx8eqRJ3N+9pY=\ngithub.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=\ngithub.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=\ngithub.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM=\ngithub.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=\ngithub.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=\ngithub.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=\ngithub.com/cloudflare/circl v1.6.3 h1:9GPOhQGF9MCYUeXyMYlqTR6a5gTrgR/fBLXvUgtVcg8=\ngithub.com/cloudflare/circl v1.6.3/go.mod h1:2eXP6Qfat4O/Yhh8BznvKnJ+uzEoTQ6jVKJRn81BiS4=\ngithub.com/codahale/rfc6979 v0.0.0-20141003034818-6a90f24967eb h1:EDmT6Q9Zs+SbUoc7Ik9EfrFqcylYqgPZ9ANSbTAntnE=\ngithub.com/codahale/rfc6979 v0.0.0-20141003034818-6a90f24967eb/go.mod h1:ZjrT6AXHbDs86ZSdt/osfBi5qfexBrKUdONk989Wnk4=\ngithub.com/containerd/cgroups/v3 v3.1.2 h1:OSosXMtkhI6Qove637tg1XgK4q+DhR0mX8Wi8EhrHa4=\ngithub.com/containerd/cgroups/v3 v3.1.2/go.mod h1:PKZ2AcWmSBsY/tJUVhtS/rluX0b1uq1GmPO1ElCmbOw=\ngithub.com/containerd/console v1.0.5 h1:R0ymNeydRqH2DmakFNdmjR2k0t7UPuiOV/N/27/qqsc=\ngithub.com/containerd/console v1.0.5/go.mod h1:YynlIjWYF8myEu6sdkwKIvGQq+cOckRm6So2avqoYAk=\ngithub.com/containerd/containerd/api v1.10.0 h1:5n0oHYVBwN4VhoX9fFykCV9dF1/BvAXeg2F8W6UYq1o=\ngithub.com/containerd/containerd/api v1.10.0/go.mod h1:NBm1OAk8ZL+LG8R0ceObGxT5hbUYj7CzTmR3xh0DlMM=\ngithub.com/containerd/containerd/v2 v2.2.1 h1:TpyxcY4AL5A+07dxETevunVS5zxqzuq7ZqJXknM11yk=\ngithub.com/containerd/containerd/v2 v2.2.1/go.mod h1:NR70yW1iDxe84F2iFWbR9xfAN0N2F0NcjTi1OVth4nU=\ngithub.com/containerd/continuity v0.4.5 h1:ZRoN1sXq9u7V6QoHMcVWGhOwDFqZ4B9i5H6un1Wh0x4=\ngithub.com/containerd/continuity v0.4.5/go.mod h1:/lNJvtJKUQStBzpVQ1+rasXO1LAWtUQssk28EZvJ3nE=\ngithub.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=\ngithub.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M=\ngithub.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE=\ngithub.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk=\ngithub.com/containerd/fifo v1.1.0 h1:4I2mbh5stb1u6ycIABlBw9zgtlK8viPI9QkQNRQEEmY=\ngithub.com/containerd/fifo v1.1.0/go.mod h1:bmC4NWMbXlt2EZ0Hc7Fx7QzTFxgPID13eH0Qu+MAb2o=\ngithub.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=\ngithub.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo=\ngithub.com/containerd/nydus-snapshotter v0.15.11 h1:YTdF4rsjFRsfyaIhnWVUSLz8FqJwOyRZ5FhvFjHh7Uc=\ngithub.com/containerd/nydus-snapshotter v0.15.11/go.mod h1:EWRd/QJ0b6UKHAqYgiV5gHlqLC2qq5cQiSlXEdVovrA=\ngithub.com/containerd/platforms v1.0.0-rc.2 h1:0SPgaNZPVWGEi4grZdV8VRYQn78y+nm6acgLGv/QzE4=\ngithub.com/containerd/platforms v1.0.0-rc.2/go.mod h1:J71L7B+aiM5SdIEqmd9wp6THLVRzJGXfNuWCZCllLA4=\ngithub.com/containerd/plugin v1.0.0 h1:c8Kf1TNl6+e2TtMHZt+39yAPDbouRH9WAToRjex483Y=\ngithub.com/containerd/plugin v1.0.0/go.mod h1:hQfJe5nmWfImiqT1q8Si3jLv3ynMUIBB47bQ+KexvO8=\ngithub.com/containerd/stargz-snapshotter/estargz v0.18.2 h1:yXkZFYIzz3eoLwlTUZKz2iQ4MrckBxJjkmD16ynUTrw=\ngithub.com/containerd/stargz-snapshotter/estargz v0.18.2/go.mod h1:XyVU5tcJ3PRpkA9XS2T5us6Eg35yM0214Y+wvrZTBrY=\ngithub.com/containerd/ttrpc v1.2.7 h1:qIrroQvuOL9HQ1X6KHe2ohc7p+HP/0VE6XPU7elJRqQ=\ngithub.com/containerd/ttrpc v1.2.7/go.mod h1:YCXHsb32f+Sq5/72xHubdiJRQY9inL4a4ZQrAbN1q9o=\ngithub.com/containerd/typeurl/v2 v2.2.3 h1:yNA/94zxWdvYACdYO8zofhrTVuQY73fFU1y++dYSw40=\ngithub.com/containerd/typeurl/v2 v2.2.3/go.mod h1:95ljDnPfD3bAbDJRugOiShd/DlAAsxGtUBhJxIn7SCk=\ngithub.com/cpuguy83/dockercfg v0.3.2 h1:DlJTyZGBDlXqUZ2Dk2Q3xHs/FtnooJJVaad2S9GKorA=\ngithub.com/cpuguy83/dockercfg v0.3.2/go.mod h1:sugsbF4//dDlL/i+S+rtpIWp+5h0BHJHfjj5/jFyUJc=\ngithub.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=\ngithub.com/cpuguy83/go-md2man/v2 v2.0.7 h1:zbFlGlXEAKlwXpmvle3d8Oe3YnkKIK4xSRTd3sHPnBo=\ngithub.com/cpuguy83/go-md2man/v2 v2.0.7/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=\ngithub.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=\ngithub.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=\ngithub.com/cyphar/filepath-securejoin v0.6.0 h1:BtGB77njd6SVO6VztOHfPxKitJvd/VPT+OFBFMOi1Is=\ngithub.com/cyphar/filepath-securejoin v0.6.0/go.mod h1:A8hd4EnAeyujCJRrICiOWqjS1AX0a9kM5XL+NwKoYSc=\ngithub.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=\ngithub.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=\ngithub.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=\ngithub.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=\ngithub.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=\ngithub.com/dnephin/pflag v1.0.7 h1:oxONGlWxhmUct0YzKTgrpQv9AUA1wtPBn7zuSjJqptk=\ngithub.com/dnephin/pflag v1.0.7/go.mod h1:uxE91IoWURlOiTUIA8Mq5ZZkAv3dPUfZNaT80Zm7OQE=\ngithub.com/docker/cli v29.2.1+incompatible h1:n3Jt0QVCN65eiVBoUTZQM9mcQICCJt3akW4pKAbKdJg=\ngithub.com/docker/cli v29.2.1+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8=\ngithub.com/docker/distribution v2.8.3+incompatible h1:AtKxIZ36LoNK51+Z6RpzLpddBirtxJnzDrHLEKxTAYk=\ngithub.com/docker/distribution v2.8.3+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w=\ngithub.com/docker/docker v28.5.2+incompatible h1:DBX0Y0zAjZbSrm1uzOkdr1onVghKaftjlSWt4AFexzM=\ngithub.com/docker/docker v28.5.2+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=\ngithub.com/docker/docker-credential-helpers v0.9.5 h1:EFNN8DHvaiK8zVqFA2DT6BjXE0GzfLOZ38ggPTKePkY=\ngithub.com/docker/docker-credential-helpers v0.9.5/go.mod h1:v1S+hepowrQXITkEfw6o4+BMbGot02wiKpzWhGUZK6c=\ngithub.com/docker/go-connections v0.6.0 h1:LlMG9azAe1TqfR7sO+NJttz1gy6KO7VJBh+pMmjSD94=\ngithub.com/docker/go-connections v0.6.0/go.mod h1:AahvXYshr6JgfUJGdDCs2b5EZG/vmaMAntpSFH5BFKE=\ngithub.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=\ngithub.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=\ngithub.com/ebitengine/purego v0.8.4 h1:CF7LEKg5FFOsASUj0+QwaXf8Ht6TlFxg09+S9wz0omw=\ngithub.com/ebitengine/purego v0.8.4/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=\ngithub.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=\ngithub.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=\ngithub.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=\ngithub.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=\ngithub.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k=\ngithub.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=\ngithub.com/fvbommel/sortorder v1.1.0 h1:fUmoe+HLsBTctBDoaBwpQo5N+nrCp8g/BjKb/6ZQmYw=\ngithub.com/fvbommel/sortorder v1.1.0/go.mod h1:uk88iVf1ovNn1iLfgUVU2F9o5eO30ui720w+kxuqRs0=\ngithub.com/getkin/kin-openapi v0.133.0 h1:pJdmNohVIJ97r4AUFtEXRXwESr8b0bD721u/Tz6k8PQ=\ngithub.com/getkin/kin-openapi v0.133.0/go.mod h1:boAciF6cXk5FhPqe/NQeBTeenbjqU4LhWBf09ILVvWE=\ngithub.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=\ngithub.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=\ngithub.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=\ngithub.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=\ngithub.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=\ngithub.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY=\ngithub.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=\ngithub.com/go-openapi/jsonpointer v0.22.4 h1:dZtK82WlNpVLDW2jlA1YCiVJFVqkED1MegOUy9kR5T4=\ngithub.com/go-openapi/jsonpointer v0.22.4/go.mod h1:elX9+UgznpFhgBuaMQ7iu4lvvX1nvNsesQ3oxmYTw80=\ngithub.com/go-openapi/swag/jsonname v0.25.4 h1:bZH0+MsS03MbnwBXYhuTttMOqk+5KcQ9869Vye1bNHI=\ngithub.com/go-openapi/swag/jsonname v0.25.4/go.mod h1:GPVEk9CWVhNvWhZgrnvRA6utbAltopbKwDu8mXNUMag=\ngithub.com/go-openapi/testify/v2 v2.0.2 h1:X999g3jeLcoY8qctY/c/Z8iBHTbwLz7R2WXd6Ub6wls=\ngithub.com/go-openapi/testify/v2 v2.0.2/go.mod h1:HCPmvFFnheKK2BuwSA0TbbdxJ3I16pjwMkYkP4Ywn54=\ngithub.com/go-test/deep v1.1.1 h1:0r/53hagsehfO4bzD2Pgr/+RgHqhmf+k1Bpse2cTu1U=\ngithub.com/go-test/deep v1.1.1/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE=\ngithub.com/gofrs/flock v0.13.0 h1:95JolYOvGMqeH31+FC7D2+uULf6mG61mEZ/A8dRYMzw=\ngithub.com/gofrs/flock v0.13.0/go.mod h1:jxeyy9R1auM5S6JYDBhDt+E2TCo7DkratH4Pgi8P+Z0=\ngithub.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=\ngithub.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=\ngithub.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 h1:f+oWsMOmNPc8JmEHVZIycC7hBoQxHH9pNKQORJNozsQ=\ngithub.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8/go.mod h1:wcDNUvekVysuuOpQKo3191zZyTpiI6se1N1ULghS0sw=\ngithub.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=\ngithub.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=\ngithub.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=\ngithub.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=\ngithub.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=\ngithub.com/google/go-containerregistry v0.21.1 h1:sOt/o9BS2b87FnR7wxXPvRKU1XVJn2QCwOS5g8zQXlc=\ngithub.com/google/go-containerregistry v0.21.1/go.mod h1:ctO5aCaewH4AK1AumSF5DPW+0+R+d2FmylMJdp5G7p0=\ngithub.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=\ngithub.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=\ngithub.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=\ngithub.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=\ngithub.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3 h1:NmZ1PKzSTQbuGHw9DGPFomqkkLWMC+vZCkfs+FHv1Vg=\ngithub.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3/go.mod h1:zQrxl1YP88HQlA6i9c63DSVPFklWpGX4OWAc9bFuaH4=\ngithub.com/hashicorp/go-version v1.7.0 h1:5tqGy27NaOTB8yJKUZELlFAS/LTKJkrmONwQKeRZfjY=\ngithub.com/hashicorp/go-version v1.7.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=\ngithub.com/in-toto/attestation v1.1.2 h1:MBFn6lsMq6dptQZJBhalXTcWMb/aJy3V+GX3VYj/V1E=\ngithub.com/in-toto/attestation v1.1.2/go.mod h1:gYFddHMZj3DiQ0b62ltNi1Vj5rC879bTmBbrv9CRHpM=\ngithub.com/in-toto/in-toto-golang v0.10.0 h1:+s2eZQSK3WmWfYV85qXVSBfqgawi/5L02MaqA4o/tpM=\ngithub.com/in-toto/in-toto-golang v0.10.0/go.mod h1:wjT4RiyFlLWCmLUJjwB8oZcjaq7HA390aMJcD3xXgmg=\ngithub.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=\ngithub.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=\ngithub.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY=\ngithub.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=\ngithub.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=\ngithub.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=\ngithub.com/klauspost/compress v1.18.4 h1:RPhnKRAQ4Fh8zU2FY/6ZFDwTVTxgJ/EMydqSTzE9a2c=\ngithub.com/klauspost/compress v1.18.4/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4=\ngithub.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=\ngithub.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=\ngithub.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=\ngithub.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=\ngithub.com/logrusorgru/aurora v2.0.3+incompatible h1:tOpm7WcpBTn4fjmVfgpQq0EfczGlG91VSDkswnjF5A8=\ngithub.com/logrusorgru/aurora v2.0.3+incompatible/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=\ngithub.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4=\ngithub.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=\ngithub.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8SYxI99mE=\ngithub.com/magiconair/properties v1.8.10/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=\ngithub.com/mailru/easyjson v0.9.0 h1:PrnmzHw7262yW8sTBwxi1PdJA3Iw/EKBa8psRf7d9a4=\ngithub.com/mailru/easyjson v0.9.0/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU=\ngithub.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=\ngithub.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=\ngithub.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=\ngithub.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=\ngithub.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=\ngithub.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=\ngithub.com/moby/buildkit v0.28.0 h1:rKulfRRSduHJPNpLTk481fHElqN9tps0VUx8YV/5zsA=\ngithub.com/moby/buildkit v0.28.0/go.mod h1:RCuOcj/bVsCriBG8NeFzRxjiCFQKnKP7KOVlNTS18t4=\ngithub.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=\ngithub.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo=\ngithub.com/moby/go-archive v0.2.0 h1:zg5QDUM2mi0JIM9fdQZWC7U8+2ZfixfTYoHL7rWUcP8=\ngithub.com/moby/go-archive v0.2.0/go.mod h1:mNeivT14o8xU+5q1YnNrkQVpK+dnNe/K6fHqnTg4qPU=\ngithub.com/moby/locker v1.0.1 h1:fOXqR41zeveg4fFODix+1Ch4mj/gT0NE1XJbp/epuBg=\ngithub.com/moby/locker v1.0.1/go.mod h1:S7SDdo5zpBK84bzzVlKr2V0hz+7x9hWbYC/kq7oQppc=\ngithub.com/moby/moby/api v1.53.0 h1:PihqG1ncw4W+8mZs69jlwGXdaYBeb5brF6BL7mPIS/w=\ngithub.com/moby/moby/api v1.53.0/go.mod h1:8mb+ReTlisw4pS6BRzCMts5M49W5M7bKt1cJy/YbAqc=\ngithub.com/moby/moby/client v0.2.2 h1:Pt4hRMCAIlyjL3cr8M5TrXCwKzguebPAc2do2ur7dEM=\ngithub.com/moby/moby/client v0.2.2/go.mod h1:2EkIPVNCqR05CMIzL1mfA07t0HvVUUOl85pasRz/GmQ=\ngithub.com/moby/patternmatcher v0.6.0 h1:GmP9lR19aU5GqSSFko+5pRqHi+Ohk1O69aFiKkVGiPk=\ngithub.com/moby/patternmatcher v0.6.0/go.mod h1:hDPoyOpDY7OrrMDLaYoY3hf52gNCR/YOUYxkhApJIxc=\ngithub.com/moby/policy-helpers v0.0.0-20260211190020-824747bfdd3c h1:hRUo0Ir9PEaa0PQCgg8WvGku0sgmTo/NgnCzMb83iII=\ngithub.com/moby/policy-helpers v0.0.0-20260211190020-824747bfdd3c/go.mod h1:2P1OGoTVIrybI4M7yhpkDpqiwOnI3yR+HnNhEyo8ovs=\ngithub.com/moby/sys/atomicwriter v0.1.0 h1:kw5D/EqkBwsBFi0ss9v1VG3wIkVhzGvLklJ+w3A14Sw=\ngithub.com/moby/sys/atomicwriter v0.1.0/go.mod h1:Ul8oqv2ZMNHOceF643P6FKPXeCmYtlQMvpizfsSoaWs=\ngithub.com/moby/sys/mountinfo v0.7.2 h1:1shs6aH5s4o5H2zQLn796ADW1wMrIwHsyJ2v9KouLrg=\ngithub.com/moby/sys/mountinfo v0.7.2/go.mod h1:1YOa8w8Ih7uW0wALDUgT1dTTSBrZ+HiBLGws92L2RU4=\ngithub.com/moby/sys/sequential v0.6.0 h1:qrx7XFUd/5DxtqcoH1h438hF5TmOvzC/lspjy7zgvCU=\ngithub.com/moby/sys/sequential v0.6.0/go.mod h1:uyv8EUTrca5PnDsdMGXhZe6CCe8U/UiTWd+lL+7b/Ko=\ngithub.com/moby/sys/signal v0.7.1 h1:PrQxdvxcGijdo6UXXo/lU/TvHUWyPhj7UOpSo8tuvk0=\ngithub.com/moby/sys/signal v0.7.1/go.mod h1:Se1VGehYokAkrSQwL4tDzHvETwUZlnY7S5XtQ50mQp8=\ngithub.com/moby/sys/user v0.4.0 h1:jhcMKit7SA80hivmFJcbB1vqmw//wU61Zdui2eQXuMs=\ngithub.com/moby/sys/user v0.4.0/go.mod h1:bG+tYYYJgaMtRKgEmuueC0hJEAZWwtIbZTB+85uoHjs=\ngithub.com/moby/sys/userns v0.1.0 h1:tVLXkFOxVu9A64/yh59slHVv9ahO9UIev4JZusOLG/g=\ngithub.com/moby/sys/userns v0.1.0/go.mod h1:IHUYgu/kao6N8YZlp9Cf444ySSvCmDlmzUcYfDHOl28=\ngithub.com/moby/term v0.5.2 h1:6qk3FJAFDs6i/q3W/pQ97SX192qKfZgGjCQqfCJkgzQ=\ngithub.com/moby/term v0.5.2/go.mod h1:d3djjFCrjnB+fl8NJux+EJzu0msscUP+f8it8hPkFLc=\ngithub.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 h1:RWengNIwukTxcDr9M+97sNutRR1RKhG96O6jWumTTnw=\ngithub.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826/go.mod h1:TaXosZuwdSHYgviHp1DAtfrULt5eUgsSMsZf+YrPgl8=\ngithub.com/morikuni/aec v1.1.0 h1:vBBl0pUnvi/Je71dsRrhMBtreIqNMYErSAbEeb8jrXQ=\ngithub.com/morikuni/aec v1.1.0/go.mod h1:xDRgiq/iw5l+zkao76YTKzKttOp2cwPEne25HDkJnBw=\ngithub.com/oasdiff/yaml v0.0.0-20250309154309-f31be36b4037 h1:G7ERwszslrBzRxj//JalHPu/3yz+De2J+4aLtSRlHiY=\ngithub.com/oasdiff/yaml v0.0.0-20250309154309-f31be36b4037/go.mod h1:2bpvgLBZEtENV5scfDFEtB/5+1M4hkQhDQrccEJ/qGw=\ngithub.com/oasdiff/yaml3 v0.0.0-20250309153720-d2182401db90 h1:bQx3WeLcUWy+RletIKwUIt4x3t8n2SxavmoclizMb8c=\ngithub.com/oasdiff/yaml3 v0.0.0-20250309153720-d2182401db90/go.mod h1:y5+oSEHCPT/DGrS++Wc/479ERge0zTFxaF8PbGKcg2o=\ngithub.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=\ngithub.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=\ngithub.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040=\ngithub.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M=\ngithub.com/opencontainers/runtime-spec v1.3.0 h1:YZupQUdctfhpZy3TM39nN9Ika5CBWT5diQ8ibYCRkxg=\ngithub.com/opencontainers/runtime-spec v1.3.0/go.mod h1:jwyrGlmzljRJv/Fgzds9SsS/C5hL+LL3ko9hs6T5lQ0=\ngithub.com/opencontainers/selinux v1.13.1 h1:A8nNeceYngH9Ow++M+VVEwJVpdFmrlxsN22F+ISDCJE=\ngithub.com/opencontainers/selinux v1.13.1/go.mod h1:S10WXZ/osk2kWOYKy1x2f/eXF5ZHJoUs8UU/2caNRbg=\ngithub.com/package-url/packageurl-go v0.1.1 h1:KTRE0bK3sKbFKAk3yy63DpeskU7Cvs/x/Da5l+RtzyU=\ngithub.com/package-url/packageurl-go v0.1.1/go.mod h1:uQd4a7Rh3ZsVg5j0lNyAfyxIeGde9yrlhjF78GzeW0c=\ngithub.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4=\ngithub.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=\ngithub.com/perimeterx/marshmallow v1.1.5 h1:a2LALqQ1BlHM8PZblsDdidgv1mWi1DgC2UmX50IvK2s=\ngithub.com/perimeterx/marshmallow v1.1.5/go.mod h1:dsXbUu8CRzfYP5a87xpp0xq9S3u0Vchtcl8we9tYaXw=\ngithub.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=\ngithub.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=\ngithub.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo=\ngithub.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=\ngithub.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=\ngithub.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=\ngithub.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=\ngithub.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=\ngithub.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=\ngithub.com/prometheus/procfs v0.17.0 h1:FuLQ+05u4ZI+SS/w9+BWEM2TXiHKsUQ9TADiRH7DuK0=\ngithub.com/prometheus/procfs v0.17.0/go.mod h1:oPQLaDAMRbA+u8H5Pbfq+dl3VDAvHxMUOVhe0wYB2zw=\ngithub.com/replicate/go v0.0.0-20250205165008-b772d7cd506b h1:GIkpkQ+xwWJ6IRUFmwCLcg+zkZVoKmVXnPjhMncZc4I=\ngithub.com/replicate/go v0.0.0-20250205165008-b772d7cd506b/go.mod h1:kUMwEnHJEvWXdu6Py/9fjp7969tsPRYN2a4+Z8BiVEE=\ngithub.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=\ngithub.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=\ngithub.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk=\ngithub.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=\ngithub.com/sabhiram/go-gitignore v0.0.0-20210923224102-525f6e181f06 h1:OkMGxebDjyw0ULyrTYWeN0UNCCkmCWfjPnIA2W6oviI=\ngithub.com/sabhiram/go-gitignore v0.0.0-20210923224102-525f6e181f06/go.mod h1:+ePHsJ1keEjQtpvf9HHw0f4ZeJ0TLRsxhunSI2hYJSs=\ngithub.com/secure-systems-lab/go-securesystemslib v0.10.0 h1:l+H5ErcW0PAehBNrBxoGv1jjNpGYdZ9RcheFkB2WI14=\ngithub.com/secure-systems-lab/go-securesystemslib v0.10.0/go.mod h1:MRKONWmRoFzPNQ9USRF9i1mc7MvAVvF1LlW8X5VWDvk=\ngithub.com/shibumi/go-pathspec v1.3.0 h1:QUyMZhFo0Md5B8zV8x2tesohbb5kfbpTi9rBnKh5dkI=\ngithub.com/shibumi/go-pathspec v1.3.0/go.mod h1:Xutfslp817l2I1cZvgcfeMQJG5QnU2lh5tVaaMCl3jE=\ngithub.com/shirou/gopsutil/v4 v4.25.6 h1:kLysI2JsKorfaFPcYmcJqbzROzsBWEOAtw6A7dIfqXs=\ngithub.com/shirou/gopsutil/v4 v4.25.6/go.mod h1:PfybzyydfZcN+JMMjkF6Zb8Mq1A/VcogFFg7hj50W9c=\ngithub.com/sigstore/sigstore v1.10.4 h1:ytOmxMgLdcUed3w1SbbZOgcxqwMG61lh1TmZLN+WeZE=\ngithub.com/sigstore/sigstore v1.10.4/go.mod h1:tDiyrdOref3q6qJxm2G+JHghqfmvifB7hw+EReAfnbI=\ngithub.com/sigstore/sigstore-go v1.1.4 h1:wTTsgCHOfqiEzVyBYA6mDczGtBkN7cM8mPpjJj5QvMg=\ngithub.com/sigstore/sigstore-go v1.1.4/go.mod h1:2U/mQOT9cjjxrtIUeKDVhL+sHBKsnWddn8URlswdBsg=\ngithub.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w=\ngithub.com/sirupsen/logrus v1.9.4/go.mod h1:ftWc9WdOfJ0a92nsE2jF5u5ZwH8Bv2zdeOC42RjbV2g=\ngithub.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82 h1:6C8qej6f1bStuePVkLSFxoU22XBS165D3klxlzRg8F4=\ngithub.com/smacker/go-tree-sitter v0.0.0-20240827094217-dd81d9e9be82/go.mod h1:xe4pgH49k4SsmkQq5OT8abwhWmnzkhpgnXeekbx2efw=\ngithub.com/spdx/tools-golang v0.5.7 h1:+sWcKGnhwp3vLdMqPcLdA6QK679vd86cK9hQWH3AwCg=\ngithub.com/spdx/tools-golang v0.5.7/go.mod h1:jg7w0LOpoNAw6OxKEzCoqPC2GCTj45LyTlVmXubDsYw=\ngithub.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=\ngithub.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=\ngithub.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=\ngithub.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk=\ngithub.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=\ngithub.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=\ngithub.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=\ngithub.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=\ngithub.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=\ngithub.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=\ngithub.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=\ngithub.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=\ngithub.com/testcontainers/testcontainers-go v0.40.0 h1:pSdJYLOVgLE8YdUY2FHQ1Fxu+aMnb6JfVz1mxk7OeMU=\ngithub.com/testcontainers/testcontainers-go v0.40.0/go.mod h1:FSXV5KQtX2HAMlm7U3APNyLkkap35zNLxukw9oBi/MY=\ngithub.com/testcontainers/testcontainers-go/modules/registry v0.40.0 h1:z+CymIuT9quh8plBbM+lpncY6diV//q0LbRk+mxMpow=\ngithub.com/testcontainers/testcontainers-go/modules/registry v0.40.0/go.mod h1:TWdy7+y7w14Ii5UCSfr7qvxPYI3GE7lc7NEP0ofxlLQ=\ngithub.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU=\ngithub.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=\ngithub.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk=\ngithub.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY=\ngithub.com/tonistiigi/fsutil v0.0.0-20251211185533-a2aa163d723f h1:Z4NEQ86qFl1mHuCu9gwcE+EYCwDKfXAYXZbdIXyxmEA=\ngithub.com/tonistiigi/fsutil v0.0.0-20251211185533-a2aa163d723f/go.mod h1:BKdcez7BiVtBvIcef90ZPc6ebqIWr4JWD7+EvLm6J98=\ngithub.com/tonistiigi/go-csvvalue v0.0.0-20240814133006-030d3b2625d0 h1:2f304B10LaZdB8kkVEaoXvAMVan2tl9AiK4G0odjQtE=\ngithub.com/tonistiigi/go-csvvalue v0.0.0-20240814133006-030d3b2625d0/go.mod h1:278M4p8WsNh3n4a1eqiFcV2FGk7wE5fwUpUom9mK9lE=\ngithub.com/tonistiigi/units v0.0.0-20180711220420-6950e57a87ea h1:SXhTLE6pb6eld/v/cCndK0AMpt1wiVFb/YYmqB3/QG0=\ngithub.com/tonistiigi/units v0.0.0-20180711220420-6950e57a87ea/go.mod h1:WPnis/6cRcDZSUvVmezrxJPkiO87ThFYsoUiMwWNDJk=\ngithub.com/tonistiigi/vt100 v0.0.0-20240514184818-90bafcd6abab h1:H6aJ0yKQ0gF49Qb2z5hI1UHxSQt4JMyxebFR15KnApw=\ngithub.com/tonistiigi/vt100 v0.0.0-20240514184818-90bafcd6abab/go.mod h1:ulncasL3N9uLrVann0m+CDlJKWsIAP34MPcOJF6VRvc=\ngithub.com/ugorji/go/codec v1.2.7 h1:YPXUKf7fYbp/y8xloBqZOw2qaVggbfwMlI8WM3wZUJ0=\ngithub.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY=\ngithub.com/vbatts/tar-split v0.12.2 h1:w/Y6tjxpeiFMR47yzZPlPj/FcPLpXbTUi/9H7d3CPa4=\ngithub.com/vbatts/tar-split v0.12.2/go.mod h1:eF6B6i6ftWQcDqEn3/iGFRFRo8cBIMSJVOpnNdfTMFA=\ngithub.com/vincent-petithory/dataurl v1.0.0 h1:cXw+kPto8NLuJtlMsI152irrVw9fRDX8AbShPRpg2CI=\ngithub.com/vincent-petithory/dataurl v1.0.0/go.mod h1:FHafX5vmDzyP+1CQATJn7WFKc9CvnvxyvZy6I1MrG/U=\ngithub.com/woodsbury/decimal128 v1.3.0 h1:8pffMNWIlC0O5vbyHWFZAt5yWvWcrHA+3ovIIjVWss0=\ngithub.com/woodsbury/decimal128 v1.3.0/go.mod h1:C5UTmyTjW3JftjUFzOVhC20BEQa2a4ZKOB5I6Zjb+ds=\ngithub.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU=\ngithub.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb h1:zGWFAtiMcyryUHoUjUJX0/lt1H2+i2Ka2n+D3DImSNo=\ngithub.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU=\ngithub.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 h1:EzJWgHovont7NscjpAxXsDA8S8BMYve8Y5+7cuRE7R0=\ngithub.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415/go.mod h1:GwrjFmJcFw6At/Gs6z4yjiIwzuJ1/+UwLxMQDVQXShQ=\ngithub.com/xeipuuv/gojsonschema v1.2.0 h1:LhYJRs+L4fBtjZUfuSZIKGeVu0QRy8e5Xi7D17UxZ74=\ngithub.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y=\ngithub.com/xeonx/timeago v1.0.0-rc5 h1:pwcQGpaH3eLfPtXeyPA4DmHWjoQt0Ea7/++FwpxqLxg=\ngithub.com/xeonx/timeago v1.0.0-rc5/go.mod h1:qDLrYEFynLO7y5Ho7w3GwgtYgpy5UfhcXIIQvMKVDkA=\ngithub.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=\ngithub.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=\ngithub.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=\ngithub.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=\ngo.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0=\ngo.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo=\ngo.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=\ngo.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=\ngo.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0 h1:YH4g8lQroajqUwWbq/tr2QX1JFmEXaDLgG+ew9bLMWo=\ngo.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0/go.mod h1:fvPi2qXDqFs8M4B4fmJhE92TyQs9Ydjlg3RvfUp+NbQ=\ngo.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace v0.63.0 h1:2pn7OzMewmYRiNtv1doZnLo3gONcnMHlFnmOR8Vgt+8=\ngo.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace v0.63.0/go.mod h1:rjbQTDEPQymPE0YnRQp9/NuPwwtL0sesz/fnqRW/v84=\ngo.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 h1:RbKq8BG0FI8OiXhBfcRtqqHcZcka+gU3cskNuf05R18=\ngo.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0/go.mod h1:h06DGIukJOevXaj/xrNjhi/2098RZzcLTbc0jDAUbsg=\ngo.opentelemetry.io/otel v1.40.0 h1:oA5YeOcpRTXq6NN7frwmwFR0Cn3RhTVZvXsP4duvCms=\ngo.opentelemetry.io/otel v1.40.0/go.mod h1:IMb+uXZUKkMXdPddhwAHm6UfOwJyh4ct1ybIlV14J0g=\ngo.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0 h1:GqRJVj7UmLjCVyVJ3ZFLdPRmhDUp2zFmQe3RHIOsw24=\ngo.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0/go.mod h1:ri3aaHSmCTVYu2AWv44YMauwAQc0aqI9gHKIcSbI1pU=\ngo.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0 h1:aTL7F04bJHUlztTsNGJ2l+6he8c+y/b//eR0jjjemT4=\ngo.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0/go.mod h1:kldtb7jDTeol0l3ewcmd8SDvx3EmIE7lyvqbasU3QC4=\ngo.opentelemetry.io/otel/metric v1.40.0 h1:rcZe317KPftE2rstWIBitCdVp89A2HqjkxR3c11+p9g=\ngo.opentelemetry.io/otel/metric v1.40.0/go.mod h1:ib/crwQH7N3r5kfiBZQbwrTge743UDc7DTFVZrrXnqc=\ngo.opentelemetry.io/otel/sdk v1.40.0 h1:KHW/jUzgo6wsPh9At46+h4upjtccTmuZCFAc9OJ71f8=\ngo.opentelemetry.io/otel/sdk v1.40.0/go.mod h1:Ph7EFdYvxq72Y8Li9q8KebuYUr2KoeyHx0DRMKrYBUE=\ngo.opentelemetry.io/otel/sdk/metric v1.40.0 h1:mtmdVqgQkeRxHgRv4qhyJduP3fYJRMX4AtAlbuWdCYw=\ngo.opentelemetry.io/otel/sdk/metric v1.40.0/go.mod h1:4Z2bGMf0KSK3uRjlczMOeMhKU2rhUqdWNoKcYrtcBPg=\ngo.opentelemetry.io/otel/trace v1.40.0 h1:WA4etStDttCSYuhwvEa8OP8I5EWu24lkOzp+ZYblVjw=\ngo.opentelemetry.io/otel/trace v1.40.0/go.mod h1:zeAhriXecNGP/s2SEG3+Y8X9ujcJOTqQ5RgdEJcawiA=\ngo.opentelemetry.io/proto/otlp v1.7.1 h1:gTOMpGDb0WTBOP8JaO72iL3auEZhVmAQg4ipjOVAtj4=\ngo.opentelemetry.io/proto/otlp v1.7.1/go.mod h1:b2rVh6rfI/s2pHWNlB7ILJcRALpcNDzKhACevjI+ZnE=\ngo.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=\ngo.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=\ngo.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc=\ngo.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=\ngo.yaml.in/yaml/v4 v4.0.0-rc.4 h1:UP4+v6fFrBIb1l934bDl//mmnoIZEDK0idg1+AIvX5U=\ngo.yaml.in/yaml/v4 v4.0.0-rc.4/go.mod h1:aZqd9kCMsGL7AuUv/m/PvWLdg5sjJsZ4oHDEnfPPfY0=\ngolang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=\ngolang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=\ngolang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=\ngolang.org/x/crypto v0.48.0 h1:/VRzVqiRSggnhY7gNRxPauEQ5Drw9haKdM0jqfcCFts=\ngolang.org/x/crypto v0.48.0/go.mod h1:r0kV5h3qnFPlQnBSrULhlsRfryS2pmewsg+XfMgkVos=\ngolang.org/x/exp v0.0.0-20250911091902-df9299821621 h1:2id6c1/gto0kaHYyrixvknJ8tUK/Qs5IsmBtrc+FtgU=\ngolang.org/x/exp v0.0.0-20250911091902-df9299821621/go.mod h1:TwQYMMnGpvZyc+JpB/UAuTNIsVJifOlSkrZkhcvpVUk=\ngolang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=\ngolang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=\ngolang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8=\ngolang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w=\ngolang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=\ngolang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=\ngolang.org/x/net v0.0.0-20200114155413-6afb5195e5aa/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=\ngolang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=\ngolang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=\ngolang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo=\ngolang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y=\ngolang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=\ngolang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=\ngolang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=\ngolang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=\ngolang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=\ngolang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=\ngolang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=\ngolang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=\ngolang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=\ngolang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=\ngolang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=\ngolang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=\ngolang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=\ngolang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=\ngolang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=\ngolang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo=\ngolang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=\ngolang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4 h1:bTLqdHv7xrGlFbvf5/TXNxy/iUwwdkjhqQTJDjW7aj0=\ngolang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4/go.mod h1:g5NllXBEermZrmR51cJDQxmJUHUOfRAaNyWBM+R+548=\ngolang.org/x/term v0.40.0 h1:36e4zGLqU4yhjlmxEaagx2KuYbJq3EwY8K943ZsHcvg=\ngolang.org/x/term v0.40.0/go.mod h1:w2P8uVp06p2iyKKuvXIm7N/y0UCRt3UfJTfZ7oOpglM=\ngolang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=\ngolang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=\ngolang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=\ngolang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA=\ngolang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=\ngolang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=\ngolang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=\ngolang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=\ngolang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=\ngolang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=\ngolang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=\ngolang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=\ngolang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=\ngolang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=\ngolang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=\ngolang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=\ngonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=\ngonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=\ngoogle.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 h1:fCvbg86sFXwdrl5LgVcTEvNC+2txB5mgROGmRL5mrls=\ngoogle.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:+rXWjjaukWZun3mLfjmVnQi18E1AsFbDN9QdJ5YXLto=\ngoogle.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b h1:Mv8VFug0MP9e5vUxfBcE3vUkV6CImK3cMNMIDFjmzxU=\ngoogle.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b/go.mod h1:j9x/tPzZkyxcgEFkiKEEGxfvyumM01BEtsW8xzOahRQ=\ngoogle.golang.org/grpc v1.79.1 h1:zGhSi45ODB9/p3VAawt9a+O/MULLl9dpizzNNpq7flY=\ngoogle.golang.org/grpc v1.79.1/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=\ngoogle.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=\ngoogle.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=\ngopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=\ngopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=\ngopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=\ngopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=\ngopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=\ngopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=\ngotest.tools/gotestsum v1.12.2 h1:eli4tu9Q2D/ogDsEGSr8XfQfl7mT0JsGOG6DFtUiZ/Q=\ngotest.tools/gotestsum v1.12.2/go.mod h1:kjRtCglPZVsSU0hFHX3M5VWBM6Y63emHuB14ER1/sow=\ngotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q=\ngotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA=\npgregory.net/rapid v1.2.0 h1:keKAYRcjm+e1F0oAuU5F5+YPAWcyxNNRK2wud503Gnk=\npgregory.net/rapid v1.2.0/go.mod h1:PY5XlDGj0+V1FCq0o192FdRhpKHGTRIWBgqjDBTrq04=\n" }, { "path": "integration-tests/.gitignore", "content": ".bin/\n" }, { "path": "integration-tests/README.md", "content": "# Cog Integration Tests\n\nThis directory contains Go-based integration tests for the Cog CLI using the [testscript](https://pkg.go.dev/github.com/rogpeppe/go-internal/testscript) framework.\n\n## Test Formats\n\nMost integration tests use the txtar format (`.txtar` files in `tests/`), which provides a simple declarative way to define test scripts and fixtures.\n\nHowever, some tests require capabilities that don't fit txtar's sequential execution model and are written as standard Go test functions instead:\n\n| Test | Location | Why Go instead of txtar |\n|------|----------|-------------------------|\n| `TestConcurrentPredictions` | `concurrent/` | Requires parallel HTTP requests with precise timing coordination |\n| `TestLogin*` | `login/` | Login requires interactive PTY input and mock HTTP servers |\n\nNote: PTY/TTY tests now use the `pty-run` command in txtar format (see Custom Commands below).\n\n## Quick Start\n\n```bash\n# Run all tests\nmake test-integration\n\n# Run fast tests only (skip slow GPU/framework tests)\ncd integration-tests && go test -short -v\n\n# Run a specific test\ncd integration-tests && go test -v -run TestIntegration/string_predictor\n\n# Run with a custom cog binary\nCOG_BINARY=/path/to/cog make test-integration\n```\n\n## Directory Structure\n\n```\nintegration-tests/\n├── README.md # This file\n├── suite_test.go # Main test runner (txtar tests)\n├── harness/\n│ ├── harness.go # Test harness core\n│ ├── command.go # Command interface\n│ └── cmd_pty.go # PTY command implementation\n├── tests/\n│ └── *.txtar # Test files (one per test case)\n├── concurrent/\n│ └── concurrent_test.go # Concurrent request tests\n├── login/\n│ └── login_test.go # Login tests with PTY\n└── .bin/\n └── cog # Cached cog binary (auto-generated)\n```\n\n## Writing Tests\n\nTests are `.txtar` files in the `tests/` directory. Each file is a self-contained test with embedded fixtures.\n\n**NOTE: if a test has the suffix `_serial` (e.g. `tests/integration_test_name_serial.txtar`) it will be run in isolation of all other tests. By default we run integration tests in Parallel using `t.Parallel()`.**\n\n### Editor Support\n\nFor syntax highlighting of `.txtar` files:\n\n**VS Code:**\n- [testscript](https://marketplace.visualstudio.com/items?itemName=twpayne.vscode-testscript) by twpayne - Syntax highlighting with embedded file support\n- [txtar](https://github.com/brody715/vscode-txtar) by brody715 - Alternative txtar extension\n\nInstall via VS Code:\n```\next install twpayne.vscode-testscript\n```\n\n**Zed:**\n- [zed-txtar](https://github.com/FollowTheProcess/zed-txtar) - Syntax highlighting for txtar files\n\nInstall via Zed extensions panel or add to your extensions.\n\n**Vim/Neovim:**\n- Use [tree-sitter-go-template](https://github.com/ngalaiko/tree-sitter-go-template) for basic support\n- Or set filetype manually: `:set ft=conf` for basic highlighting\n\n### Basic Test Structure\n\n```txtar\n# Comments describe what the test does\n# This is a test for basic string prediction\n\n# Build the Docker image\ncog build -t $TEST_IMAGE\n\n# Run a prediction\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n# Test that invalid input fails\n! cog predict $TEST_IMAGE -i wrong=value\nstderr 'Field required'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def predict(self, s: str) -> str:\n return \"hello \" + s\n```\n\n### Test File Format\n\n- Lines starting with `#` are comments\n- Lines starting with `-- filename --` begin embedded files\n- Commands are executed in order\n- Use `!` prefix for commands expected to fail\n- Use `stdout` and `stderr` to assert on command output\n\n## Environment Variables\n\nThe harness automatically sets these environment variables:\n\n| Variable | Description |\n|----------|-------------|\n| `$TEST_IMAGE` | Unique Docker image name for test isolation |\n| `$WORK` | Test's temporary working directory |\n| `$SERVER_URL` | URL of running cog server (after `cog serve`) |\n| `$HOME` | Real home directory (for Docker credentials) |\n\nYou can also use:\n\n| Variable | Description |\n|----------|-------------|\n| `COG_BINARY` | Path to cog binary (defaults to auto-build) |\n| `TEST_PARALLEL` | Number of parallel tests (default: 4) |\n\nUse `go test -short` to skip slow tests.\n\n## Custom Commands\n\n### `cog` - Run cog CLI commands\n\n```txtar\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE -i name=value\n! cog predict $TEST_IMAGE -i invalid=input # Expected to fail\n```\n\nSpecial handling for `cog serve`:\n- Runs in background automatically\n- Allocates a random port\n- Waits for health check before continuing\n- Sets `$SERVER_URL` for subsequent commands\n- Cleans up on test completion\n\n### `curl` - Make HTTP requests to cog server\n\n```txtar\ncog serve\ncurl GET /health-check\ncurl POST /predictions '{\"input\":{\"s\":\"hello\"}}'\nstdout '\"output\":\"hello\"'\n```\n\nUsage: `curl METHOD PATH [BODY]`\n\nThe `curl` command includes built-in retry logic (10 attempts, 500ms delay) for resilience against timing issues in integration tests.\n\n### `wait-for` - Wait for conditions\n\n**Note**: This command waits for conditions on the **host machine**, not inside Docker containers. For Docker-based tests, use `curl` instead (which has built-in retry logic).\n\n```txtar\n# Wait for a file to exist (host filesystem only)\nwait-for file output.txt 30s\n\n# Wait for HTTP endpoint (host-accessible URLs only)\nwait-for http http://localhost:8080/health 200 60s\n\n# Wait for file with content\nwait-for not-empty results.json 30s\n```\n\nUsage: `wait-for CONDITION TARGET [ARGS] [TIMEOUT]`\n\n### `pty-run` - Run commands with PTY\n\nRun a command with a pseudo-terminal (PTY), sending input from a file and capturing output.\n\n```txtar\n# Run bash interactively with commands from input file\npty-run input.txt cog run $TEST_IMAGE /bin/bash\nstdout 'expected output'\n\n# Run a simple command (no input needed)\npty-run /dev/null cog run $TEST_IMAGE echo \"hello\"\nstdout 'hello'\n```\n\nUsage: `pty-run [args...]`\n\n- The input file contents are written to the PTY as terminal input\n- Use `/dev/null` if no input is needed\n- Output is captured and can be matched with `stdout` command\n- Uses `github.com/creack/pty` which works on both Linux and macOS\n\n## Conditions\n\nUse conditions to control when tests run based on environment. Conditions are evaluated by the test runner and can be used with `skip` to conditionally skip tests.\n\n### Available Conditions\n\n| Condition | Evaluates to True When | Negated | Example Use Case |\n|-----------|------------------------|---------|------------------|\n| `[short]` | `go test -short` is used | `[!short]` | Use `[short] skip` to skip GPU tests, long builds, or slow framework installs when running in short mode |\n| `[linux]` | Running on Linux | `[!linux]` | Tests requiring Linux-specific features |\n| `[amd64]` | Running on amd64/x86_64 architecture | `[!amd64]` | Tests requiring specific CPU architecture |\n| `[linux_amd64]` | Running on Linux AND amd64 | `[!linux_amd64]` | Tests requiring both Linux and amd64 (e.g., `--use-cog-base-image` builds) |\n\n### Usage Examples\n\n**Skip slow tests:**\n\n```txtar\n[short] skip 'requires GPU or long build time'\n\ncog build -t $TEST_IMAGE\n# ... rest of test\n```\n\nSkip slow tests with: `go test -short ./...`\n\n**Platform-specific tests:**\n\n```txtar\n[!linux] skip 'requires Linux'\n\n# Linux-specific test\ncog build -t $TEST_IMAGE\n```\n\n**Architecture-specific tests:**\n\n```txtar\n[!amd64] skip 'requires amd64 architecture'\n\n# amd64-specific test\ncog build -t $TEST_IMAGE\n```\n\n**Combined platform and architecture:**\n\n```txtar\n[!linux_amd64] skip 'requires Linux on amd64'\n\n# Test that requires both (e.g., --use-cog-base-image builds)\ncog build -t $TEST_IMAGE --use-cog-base-image\n```\n\n### Condition Logic\n\nConditions can be negated with `!`:\n- `[short]` - True when `go test -short` is used\n - Use `[short] skip` to skip a slow test when running in short mode\n- `[!short]` - True when NOT running with `-short` flag\n - Use `[!short] skip` to only run a test in short mode (rare)\n- `[!linux]` - True when NOT on Linux\n - Use `[!linux] skip` to skip non-Linux tests\n- `[linux_amd64]` - True when on Linux AND amd64\n - Use `[!linux_amd64] skip` to skip tests that need this specific platform\n\nMultiple conditions can be used on separate lines:\n\n```txtar\n[short] skip 'requires long build time'\n[!linux] skip 'requires Linux'\n\n# Only runs on Linux when not using -short flag\ncog build -t $TEST_IMAGE\n```\n\n## Built-in Commands\n\nThese are provided by testscript itself:\n\n| Command | Description |\n|---------|-------------|\n| `exec` | Run an arbitrary command |\n| `stdout PATTERN` | Assert stdout matches regex |\n| `stderr PATTERN` | Assert stderr matches regex |\n| `exists FILE` | Assert file exists |\n| `! exists FILE` | Assert file does not exist |\n| `cp SRC DST` | Copy file |\n| `rm FILE` | Remove file |\n| `mkdir DIR` | Create directory |\n| `cd DIR` | Change directory |\n| `env KEY=VALUE` | Set environment variable |\n| `skip MESSAGE` | Skip the test |\n| `stop MESSAGE` | Stop test early (success) |\n\nSee [testscript documentation](https://pkg.go.dev/github.com/rogpeppe/go-internal/testscript) for the full list.\n\n## Test Patterns\n\n### Testing predictions\n\n```txtar\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE -i s=hello\nstdout 'hello'\n```\n\n### Testing server endpoints\n\n```txtar\ncog build -t $TEST_IMAGE\ncog serve\ncurl POST /predictions '{\"input\":{\"s\":\"test\"}}'\nstdout '\"output\":'\n```\n\n### Testing expected failures\n\n```txtar\n# Build should fail without predictor\n! cog build -t $TEST_IMAGE\nstderr 'predict'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.12\"\n# Note: no predict field\n```\n\n### Testing with subprocess initialization\n\n```txtar\ncog build -t $TEST_IMAGE\ncog serve\n\n# curl has built-in retry logic for timing resilience\ncurl POST /predictions '{\"input\":{\"s\":\"test\"}}'\nstdout '\"output\":\"hello test\"'\n\n-- predict.py --\nimport subprocess\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def setup(self):\n self.bg = subprocess.Popen([\"./background.sh\"])\n \n def predict(self, s: str) -> str:\n return \"hello \" + s\n```\n\n### Slow tests (GPU/frameworks)\n\n```txtar\n[fast] skip 'requires long build time'\n\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE\nstdout 'torch'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.12\"\n gpu: true\n python_packages:\n - torch==2.7.0\n```\n\n## How It Works\n\n1. **Test Discovery**: The test runner finds all `.txtar` files in `tests/`\n2. **Setup**: For each test, the harness:\n - Creates a fresh temporary directory\n - Extracts embedded files from the txtar\n - Sets environment variables (`$TEST_IMAGE`, etc.)\n - Registers cleanup (Docker image removal, server shutdown)\n3. **Execution**: Commands run sequentially in the temp directory\n4. **Assertions**: `stdout`/`stderr` commands verify output\n5. **Cleanup**: Docker images are removed, servers are stopped\n\n## Debugging Failed Tests\n\n### View verbose output\n\n```bash\ngo test -v -run TestIntegration/test_name\n```\n\n### Keep work directory\n\n```bash\n# Add to test or set in harness\nenv TESTWORK=1\n```\n\n### Run single test interactively\n\n```bash\ncd integration-tests\ngo test -v -run TestIntegration/string_predictor -timeout 10m\n```\n\n### Check Docker images\n\n```bash\n# List test images (should be cleaned up)\ndocker images | grep cog-test\n```\n\n## Adding New Tests\n\n1. Create a new `.txtar` file in `tests/`\n2. Name it descriptively (e.g., `async_predictor.txtar`)\n3. Add comments explaining what's being tested\n4. Include all necessary fixture files inline\n5. Run your test: `go test -v -run TestIntegration/your_test_name`\n\n## Common Issues\n\n### Test times out waiting for server\n\nThe server health check has a 30-second timeout. If your model takes longer to load:\n- Consider if it should be a `[slow]` test\n- Check for errors in the predictor's `setup()` method\n\n### \"SERVER_URL not set\" error\n\nMake sure `cog serve` is called before `curl`.\n\n### Docker build output cluttering logs\n\nBuild output is suppressed by default (`BUILDKIT_PROGRESS=quiet`). Errors are still shown.\n\n### Files created in container not visible\n\nThe `wait-for file` command checks the **host** filesystem, not inside Docker containers. Use `curl` for Docker-based synchronization (it has built-in retry logic).\n\n### Test works locally but fails in CI\n\n- CI environments may be slower - increase retry attempts\n- Check for hardcoded paths or assumptions about the environment\n- Make sure the test is properly isolated (no shared state)\n" }, { "path": "integration-tests/concurrent/concurrent_test.go", "content": "//go:build integration\n\npackage concurrent_test\n\nimport (\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"net\"\n\t\"net/http\"\n\t\"os\"\n\t\"os/exec\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"sync\"\n\t\"syscall\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/integration-tests/harness\"\n)\n\n// TestConcurrentPredictions tests that concurrent async predictions complete properly\n// with server shutdown.\n//\n// This test verifies:\n// 1. Multiple predictions can run concurrently\n// 2. Server shutdown waits for running predictions to complete\n// 3. All predictions return correct results\n//\n// This test is written in Go (not txtar) because it requires parallel HTTP requests\n// with precise timing coordination that doesn't fit txtar's sequential execution model.\nfunc TestConcurrentPredictions(t *testing.T) {\n\tif testing.Short() {\n\t\tt.Skip(\"skipping slow test in short mode\")\n\t}\n\n\t// Create a temp directory for our test project\n\ttmpDir, err := os.MkdirTemp(\"\", \"cog-concurrent-test-*\")\n\trequire.NoError(t, err, \"failed to create temp dir\")\n\tdefer os.RemoveAll(tmpDir)\n\n\t// Write the async-sleep predictor fixture\n\terr = os.WriteFile(filepath.Join(tmpDir, \"cog.yaml\"), []byte(cogYAML), 0o644)\n\trequire.NoError(t, err, \"failed to write cog.yaml\")\n\terr = os.WriteFile(filepath.Join(tmpDir, \"predict.py\"), []byte(predictPy), 0o644)\n\trequire.NoError(t, err, \"failed to write predict.py\")\n\n\t// Get the cog binary\n\tcogBinary, err := harness.ResolveCogBinary()\n\trequire.NoError(t, err, \"failed to resolve cog binary\")\n\n\t// Generate unique image name\n\timageName := fmt.Sprintf(\"cog-concurrent-test-%d\", time.Now().UnixNano())\n\tdefer func() {\n\t\texec.Command(\"docker\", \"rmi\", \"-f\", imageName).Run()\n\t}()\n\n\t// Build the image\n\tt.Log(\"Building image...\")\n\tbuildCmd := exec.Command(cogBinary, \"build\", \"-t\", imageName)\n\tbuildCmd.Dir = tmpDir\n\tbuildCmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\toutput, err := buildCmd.CombinedOutput()\n\trequire.NoError(t, err, \"failed to build image\\n%s\", output)\n\n\t// Start the server\n\tt.Log(\"Starting server...\")\n\tport, err := allocatePort()\n\trequire.NoError(t, err, \"failed to allocate port\")\n\n\tserveCmd := exec.Command(cogBinary, \"serve\", \"-p\", fmt.Sprintf(\"%d\", port))\n\tserveCmd.Dir = tmpDir\n\tserveCmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\n\terr = serveCmd.Start()\n\trequire.NoError(t, err, \"failed to start server\")\n\tdefer func() {\n\t\tserveCmd.Process.Kill()\n\t\tserveCmd.Wait()\n\t}()\n\n\tserverURL := fmt.Sprintf(\"http://127.0.0.1:%d\", port)\n\n\t// Wait for server to be ready\n\tt.Log(\"Waiting for server to be ready...\")\n\trequire.True(t, waitForServerReady(serverURL, 60*time.Second), \"server did not become ready within timeout\")\n\n\t// Fire 5 concurrent predictions\n\tt.Log(\"Starting concurrent predictions...\")\n\tconst numPredictions = 5\n\tvar wg sync.WaitGroup\n\tresults := make([]predictionResult, numPredictions)\n\n\tstart := time.Now()\n\n\tfor i := range numPredictions {\n\t\twg.Add(1)\n\t\tgo func(idx int) {\n\t\t\tdefer wg.Done()\n\t\t\tresults[idx] = makePrediction(serverURL, idx)\n\t\t}(i)\n\t}\n\n\t// Wait a bit for all predictions to be accepted but not completed\n\ttime.Sleep(200 * time.Millisecond)\n\n\t// Shutdown the server while predictions are in-flight\n\tt.Log(\"Sending shutdown request...\")\n\tshutdownResp, err := http.Post(serverURL+\"/shutdown\", \"application/json\", nil)\n\tif err != nil {\n\t\tt.Logf(\"shutdown request error (may be expected): %v\", err)\n\t} else {\n\t\tshutdownResp.Body.Close()\n\t}\n\n\t// Wait for all predictions to complete\n\twg.Wait()\n\telapsed := time.Since(start)\n\n\tt.Logf(\"All predictions completed in %v\", elapsed)\n\n\t// Verify timing - should be < 3s if running concurrently (each sleeps 1s)\n\tassert.Less(t, elapsed, 3*time.Second, \"predictions took too long (%v), expected < 3s for concurrent execution\", elapsed)\n\n\t// Verify all predictions succeeded with correct output\n\tfor i, result := range results {\n\t\tif !assert.NoError(t, result.err, \"prediction %d failed\", i) {\n\t\t\tcontinue\n\t\t}\n\t\tif !assert.Equal(t, http.StatusOK, result.statusCode, \"prediction %d returned unexpected status\", i) {\n\t\t\tcontinue\n\t\t}\n\t\texpectedOutput := fmt.Sprintf(\"wake up sleepyhead%d\", i)\n\t\tassert.Equal(t, expectedOutput, result.output, \"prediction %d output mismatch\", i)\n\t}\n}\n\ntype predictionResult struct {\n\tstatusCode int\n\toutput string\n\terr error\n}\n\nfunc makePrediction(serverURL string, idx int) predictionResult {\n\treqBody := fmt.Sprintf(`{\"id\":\"id-%d\",\"input\":{\"s\":\"sleepyhead%d\",\"sleep\":1.0}}`, idx, idx)\n\n\tresp, err := http.Post(\n\t\tserverURL+\"/predictions\",\n\t\t\"application/json\",\n\t\tstrings.NewReader(reqBody),\n\t)\n\tif err != nil {\n\t\treturn predictionResult{err: err}\n\t}\n\tdefer resp.Body.Close()\n\n\tvar response struct {\n\t\tOutput string `json:\"output\"`\n\t\tStatus string `json:\"status\"`\n\t}\n\tif err := json.NewDecoder(resp.Body).Decode(&response); err != nil {\n\t\treturn predictionResult{statusCode: resp.StatusCode, err: err}\n\t}\n\n\treturn predictionResult{\n\t\tstatusCode: resp.StatusCode,\n\t\toutput: response.Output,\n\t}\n}\n\nfunc waitForServerReady(serverURL string, timeout time.Duration) bool {\n\tclient := &http.Client{Timeout: 2 * time.Second}\n\tdeadline := time.Now().Add(timeout)\n\n\tfor time.Now().Before(deadline) {\n\t\tresp, err := client.Get(serverURL + \"/health-check\")\n\t\tif err != nil {\n\t\t\ttime.Sleep(200 * time.Millisecond)\n\t\t\tcontinue\n\t\t}\n\n\t\tvar health struct {\n\t\t\tStatus string `json:\"status\"`\n\t\t}\n\t\tif err := json.NewDecoder(resp.Body).Decode(&health); err != nil {\n\t\t\tresp.Body.Close()\n\t\t\ttime.Sleep(200 * time.Millisecond)\n\t\t\tcontinue\n\t\t}\n\t\tresp.Body.Close()\n\n\t\tif health.Status == \"READY\" {\n\t\t\treturn true\n\t\t}\n\t\tif health.Status == \"SETUP_FAILED\" || health.Status == \"DEFUNCT\" {\n\t\t\treturn false\n\t\t}\n\n\t\ttime.Sleep(200 * time.Millisecond)\n\t}\n\n\treturn false\n}\n\n// waitForServerStatus polls /health-check until the server reports the given status.\n// Unlike waitForServerReady which waits for READY, this can wait for intermediate\n// states like STARTING (useful for testing signals during setup).\nfunc waitForServerStatus(serverURL string, targetStatus string, timeout time.Duration) bool {\n\tclient := &http.Client{Timeout: 2 * time.Second}\n\tdeadline := time.Now().Add(timeout)\n\n\tfor time.Now().Before(deadline) {\n\t\tresp, err := client.Get(serverURL + \"/health-check\")\n\t\tif err != nil {\n\t\t\ttime.Sleep(200 * time.Millisecond)\n\t\t\tcontinue\n\t\t}\n\n\t\tvar health struct {\n\t\t\tStatus string `json:\"status\"`\n\t\t}\n\t\tif err := json.NewDecoder(resp.Body).Decode(&health); err != nil {\n\t\t\tresp.Body.Close()\n\t\t\ttime.Sleep(200 * time.Millisecond)\n\t\t\tcontinue\n\t\t}\n\t\tresp.Body.Close()\n\n\t\tif health.Status == targetStatus {\n\t\t\treturn true\n\t\t}\n\t\tif health.Status == \"SETUP_FAILED\" || health.Status == \"DEFUNCT\" {\n\t\t\treturn false\n\t\t}\n\n\t\ttime.Sleep(200 * time.Millisecond)\n\t}\n\n\treturn false\n}\n\n// allocatePort finds an available TCP port by letting the OS assign one.\nfunc allocatePort() (int, error) {\n\tlistener, err := net.Listen(\"tcp\", \"127.0.0.1:0\")\n\tif err != nil {\n\t\treturn 0, err\n\t}\n\tdefer listener.Close()\n\treturn listener.Addr().(*net.TCPAddr).Port, nil\n}\n\n// Embedded fixture files\n\nconst cogYAML = `build:\n python_version: \"3.11\"\npredict: \"predict.py:Predictor\"\nconcurrency:\n max: 5\n`\n\nconst predictPy = `import asyncio\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n async def predict(self, s: str, sleep: float) -> str:\n await asyncio.sleep(sleep)\n return f\"wake up {s}\"\n`\n\n// TestConcurrentAboveLimit tests that sending more predictions than max_concurrency\n// returns a 409 Conflict for the excess prediction.\nfunc TestConcurrentAboveLimit(t *testing.T) {\n\tif testing.Short() {\n\t\tt.Skip(\"skipping slow test in short mode\")\n\t}\n\n\ttmpDir, err := os.MkdirTemp(\"\", \"cog-above-limit-test-*\")\n\trequire.NoError(t, err, \"failed to create temp dir\")\n\tdefer os.RemoveAll(tmpDir)\n\n\terr = os.WriteFile(filepath.Join(tmpDir, \"cog.yaml\"), []byte(aboveLimitCogYAML), 0o644)\n\trequire.NoError(t, err, \"failed to write cog.yaml\")\n\terr = os.WriteFile(filepath.Join(tmpDir, \"predict.py\"), []byte(predictPy), 0o644)\n\trequire.NoError(t, err, \"failed to write predict.py\")\n\n\tcogBinary, err := harness.ResolveCogBinary()\n\trequire.NoError(t, err, \"failed to resolve cog binary\")\n\n\timageName := fmt.Sprintf(\"cog-above-limit-test-%d\", time.Now().UnixNano())\n\tdefer func() {\n\t\texec.Command(\"docker\", \"rmi\", \"-f\", imageName).Run()\n\t}()\n\n\tt.Log(\"Building image...\")\n\tbuildCmd := exec.Command(cogBinary, \"build\", \"-t\", imageName)\n\tbuildCmd.Dir = tmpDir\n\tbuildCmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\toutput, err := buildCmd.CombinedOutput()\n\trequire.NoError(t, err, \"failed to build image\\n%s\", output)\n\n\tt.Log(\"Starting server...\")\n\tport, err := allocatePort()\n\trequire.NoError(t, err, \"failed to allocate port\")\n\n\tserveCmd := exec.Command(cogBinary, \"serve\", \"-p\", fmt.Sprintf(\"%d\", port))\n\tserveCmd.Dir = tmpDir\n\tserveCmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\n\terr = serveCmd.Start()\n\trequire.NoError(t, err, \"failed to start server\")\n\tdefer func() {\n\t\tserveCmd.Process.Kill()\n\t\tserveCmd.Wait()\n\t}()\n\n\tserverURL := fmt.Sprintf(\"http://127.0.0.1:%d\", port)\n\n\tt.Log(\"Waiting for server to be ready...\")\n\trequire.True(t, waitForServerReady(serverURL, 60*time.Second), \"server did not become ready within timeout\")\n\n\t// Fill all 2 slots with long-running predictions (each sleeps 1s)\n\tconst maxConcurrency = 2\n\tvar wg sync.WaitGroup\n\tfor i := range maxConcurrency {\n\t\twg.Add(1)\n\t\tgo func(idx int) {\n\t\t\tdefer wg.Done()\n\t\t\tmakePrediction(serverURL, idx)\n\t\t}(i)\n\t}\n\n\t// Poll with an overflow request until we get a 409, meaning both slots\n\t// are occupied. This avoids a fixed sleep that can flake on slow CI.\n\tt.Log(\"Polling for 409 (all slots occupied)...\")\n\tdeadline := time.Now().Add(10 * time.Second)\n\tvar resp *http.Response\n\tfor time.Now().Before(deadline) {\n\t\textraBody := `{\"id\":\"extra\",\"input\":{\"s\":\"overflow\",\"sleep\":1.0}}`\n\t\tresp, err = http.Post(\n\t\t\tserverURL+\"/predictions\",\n\t\t\t\"application/json\",\n\t\t\tstrings.NewReader(extraBody),\n\t\t)\n\t\trequire.NoError(t, err, \"failed to send extra prediction\")\n\t\tif resp.StatusCode == http.StatusConflict {\n\t\t\tbreak\n\t\t}\n\t\t// Got 200 — slots weren't full yet, close and retry\n\t\tresp.Body.Close()\n\t\ttime.Sleep(100 * time.Millisecond)\n\t}\n\tdefer resp.Body.Close()\n\n\trequire.Equal(t, http.StatusConflict, resp.StatusCode, \"extra prediction status = %d, want %d (409 Conflict); slots never filled within timeout\", resp.StatusCode, http.StatusConflict)\n\n\tvar errResp struct {\n\t\tError string `json:\"error\"`\n\t\tStatus string `json:\"status\"`\n\t}\n\terr = json.NewDecoder(resp.Body).Decode(&errResp)\n\trequire.NoError(t, err, \"failed to decode error response\")\n\tassert.Equal(t, \"failed\", errResp.Status, \"error response status mismatch\")\n\tassert.Contains(t, strings.ToLower(errResp.Error), \"capacity\", \"error response error = %q, want string containing \\\"capacity\\\"\", errResp.Error)\n\n\twg.Wait()\n}\n\nconst aboveLimitCogYAML = `build:\n python_version: \"3.11\"\npredict: \"predict.py:Predictor\"\nconcurrency:\n max: 2\n`\n\n// TestSIGTERMDuringSetup tests that SIGTERM during setup() causes clean shutdown.\nfunc TestSIGTERMDuringSetup(t *testing.T) {\n\tif testing.Short() {\n\t\tt.Skip(\"skipping slow test in short mode\")\n\t}\n\n\ttmpDir, err := os.MkdirTemp(\"\", \"cog-sigterm-setup-test-*\")\n\trequire.NoError(t, err, \"failed to create temp dir\")\n\tdefer os.RemoveAll(tmpDir)\n\n\tslowSetupCogYAML := `build:\n python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n`\n\tslowSetupPredictPy := `import time\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def setup(self) -> None:\n time.sleep(30)\n\n def predict(self, s: str) -> str:\n return \"hello \" + s\n`\n\n\terr = os.WriteFile(filepath.Join(tmpDir, \"cog.yaml\"), []byte(slowSetupCogYAML), 0o644)\n\trequire.NoError(t, err, \"failed to write cog.yaml\")\n\terr = os.WriteFile(filepath.Join(tmpDir, \"predict.py\"), []byte(slowSetupPredictPy), 0o644)\n\trequire.NoError(t, err, \"failed to write predict.py\")\n\n\tcogBinary, err := harness.ResolveCogBinary()\n\trequire.NoError(t, err, \"failed to resolve cog binary\")\n\n\tt.Log(\"Building image...\")\n\timageName := fmt.Sprintf(\"cog-sigterm-setup-test-%d\", time.Now().UnixNano())\n\tdefer func() {\n\t\texec.Command(\"docker\", \"rmi\", \"-f\", imageName).Run()\n\t}()\n\n\tbuildCmd := exec.Command(cogBinary, \"build\", \"-t\", imageName)\n\tbuildCmd.Dir = tmpDir\n\tbuildCmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\toutput, err := buildCmd.CombinedOutput()\n\trequire.NoError(t, err, \"failed to build image\\n%s\", output)\n\n\tt.Log(\"Starting server...\")\n\tport, err := allocatePort()\n\trequire.NoError(t, err, \"failed to allocate port\")\n\n\tserveCmd := exec.Command(cogBinary, \"serve\", \"-p\", fmt.Sprintf(\"%d\", port))\n\tserveCmd.Dir = tmpDir\n\tserveCmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\n\terr = serveCmd.Start()\n\trequire.NoError(t, err, \"failed to start server\")\n\n\t// Poll health-check until setup has begun (status STARTING),\n\t// rather than a fixed sleep that can be too short on cold Docker pulls.\n\tt.Log(\"Waiting for setup to begin (STARTING status)...\")\n\tif !waitForServerStatus(fmt.Sprintf(\"http://127.0.0.1:%d\", port), \"STARTING\", 60*time.Second) {\n\t\tserveCmd.Process.Kill()\n\t\tserveCmd.Wait()\n\t\tt.Fatal(\"server did not reach STARTING status within timeout\")\n\t}\n\n\t// Send SIGTERM\n\tt.Log(\"Sending SIGTERM during setup...\")\n\terr = serveCmd.Process.Signal(syscall.SIGTERM)\n\trequire.NoError(t, err, \"failed to send signal\")\n\n\t// Wait for process to exit with a timeout\n\tdone := make(chan error, 1)\n\tgo func() {\n\t\tdone <- serveCmd.Wait()\n\t}()\n\n\tselect {\n\tcase err := <-done:\n\t\tif err == nil {\n\t\t\tt.Fatal(\"server exited cleanly after SIGTERM; expected termination by signal\")\n\t\t}\n\n\t\tvar exitErr *exec.ExitError\n\t\tif !errors.As(err, &exitErr) {\n\t\t\tt.Fatalf(\"server exited with unexpected error type after SIGTERM: %T (%v)\", err, err)\n\t\t}\n\n\t\tws, ok := exitErr.Sys().(syscall.WaitStatus)\n\t\tif !ok {\n\t\t\tt.Fatalf(\"server exited after SIGTERM but wait status was unavailable: %v\", err)\n\t\t}\n\t\tif !ws.Signaled() || ws.Signal() != syscall.SIGTERM {\n\t\t\tt.Fatalf(\"server exit = %v, want signal %v\", ws, syscall.SIGTERM)\n\t\t}\n\tcase <-time.After(15 * time.Second):\n\t\tserveCmd.Process.Kill()\n\t\tt.Fatal(\"server did not exit within 15s after SIGTERM\")\n\t}\n}\n" }, { "path": "integration-tests/harness/cmd_pty.go", "content": "package harness\n\nimport (\n\t\"bytes\"\n\t\"io\"\n\t\"os\"\n\t\"os/exec\"\n\t\"strings\"\n\t\"sync\"\n\t\"time\"\n\n\t\"github.com/creack/pty\"\n\t\"github.com/rogpeppe/go-internal/testscript\"\n)\n\n// PtyRunCommand implements the 'pty-run' command for testscript.\ntype PtyRunCommand struct {\n\tharness *Harness\n}\n\nfunc (c *PtyRunCommand) Name() string { return \"pty-run\" }\n\n// Run executes a command with a PTY, sending input from a file and capturing output.\n//\n// Usage: pty-run [args...]\n//\n// The input file contents are written to the PTY as terminal input.\n// Use /dev/null or an empty file if no input is needed.\n// The command's output is written to stdout for matching with 'stdout' command.\n//\n// This uses github.com/creack/pty which works on both Linux and macOS,\n// unlike testscript's native ttyin/ttyout which hangs on macOS due to\n// Go bug https://github.com/golang/go/issues/61779.\n//\n// TODO: Remove this implementation and use testscript's native ttyin/ttyout\n// once the Go bug is fixed (check Go 1.26+).\nfunc (c *PtyRunCommand) Run(ts *testscript.TestScript, neg bool, args []string) {\n\tif len(args) < 2 {\n\t\tts.Fatalf(\"pty-run: usage: pty-run [args...]\")\n\t}\n\n\tinputFile := args[0]\n\tcmdName := args[1]\n\tcmdArgs := args[2:]\n\n\t// Read input file\n\tvar input string\n\tif inputFile != \"/dev/null\" {\n\t\tinput = ts.ReadFile(inputFile)\n\t}\n\n\t// Expand environment variables in command and args\n\tcmdName = os.Expand(cmdName, ts.Getenv)\n\n\t// Handle \"cog\" command specially - use the resolved binary\n\tif cmdName == \"cog\" {\n\t\tcmdName = c.harness.CogBinary\n\t}\n\n\texpandedArgs := make([]string, len(cmdArgs))\n\tfor i, arg := range cmdArgs {\n\t\texpandedArgs[i] = os.Expand(arg, ts.Getenv)\n\t}\n\n\t// Create the command\n\tcmd := exec.Command(cmdName, expandedArgs...)\n\tcmd.Dir = ts.Getenv(\"WORK\")\n\n\t// Build environment\n\tcmd.Env = []string{\n\t\t\"HOME=\" + ts.Getenv(\"HOME\"),\n\t\t\"PATH=\" + ts.Getenv(\"PATH\"),\n\t\t\"COG_NO_UPDATE_CHECK=1\",\n\t}\n\tif v := ts.Getenv(\"REPO_ROOT\"); v != \"\" {\n\t\tcmd.Env = append(cmd.Env, \"REPO_ROOT=\"+v)\n\t}\n\tif v := ts.Getenv(\"COG_SDK_WHEEL\"); v != \"\" {\n\t\tcmd.Env = append(cmd.Env, \"COG_SDK_WHEEL=\"+v)\n\t}\n\tif v := ts.Getenv(\"COGLET_WHEEL\"); v != \"\" {\n\t\tcmd.Env = append(cmd.Env, \"COGLET_WHEEL=\"+v)\n\t}\n\n\t// Start command with PTY\n\tptmx, err := pty.Start(cmd)\n\tif err != nil {\n\t\tts.Fatalf(\"pty-run: failed to start command with PTY: %v\", err)\n\t}\n\tdefer ptmx.Close()\n\n\t// Set terminal size\n\tif err := pty.Setsize(ptmx, &pty.Winsize{Rows: 24, Cols: 80}); err != nil {\n\t\tts.Logf(\"pty-run: failed to set terminal size: %v\", err)\n\t}\n\n\t// Use shared buffer pattern for reading (avoids race conditions)\n\tvar buf bytes.Buffer\n\tvar mu sync.Mutex\n\tdone := make(chan struct{})\n\n\tgo func() {\n\t\ttmp := make([]byte, 1024)\n\t\tfor {\n\t\t\tselect {\n\t\t\tcase <-done:\n\t\t\t\treturn\n\t\t\tdefault:\n\t\t\t\tn, err := ptmx.Read(tmp)\n\t\t\t\tif n > 0 {\n\t\t\t\t\tmu.Lock()\n\t\t\t\t\tbuf.Write(tmp[:n])\n\t\t\t\t\tmu.Unlock()\n\t\t\t\t}\n\t\t\t\tif err != nil {\n\t\t\t\t\tif err != io.EOF {\n\t\t\t\t\t\t// Log non-EOF errors but don't fail - PTY may close unexpectedly\n\t\t\t\t\t\tts.Logf(\"pty-run: read error (may be normal): %v\", err)\n\t\t\t\t\t}\n\t\t\t\t\treturn\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}()\n\n\t// Write input to PTY with small delays between lines for reliability\n\tif input != \"\" {\n\t\tfor line := range strings.SplitSeq(input, \"\\n\") {\n\t\t\tif line == \"\" {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\t_, err := ptmx.Write([]byte(line + \"\\n\"))\n\t\t\tif err != nil {\n\t\t\t\tts.Logf(\"pty-run: failed to write input (may be normal if command exited): %v\", err)\n\t\t\t\tbreak\n\t\t\t}\n\t\t\t// Small delay to let the shell process the line\n\t\t\ttime.Sleep(50 * time.Millisecond)\n\t\t}\n\t}\n\n\t// Wait for command to finish with timeout\n\tcmdDone := make(chan error, 1)\n\tgo func() {\n\t\tcmdDone <- cmd.Wait()\n\t}()\n\n\ttimeout := 60 * time.Second\n\tvar cmdErr error\n\n\tselect {\n\tcase cmdErr = <-cmdDone:\n\t\t// Command finished\n\tcase <-time.After(timeout):\n\t\t// Timeout - kill the process\n\t\t_ = cmd.Process.Kill()\n\t\tmu.Lock()\n\t\toutput := buf.String()\n\t\tmu.Unlock()\n\t\tts.Logf(\"pty-run: timeout after %v, partial output: %q\", timeout, output)\n\t\tts.Fatalf(\"pty-run: command timed out after %v\", timeout)\n\t\treturn\n\t}\n\n\t// Give a moment for final output to be captured\n\ttime.Sleep(100 * time.Millisecond)\n\tclose(done)\n\n\t// Get final output\n\tmu.Lock()\n\toutput := buf.String()\n\tmu.Unlock()\n\n\t// Handle negation\n\tif neg {\n\t\tif cmdErr == nil {\n\t\t\tts.Fatalf(\"pty-run: command succeeded unexpectedly\")\n\t\t}\n\t\t// Command failed as expected - write output for potential pattern matching\n\t\t_, _ = ts.Stdout().Write([]byte(output))\n\t\treturn\n\t}\n\n\tif cmdErr != nil {\n\t\tts.Logf(\"pty-run: command output: %q\", output)\n\t\tts.Fatalf(\"pty-run: command failed: %v\", cmdErr)\n\t}\n\n\t// Write output to stdout for pattern matching\n\t_, _ = ts.Stdout().Write([]byte(output))\n}\n" }, { "path": "integration-tests/harness/command.go", "content": "package harness\n\nimport \"github.com/rogpeppe/go-internal/testscript\"\n\n// Command defines the interface for testscript commands.\ntype Command interface {\n\t// Name returns the command name as used in txtar scripts.\n\tName() string\n\n\t// Run executes the command.\n\t// neg is true if the command was prefixed with '!' (expecting failure).\n\t// args are the command arguments.\n\tRun(ts *testscript.TestScript, neg bool, args []string)\n}\n\n// CommandFunc adapts a function to the Command interface.\ntype CommandFunc struct {\n\tname string\n\tfn func(ts *testscript.TestScript, neg bool, args []string)\n}\n\nfunc (c CommandFunc) Name() string { return c.name }\nfunc (c CommandFunc) Run(ts *testscript.TestScript, neg bool, args []string) {\n\tc.fn(ts, neg, args)\n}\n\n// NewCommand creates a Command from a name and function.\nfunc NewCommand(name string, fn func(ts *testscript.TestScript, neg bool, args []string)) Command {\n\treturn CommandFunc{name: name, fn: fn}\n}\n" }, { "path": "integration-tests/harness/harness.go", "content": "// Package harness provides utilities for running cog integration tests.\npackage harness\n\nimport (\n\t\"context\"\n\tcryptorand \"crypto/rand\"\n\t\"crypto/sha256\"\n\t\"encoding/hex\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"io\"\n\tmathrand \"math/rand/v2\"\n\t\"net\"\n\t\"net/http\"\n\t\"os\"\n\t\"os/exec\"\n\t\"path/filepath\"\n\t\"strconv\"\n\t\"strings\"\n\t\"sync\"\n\t\"time\"\n\n\t\"github.com/google/go-containerregistry/pkg/crane\"\n\t\"github.com/rogpeppe/go-internal/testscript\"\n\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/registry_testhelpers\"\n)\n\n// propagatedEnvVars lists host environment variables that should be propagated\n// into testscript environments (Setup) and background processes (cmdCogServe).\n// Keep this list in sync: if you add a new env var to propagate, add it here.\nvar propagatedEnvVars = []string{\n\t\"COG_SDK_WHEEL\", // SDK wheel override\n\t\"COGLET_WHEEL\", // coglet wheel override\n\t\"RUST_LOG\", // Rust logging control\n\t\"COG_CA_CERT\", // custom CA certificates (e.g. Cloudflare WARP)\n\t\"BUILDKIT_PROGRESS\", // Docker build output format\n}\n\n// Harness provides utilities for running cog integration tests.\n// serverInfo tracks a running cog serve process and its port\ntype serverInfo struct {\n\tcmd *exec.Cmd\n\tport int\n}\n\n// registryInfo tracks a running test registry container\ntype registryInfo struct {\n\tcontainer *registry_testhelpers.RegistryContainer\n\tcleanup func()\n\thost string // e.g., \"localhost:5432\"\n}\n\n// mockUploadRecord records a single upload received by the mock upload server.\ntype mockUploadRecord struct {\n\tPath string\n\tContentType string\n\tSize int\n}\n\n// mockUploadServer is a lightweight HTTP server that accepts PUT requests\n// and records what was uploaded.\ntype mockUploadServer struct {\n\tserver *http.Server\n\tport int\n\tmu sync.Mutex\n\tuploads []mockUploadRecord\n}\n\n// webhookResult is the summary written to stdout by webhook-server-wait.\ntype webhookResult struct {\n\tStatus string `json:\"status\"`\n\tOutputSize int `json:\"output_size\"`\n\tHasError bool `json:\"has_error\"`\n\tErrorMessage string `json:\"error_message,omitempty\"`\n\tMetrics json.RawMessage `json:\"metrics,omitempty\"`\n}\n\n// webhookServer accepts prediction webhook callbacks from coglet.\n// It parses the JSON payload to extract status and output size, without\n// ever exposing the (potentially huge) output to testscript's log buffer.\ntype webhookServer struct {\n\tserver *http.Server\n\tport int\n\tmu sync.Mutex\n\tresult *webhookResult\n\tdone chan struct{} // closed on first terminal webhook\n}\n\ntype Harness struct {\n\tCogBinary string\n\t// realHome is captured at creation time before testscript overrides HOME\n\trealHome string\n\t// repoRoot is the path to the cog repository root\n\trepoRoot string\n\t// serverProcs tracks background cog serve processes for cleanup, keyed by work directory\n\tserverProcs map[string]*serverInfo\n\tserverProcsMu sync.Mutex\n\t// registries tracks test registry containers for cleanup, keyed by work directory\n\tregistries map[string]*registryInfo\n\tregistriesMu sync.Mutex\n\t// uploadServers tracks mock upload servers for cleanup, keyed by work directory\n\tuploadServers map[string]*mockUploadServer\n\tuploadServersMu sync.Mutex\n\t// webhookServers tracks webhook receiver servers for cleanup, keyed by work directory\n\twebhookServers map[string]*webhookServer\n\twebhookServersMu sync.Mutex\n}\n\n// New creates a new Harness, resolving the cog binary location.\nfunc New() (*Harness, error) {\n\tcogBinary, err := ResolveCogBinary()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\trepoRoot, err := findRepoRoot()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn &Harness{\n\t\tCogBinary: cogBinary,\n\t\trealHome: os.Getenv(\"HOME\"),\n\t\trepoRoot: repoRoot,\n\t\tserverProcs: make(map[string]*serverInfo),\n\t\tregistries: make(map[string]*registryInfo),\n\t\tuploadServers: make(map[string]*mockUploadServer),\n\t\twebhookServers: make(map[string]*webhookServer),\n\t}, nil\n}\n\n// ResolveCogBinary finds the cog binary to use for tests.\n// It checks (in order):\n// 1. COG_BINARY environment variable\n// 2. Build from source (if in cog repository)\nfunc ResolveCogBinary() (string, error) {\n\tif cogBinary := os.Getenv(\"COG_BINARY\"); cogBinary != \"\" {\n\t\tif !filepath.IsAbs(cogBinary) {\n\t\t\t// Resolve relative paths from repo root, not the test package directory.\n\t\t\trepoRoot, err := findRepoRoot()\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", err\n\t\t\t}\n\t\t\tcogBinary = filepath.Join(repoRoot, cogBinary)\n\t\t}\n\t\treturn cogBinary, nil\n\t}\n\n\t// Build from source\n\treturn buildCogBinary()\n}\n\n// buildCogBinary builds the cog binary from source.\n// It finds the repository root, builds wheels if needed, and compiles the binary.\n// If the binary already exists, it returns the cached path.\nfunc buildCogBinary() (string, error) {\n\t// Find repository root (where go.mod with module github.com/replicate/cog exists)\n\trepoRoot, err := findRepoRoot()\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to find cog repository root: %w\", err)\n\t}\n\n\t// Check if binary already exists\n\tbinPath := filepath.Join(repoRoot, \"integration-tests\", \".bin\", \"cog\")\n\tif _, err := os.Stat(binPath); err == nil {\n\t\tfmt.Printf(\"Using cached cog binary: %s\\n\", binPath)\n\t\treturn binPath, nil\n\t}\n\n\t// Check if wheels exist, build if not\n\tvar (\n\t\twheelsDir = filepath.Join(repoRoot, \"pkg\", \"wheels\")\n\t\tcogWheelExists, _ = filepath.Glob(filepath.Join(wheelsDir, \"cog-*.whl\"))\n\t\tcogletWheelExists, _ = filepath.Glob(filepath.Join(wheelsDir, \"coglet-*.whl\"))\n\t)\n\n\tif len(cogWheelExists) == 0 || len(cogletWheelExists) == 0 {\n\t\tfmt.Println(\"Building Python wheels...\")\n\t\tif err := runCommand(repoRoot, \"mise\", \"run\", \"build:wheels\"); err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to build wheels: %w\", err)\n\t\t}\n\n\t\tfmt.Println(\"Generating wheel embeds...\")\n\t\tif err := runCommand(repoRoot, \"go\", \"generate\", \"./pkg/wheels\"); err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to generate wheel embeds: %w\", err)\n\t\t}\n\t}\n\n\t// Build the cog binary\n\tif err := os.MkdirAll(filepath.Dir(binPath), 0o755); err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to create bin directory: %w\", err)\n\t}\n\n\tfmt.Println(\"Building cog binary...\")\n\tif err := runCommand(repoRoot, \"go\", \"build\", \"-o\", binPath, \"./cmd/cog\"); err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to build cog: %w\", err)\n\t}\n\n\treturn binPath, nil\n}\n\n// findRepoRoot finds the cog repository root by looking for go.mod with the main module\nfunc findRepoRoot() (string, error) {\n\t// Start from current working directory\n\tdir, err := os.Getwd()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tfor {\n\t\tgoMod := filepath.Join(dir, \"go.mod\")\n\t\tif _, err := os.Stat(goMod); err == nil {\n\t\t\t// Verify it's the cog repo root (matches the expected module path)\n\t\t\tcontent, err := os.ReadFile(goMod)\n\t\t\tif err == nil && strings.Contains(string(content), \"module github.com/replicate/cog\\n\") {\n\t\t\t\treturn dir, nil\n\t\t\t}\n\t\t}\n\n\t\tparent := filepath.Dir(dir)\n\t\tif parent == dir {\n\t\t\tbreak\n\t\t}\n\t\tdir = parent\n\t}\n\n\treturn \"\", fmt.Errorf(\"could not find cog repository root\")\n}\n\n// runCommand runs a command in the specified directory\nfunc runCommand(dir string, name string, args ...string) error {\n\tcmd := exec.Command(name, args...)\n\tcmd.Dir = dir\n\tcmd.Stdout = os.Stdout\n\tcmd.Stderr = os.Stderr\n\treturn cmd.Run()\n}\n\n// Commands returns the custom testscript commands provided by this harness.\nfunc (h *Harness) Commands() map[string]func(ts *testscript.TestScript, neg bool, args []string) {\n\t// Register all commands\n\tcommands := []Command{\n\t\t// Built-in commands (defined in this file)\n\t\tNewCommand(\"cog\", h.cmdCog),\n\t\tNewCommand(\"curl\", h.cmdCurl),\n\t\tNewCommand(\"wait-for\", h.cmdWaitFor),\n\t\tNewCommand(\"docker-run\", h.cmdDockerRun),\n\n\t\t// Registry and OCI bundle testing commands\n\t\tNewCommand(\"registry-start\", h.cmdRegistryStart),\n\t\tNewCommand(\"registry-inspect\", h.cmdRegistryInspect),\n\t\tNewCommand(\"registry-seed\", h.cmdRegistrySeed),\n\t\tNewCommand(\"docker-push\", h.cmdDockerPush),\n\t\tNewCommand(\"mock-weights\", h.cmdMockWeights),\n\n\t\t// Mock upload server commands\n\t\tNewCommand(\"upload-server-start\", h.cmdUploadServerStart),\n\t\tNewCommand(\"upload-server-count\", h.cmdUploadServerCount),\n\n\t\t// Webhook receiver commands\n\t\tNewCommand(\"webhook-server-start\", h.cmdWebhookServerStart),\n\t\tNewCommand(\"webhook-server-wait\", h.cmdWebhookServerWait),\n\n\t\t// PTY command (defined in cmd_pty.go)\n\t\t&PtyRunCommand{harness: h},\n\t}\n\n\t// Build the command map\n\tresult := make(map[string]func(ts *testscript.TestScript, neg bool, args []string))\n\tfor _, cmd := range commands {\n\t\tresult[cmd.Name()] = cmd.Run\n\t}\n\treturn result\n}\n\n// cmdCog implements the 'cog' command for testscript.\n// It handles all cog subcommands, with special handling for certain commands.\nfunc (h *Harness) cmdCog(ts *testscript.TestScript, neg bool, args []string) {\n\t// Check for subcommands that need special handling\n\tif len(args) > 0 && args[0] == \"serve\" {\n\t\t// Special handling for 'cog serve' - run in background\n\t\th.cmdCogServe(ts, neg, args[1:])\n\t\treturn\n\t}\n\n\t// Default: run cog command normally\n\texpandedArgs := make([]string, len(args))\n\tfor i, arg := range args {\n\t\texpandedArgs[i] = os.Expand(arg, ts.Getenv)\n\t}\n\n\terr := ts.Exec(h.CogBinary, expandedArgs...)\n\tif neg {\n\t\tif err == nil {\n\t\t\tts.Fatalf(\"cog command succeeded unexpectedly\")\n\t\t}\n\t\treturn\n\t}\n\n\tif err != nil {\n\t\tts.Fatalf(\"cog command failed: %v\", err)\n\t}\n}\n\n// Setup returns a testscript Setup function that configures the test environment.\n// Fixtures are embedded in the txtar files themselves, so no file copying is needed.\nfunc (h *Harness) Setup(env *testscript.Env) error {\n\t// Restore real HOME for Docker credential helpers.\n\t// Docker credential helpers (e.g., docker-credential-desktop) need the real HOME\n\t// to access the macOS keychain.\n\tenv.Setenv(\"HOME\", h.realHome)\n\n\t// Export repo root for tests that need to reference files outside the work directory\n\tenv.Setenv(\"REPO_ROOT\", h.repoRoot)\n\n\t// Disable update checks during tests\n\tenv.Setenv(\"COG_NO_UPDATE_CHECK\", \"1\")\n\n\t// Propagate host env vars listed in propagatedEnvVars\n\tfor _, key := range propagatedEnvVars {\n\t\tif val := os.Getenv(key); val != \"\" {\n\t\t\tenv.Setenv(key, val)\n\t\t}\n\t}\n\n\t// Auto-detect wheels from dist/ if not explicitly set via env vars.\n\t// CI sets these env vars; locally we need to find them ourselves.\n\tdistDir := filepath.Join(h.repoRoot, \"dist\")\n\tif os.Getenv(\"COGLET_WHEEL\") == \"\" {\n\t\tif matches, _ := filepath.Glob(filepath.Join(distDir, \"coglet-*.whl\")); len(matches) > 0 {\n\t\t\tenv.Setenv(\"COGLET_WHEEL\", distDir)\n\t\t}\n\t}\n\tif os.Getenv(\"COG_SDK_WHEEL\") == \"\" {\n\t\tif matches, _ := filepath.Glob(filepath.Join(distDir, \"cog-*.whl\")); len(matches) > 0 {\n\t\t\tenv.Setenv(\"COG_SDK_WHEEL\", distDir)\n\t\t}\n\t}\n\n\t// Generate unique image name for this test run\n\timageName := generateUniqueImageName()\n\tenv.Setenv(\"TEST_IMAGE\", imageName)\n\n\t// Capture the work directory for this test (used as key for server tracking)\n\tworkDir := env.WorkDir\n\n\t// Register cleanup to remove the Docker image, stop servers, and cleanup registries\n\tenv.Defer(func() {\n\t\t// Stop the server for this specific test (if any)\n\t\th.stopServerByWorkDir(workDir)\n\t\t// Stop the registry for this specific test (if any)\n\t\th.stopRegistryByWorkDir(workDir)\n\t\t// Stop the upload server for this specific test (if any)\n\t\th.stopUploadServerByWorkDir(workDir)\n\t\t// Stop the webhook server for this specific test (if any)\n\t\th.stopWebhookServerByWorkDir(workDir)\n\t\tremoveDockerImage(imageName)\n\t})\n\n\treturn nil\n}\n\n// generateUniqueImageName creates a unique Docker image name for test isolation.\nfunc generateUniqueImageName() string {\n\tb := make([]byte, 5)\n\tif _, err := cryptorand.Read(b); err != nil {\n\t\t// Fall back to a less random but still unique name\n\t\treturn fmt.Sprintf(\"cog-test-%d\", os.Getpid())\n\t}\n\treturn fmt.Sprintf(\"cog-test-%s\", hex.EncodeToString(b))\n}\n\n// removeDockerImage attempts to remove a Docker image by name.\n// It silently ignores errors (image may not exist if test failed early).\nfunc removeDockerImage(imageName string) {\n\t// Remove all images that match the prefix (base and final images)\n\tcmd := exec.Command(\"docker\", \"images\", \"--format\", \"{{.Repository}}:{{.Tag}}\", \"--filter\", fmt.Sprintf(\"reference=%s*\", imageName))\n\toutput, err := cmd.Output()\n\tif err != nil {\n\t\treturn\n\t}\n\n\tfor img := range strings.SplitSeq(strings.TrimSpace(string(output)), \"\\n\") {\n\t\tif img == \"\" {\n\t\t\tcontinue\n\t\t}\n\t\texec.Command(\"docker\", \"rmi\", \"-f\", img).Run() //nolint:errcheck,gosec\n\t}\n}\n\n// cmdCogServe implements background 'cog serve' for testscript.\n// It starts a cog serve process in the background and waits for it to be healthy.\n// Usage: cog serve [flags]\n// Exports $SERVER_URL environment variable with the server address.\nfunc (h *Harness) cmdCogServe(ts *testscript.TestScript, neg bool, args []string) {\n\tworkDir := ts.Getenv(\"WORK\")\n\n\t// Check if server is already running\n\th.serverProcsMu.Lock()\n\tif _, exists := h.serverProcs[workDir]; exists {\n\t\th.serverProcsMu.Unlock()\n\t\tts.Fatalf(\"server already running\")\n\t}\n\th.serverProcsMu.Unlock()\n\n\t// Allocate a random available port\n\tport, err := allocatePort()\n\tif err != nil {\n\t\tts.Fatalf(\"failed to allocate port: %v\", err)\n\t}\n\n\t// Build command arguments\n\tcmdArgs := []string{\"serve\", \"-p\", strconv.Itoa(port)}\n\tcmdArgs = append(cmdArgs, args...)\n\n\t// Expand environment variables in arguments\n\texpandedArgs := make([]string, len(cmdArgs))\n\tfor i, arg := range cmdArgs {\n\t\texpandedArgs[i] = os.Expand(arg, ts.Getenv)\n\t}\n\n\t// Start the server process\n\tcmd := exec.Command(h.CogBinary, expandedArgs...)\n\tcmd.Dir = workDir\n\n\t// Build environment from testscript.\n\t// Always include core vars, plus everything from propagatedEnvVars.\n\tvar env []string\n\tfor _, key := range []string{\"HOME\", \"PATH\", \"REPO_ROOT\", \"COG_NO_UPDATE_CHECK\", \"TEST_IMAGE\"} {\n\t\tif val := ts.Getenv(key); val != \"\" {\n\t\t\tenv = append(env, key+\"=\"+val)\n\t\t}\n\t}\n\tfor _, key := range propagatedEnvVars {\n\t\tif val := ts.Getenv(key); val != \"\" {\n\t\t\tenv = append(env, key+\"=\"+val)\n\t\t}\n\t}\n\tcmd.Env = env\n\n\t// Capture server output for debugging\n\tcmd.Stdout = ts.Stdout()\n\tcmd.Stderr = ts.Stderr()\n\n\tif err := cmd.Start(); err != nil {\n\t\tts.Fatalf(\"failed to start server: %v\", err)\n\t}\n\n\t// Store the process for cleanup (keyed by work directory)\n\th.serverProcsMu.Lock()\n\th.serverProcs[workDir] = &serverInfo{cmd: cmd, port: port}\n\th.serverProcsMu.Unlock()\n\n\t// Wait for server to be healthy\n\tserverURL := fmt.Sprintf(\"http://127.0.0.1:%d\", port)\n\tts.Setenv(\"SERVER_URL\", serverURL)\n\n\tif !waitForServer(serverURL, 60*time.Second) {\n\t\tif neg {\n\t\t\t// Test expected the server to fail setup — keep it running\n\t\t\t// so the test can inspect the health-check status.\n\t\t\treturn\n\t\t}\n\t\t// Try to get server output for debugging\n\t\t_ = cmd.Process.Kill()\n\t\tts.Fatalf(\"server did not become healthy within timeout\")\n\t}\n\n\tif neg {\n\t\tts.Fatalf(\"server became healthy, but expected setup failure\")\n\t}\n}\n\n// cmdCurl implements the 'curl' command for testscript.\n// It makes HTTP requests to the server started with 'serve'.\n// Includes built-in retry logic (10 attempts, 500ms delay) for resilience.\n// Usage: curl [-H key:value]... [method] [path] [body]\n// Examples:\n//\n//\tcurl GET /health-check\n//\tcurl POST /predictions '{\"input\":{\"s\":\"hello\"}}'\n//\tcurl -H Prefer:respond-async POST /predictions '{\"input\":{}}'\nfunc (h *Harness) cmdCurl(ts *testscript.TestScript, neg bool, args []string) {\n\tif len(args) < 2 {\n\t\tts.Fatalf(\"curl: usage: curl [-H key:value]... [method] [path] [body | @file]\")\n\t}\n\n\t// Parse -H flags for extra headers\n\tvar extraHeaders [][2]string\n\tfor len(args) >= 2 && args[0] == \"-H\" {\n\t\tkv := args[1]\n\t\tparts := strings.SplitN(kv, \":\", 2)\n\t\tif len(parts) != 2 {\n\t\t\tts.Fatalf(\"curl: invalid header %q (expected key:value)\", kv)\n\t\t}\n\t\textraHeaders = append(extraHeaders, [2]string{\n\t\t\tstrings.TrimSpace(parts[0]),\n\t\t\tstrings.TrimSpace(parts[1]),\n\t\t})\n\t\targs = args[2:]\n\t}\n\n\tif len(args) < 2 {\n\t\tts.Fatalf(\"curl: usage: curl [-H key:value]... [method] [path] [body | @file]\")\n\t}\n\n\tserverURL := ts.Getenv(\"SERVER_URL\")\n\tif serverURL == \"\" {\n\t\tts.Fatalf(\"curl: SERVER_URL not set (did you call 'cog serve' first?)\")\n\t}\n\n\tmethod := args[0]\n\tpath := args[1]\n\tvar body string\n\tif len(args) > 2 {\n\t\tbody = os.Expand(args[2], ts.Getenv)\n\t\tif strings.HasPrefix(body, \"@\") {\n\t\t\tfilename := body[1:]\n\t\t\tdata, err := os.ReadFile(ts.MkAbs(filename))\n\t\t\tif err != nil {\n\t\t\t\tts.Fatalf(\"curl: failed to read body file %q: %v\", filename, err)\n\t\t\t}\n\t\t\tbody = string(data)\n\t\t}\n\t}\n\n\t// Retry settings\n\tmaxAttempts := 10\n\tretryDelay := 500 * time.Millisecond\n\n\tclient := &http.Client{Timeout: 10 * time.Second}\n\n\tvar (\n\t\tlastErr error\n\t\tlastStatus int\n\t\tlastBody string\n\t)\n\n\tfor attempt := 1; attempt <= maxAttempts; attempt++ {\n\t\treq, err := http.NewRequest(method, serverURL+path, strings.NewReader(body))\n\t\tif err != nil {\n\t\t\tlastErr = err\n\t\t\ttime.Sleep(retryDelay)\n\t\t\tcontinue\n\t\t}\n\n\t\tif body != \"\" {\n\t\t\treq.Header.Set(\"Content-Type\", \"application/json\")\n\t\t}\n\t\tfor _, h := range extraHeaders {\n\t\t\treq.Header.Set(h[0], h[1])\n\t\t}\n\n\t\tresp, err := client.Do(req) //nolint:gosec // G704: URL from test harness, not user input\n\t\tif err != nil {\n\t\t\tlastErr = err\n\t\t\ttime.Sleep(retryDelay)\n\t\t\tcontinue\n\t\t}\n\n\t\t// Read response body\n\t\trespBodyBytes, readErr := io.ReadAll(resp.Body)\n\t\tif readErr != nil {\n\t\t\tts.Fatalf(\"curl: failed to read response: %v\", readErr)\n\t\t}\n\t\trespBody := string(respBodyBytes)\n\t\t_ = resp.Body.Close()\n\n\t\tlastStatus = resp.StatusCode\n\t\tlastBody = respBody\n\t\tlastErr = nil\n\n\t\t// Check if this is a successful response\n\t\tstatusOK := resp.StatusCode >= 200 && resp.StatusCode < 300\n\n\t\tif neg {\n\t\t\tif !statusOK {\n\t\t\t\t// Expected to fail — write body to stderr so tests can assert\n\t\t\t\t_, _ = ts.Stderr().Write([]byte(respBody))\n\t\t\t\treturn\n\t\t\t}\n\t\t} else {\n\t\t\tif statusOK {\n\t\t\t\t// Success - write body to stdout\n\t\t\t\t_, _ = ts.Stdout().Write([]byte(respBody))\n\t\t\t\treturn\n\t\t\t}\n\t\t}\n\n\t\t// If this isn't the last attempt, wait before retrying\n\t\tif attempt < maxAttempts {\n\t\t\ttime.Sleep(retryDelay)\n\t\t}\n\t}\n\n\t// All attempts failed\n\tif neg {\n\t\tts.Fatalf(\"curl: expected failure but got status %d after %d attempts\", lastStatus, maxAttempts)\n\t\treturn\n\t}\n\n\tif lastErr != nil {\n\t\tts.Fatalf(\"curl: all %d attempts failed with error: %v\", maxAttempts, lastErr)\n\t\treturn\n\t}\n\n\terrorMsg := lastBody\n\tif len(errorMsg) > 500 {\n\t\terrorMsg = errorMsg[:500] + \"...\"\n\t}\n\tts.Logf(\"curl: full response body: %s\", lastBody)\n\tts.Fatalf(\"curl: all %d attempts failed with status %d: %s\", maxAttempts, lastStatus, errorMsg)\n}\n\n// StopServer stops the background server process for a test script.\nfunc (h *Harness) StopServer(ts *testscript.TestScript) {\n\tworkDir := ts.Getenv(\"WORK\")\n\th.stopServerByWorkDir(workDir)\n}\n\n// stopServerByWorkDir stops the server process associated with a work directory.\nfunc (h *Harness) stopServerByWorkDir(workDir string) {\n\th.serverProcsMu.Lock()\n\tinfo, exists := h.serverProcs[workDir]\n\tif !exists {\n\t\th.serverProcsMu.Unlock()\n\t\treturn\n\t}\n\tdelete(h.serverProcs, workDir)\n\th.serverProcsMu.Unlock()\n\n\t// Try graceful shutdown first via /shutdown endpoint\n\tserverURL := fmt.Sprintf(\"http://127.0.0.1:%d\", info.port)\n\tshutdownURL := serverURL + \"/shutdown\"\n\tresp, err := http.Post(shutdownURL, \"application/json\", nil) //nolint:gosec,noctx\n\tif err == nil {\n\t\t_ = resp.Body.Close()\n\t}\n\n\t// Force kill the cog process if still running\n\tif info.cmd.Process != nil {\n\t\t_ = info.cmd.Process.Kill()\n\t}\n\t_ = info.cmd.Wait()\n\n\t// Also kill any Docker container that may still be running on this port\n\t// Find container by port and kill it\n\toutput, err := exec.Command(\"docker\", \"ps\", \"-q\", \"--filter\", fmt.Sprintf(\"publish=%d\", info.port)).Output()\n\tif err == nil && len(output) > 0 {\n\t\tcontainerID := strings.TrimSpace(string(output))\n\t\tif containerID != \"\" {\n\t\t\texec.Command(\"docker\", \"kill\", containerID).Run() //nolint:errcheck,gosec\n\t\t}\n\t}\n}\n\n// allocatePort finds an available TCP port.\nfunc allocatePort() (int, error) {\n\tlistener, err := net.Listen(\"tcp\", \"127.0.0.1:0\")\n\tif err != nil {\n\t\treturn 0, err\n\t}\n\tdefer listener.Close()\n\treturn listener.Addr().(*net.TCPAddr).Port, nil\n}\n\n// healthCheckResponse represents the JSON response from /health-check\ntype healthCheckResponse struct {\n\tStatus string `json:\"status\"`\n}\n\n// waitForServer polls the server's health-check endpoint until it returns READY status.\n// The server may return HTTP 200 while still in STARTING state (during setup),\n// so we must check the actual status field in the response.\nfunc waitForServer(serverURL string, timeout time.Duration) bool {\n\tclient := &http.Client{Timeout: 5 * time.Second}\n\tdeadline := time.Now().Add(timeout)\n\n\tfor time.Now().Before(deadline) {\n\t\tresp, err := client.Get(serverURL + \"/health-check\")\n\t\tif err != nil {\n\t\t\ttime.Sleep(200 * time.Millisecond)\n\t\t\tcontinue\n\t\t}\n\n\t\tif resp.StatusCode == http.StatusOK {\n\t\t\tbody, err := io.ReadAll(resp.Body)\n\t\t\t_ = resp.Body.Close()\n\t\t\tif err != nil {\n\t\t\t\ttime.Sleep(200 * time.Millisecond)\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tvar health healthCheckResponse\n\t\t\tif err := json.Unmarshal(body, &health); err != nil {\n\t\t\t\ttime.Sleep(200 * time.Millisecond)\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\t// Return success when the server has completed setup\n\t\t\t// READY = setup completed, healthcheck passed (or no healthcheck)\n\t\t\t// UNHEALTHY = setup completed, but user healthcheck failed\n\t\t\t// BUSY = setup completed, prediction in progress\n\t\t\tif health.Status == \"READY\" || health.Status == \"UNHEALTHY\" || health.Status == \"BUSY\" {\n\t\t\t\treturn true\n\t\t\t}\n\n\t\t\t// If setup failed, no point waiting\n\t\t\tif health.Status == \"SETUP_FAILED\" || health.Status == \"DEFUNCT\" {\n\t\t\t\treturn false\n\t\t\t}\n\t\t} else {\n\t\t\t_ = resp.Body.Close()\n\t\t}\n\n\t\ttime.Sleep(200 * time.Millisecond)\n\t}\n\n\treturn false\n}\n\n// cmdWaitFor implements the 'wait-for' command for testscript.\n// It waits for a specific condition to become true with retries.\n// Usage:\n//\n//\twait-for file [timeout] - Wait for file to exist\n//\twait-for http [status] [timeout] - Wait for HTTP endpoint\n//\twait-for not-empty [timeout] - Wait for file with content\nfunc (h *Harness) cmdWaitFor(ts *testscript.TestScript, neg bool, args []string) {\n\tif len(args) < 2 {\n\t\tts.Fatalf(\"wait-for: usage: wait-for [file|http|not-empty] [timeout]\")\n\t}\n\n\tvar (\n\t\tcondition = args[0]\n\t\ttarget = args[1]\n\n\t\t// Default timeout of 30 seconds, can be overridden\n\t\ttimeout = 30 * time.Second\n\t)\n\n\tif len(args) > 2 {\n\t\tif duration, err := time.ParseDuration(args[len(args)-1]); err == nil {\n\t\t\ttimeout = duration\n\t\t}\n\t}\n\n\tdeadline := time.Now().Add(timeout)\n\n\tfor time.Now().Before(deadline) {\n\t\tvar conditionMet bool\n\n\t\tswitch condition {\n\t\tcase \"file\":\n\t\t\t// Wait for file to exist\n\t\t\ttargetPath := filepath.Join(ts.Getenv(\"WORK\"), target)\n\t\t\t_, err := os.Stat(targetPath)\n\t\t\tconditionMet = err == nil\n\n\t\tcase \"not-empty\":\n\t\t\t// Wait for file to exist with non-empty content\n\t\t\ttargetPath := filepath.Join(ts.Getenv(\"WORK\"), target)\n\t\t\tdata, err := os.ReadFile(targetPath)\n\t\t\tconditionMet = err == nil && len(data) > 0\n\n\t\tcase \"http\":\n\t\t\t// Wait for HTTP endpoint to return expected status\n\t\t\texpectedStatus := http.StatusOK\n\t\t\tif len(args) > 2 {\n\t\t\t\tif status, err := strconv.Atoi(args[2]); err == nil {\n\t\t\t\t\texpectedStatus = status\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tclient := &http.Client{Timeout: 2 * time.Second}\n\t\t\tresp, err := client.Get(target)\n\t\t\tif err == nil {\n\t\t\t\tconditionMet = resp.StatusCode == expectedStatus\n\t\t\t\t_ = resp.Body.Close()\n\t\t\t}\n\n\t\tdefault:\n\t\t\tts.Fatalf(\"wait-for: unknown condition: %s\", condition)\n\t\t}\n\n\t\tif neg {\n\t\t\t// For negation, we want the condition to remain false\n\t\t\tif !conditionMet {\n\t\t\t\treturn\n\t\t\t}\n\t\t} else {\n\t\t\t// Normal case: condition should become true\n\t\t\tif conditionMet {\n\t\t\t\treturn\n\t\t\t}\n\t\t}\n\n\t\ttime.Sleep(200 * time.Millisecond)\n\t}\n\n\tif neg {\n\t\tts.Fatalf(\"wait-for: condition became true (expected to remain false)\")\n\t\treturn\n\t}\n\n\tts.Fatalf(\"wait-for: timeout waiting for condition: %s %s\", condition, target)\n}\n\n// cmdDockerRun implements the 'docker-run' command for testscript.\n// It runs a command inside a Docker container.\n// Usage:\n//\n//\tdocker-run [args...]\n//\n// The container is run with:\n// - --rm (auto-remove after exit)\n// - --add-host=host.docker.internal:host-gateway (for Linux compatibility)\n// - Working directory mounted if needed\nfunc (h *Harness) cmdDockerRun(ts *testscript.TestScript, neg bool, args []string) {\n\tif len(args) < 2 {\n\t\tts.Fatalf(\"docker-run: usage: docker-run [args...]\")\n\t}\n\n\tvar (\n\t\timage = os.Expand(args[0], ts.Getenv)\n\t\tcontainerArgs = make([]string, len(args)-1)\n\t)\n\n\tfor i, arg := range args[1:] {\n\t\tcontainerArgs[i] = os.Expand(arg, ts.Getenv)\n\t}\n\n\t// Build docker run command\n\tdockerArgs := []string{\n\t\t\"run\", \"--rm\",\n\t\t\"--add-host=host.docker.internal:host-gateway\",\n\t\timage,\n\t}\n\tdockerArgs = append(dockerArgs, containerArgs...)\n\n\tcmd := exec.Command(\"docker\", dockerArgs...)\n\tcmd.Stdout = ts.Stdout()\n\tcmd.Stderr = ts.Stderr()\n\n\terr := cmd.Run()\n\tif neg {\n\t\tif err == nil {\n\t\t\tts.Fatalf(\"docker-run: command succeeded unexpectedly\")\n\t\t}\n\t\treturn\n\t}\n\n\tif err != nil {\n\t\tts.Fatalf(\"docker-run: command failed: %v\", err)\n\t}\n}\n\n// =============================================================================\n// Registry commands\n// =============================================================================\n\n// cmdRegistryStart starts a test registry container.\n// The registry is automatically cleaned up when the test ends.\n// Usage: registry-start\n// Exports $TEST_REGISTRY environment variable with the registry address.\nfunc (h *Harness) cmdRegistryStart(ts *testscript.TestScript, neg bool, args []string) {\n\tif neg {\n\t\tts.Fatalf(\"registry-start: does not support negation\")\n\t}\n\n\tworkDir := ts.Getenv(\"WORK\")\n\n\t// Check if registry is already running (idempotent)\n\th.registriesMu.Lock()\n\tif info, exists := h.registries[workDir]; exists {\n\t\th.registriesMu.Unlock()\n\t\t// Already started, just ensure env is set\n\t\tts.Setenv(\"TEST_REGISTRY\", info.host)\n\t\treturn\n\t}\n\th.registriesMu.Unlock()\n\n\t// Start new registry\n\tcontainer, cleanup, err := registry_testhelpers.StartTestRegistryWithCleanup(context.Background())\n\tif err != nil {\n\t\tts.Fatalf(\"registry-start: failed to start registry: %v\", err)\n\t}\n\n\thost := container.RegistryHost()\n\n\t// Store for cleanup\n\th.registriesMu.Lock()\n\th.registries[workDir] = ®istryInfo{\n\t\tcontainer: container,\n\t\tcleanup: cleanup,\n\t\thost: host,\n\t}\n\th.registriesMu.Unlock()\n\n\tts.Setenv(\"TEST_REGISTRY\", host)\n\tts.Logf(\"registry-start: started registry at %s\", host)\n}\n\n// stopRegistryByWorkDir stops the registry container associated with a work directory.\nfunc (h *Harness) stopRegistryByWorkDir(workDir string) {\n\th.registriesMu.Lock()\n\tinfo, exists := h.registries[workDir]\n\tif !exists {\n\t\th.registriesMu.Unlock()\n\t\treturn\n\t}\n\tdelete(h.registries, workDir)\n\th.registriesMu.Unlock()\n\n\tif info.cleanup != nil {\n\t\tinfo.cleanup()\n\t}\n}\n\n// cmdRegistrySeed copies an image into the test registry under a new repository:tag.\n// The source can be a local reference (relative to $TEST_REGISTRY) or an absolute\n// reference to an external registry (e.g., docker.io/library/python:3.12-slim).\n// The destination is always relative to $TEST_REGISTRY.\n//\n// Usage: registry-seed \n// Examples:\n//\n//\tregistry-seed alpine:latest cog-base:cuda11.8-python3.10-torch2.0.1\n//\tregistry-seed docker.io/library/python:3.12-slim cog-base:python3.12\nfunc (h *Harness) cmdRegistrySeed(ts *testscript.TestScript, neg bool, args []string) {\n\tif neg {\n\t\tts.Fatalf(\"registry-seed: does not support negation\")\n\t}\n\tif len(args) < 2 {\n\t\tts.Fatalf(\"registry-seed: usage: registry-seed \")\n\t}\n\n\tsrc := os.Expand(args[0], ts.Getenv)\n\tdst := os.Expand(args[1], ts.Getenv)\n\n\ttestRegistry := ts.Getenv(\"TEST_REGISTRY\")\n\tif testRegistry == \"\" {\n\t\tts.Fatalf(\"registry-seed: TEST_REGISTRY not set (call registry-start first)\")\n\t}\n\n\t// If the source looks like an absolute reference (contains a registry host\n\t// with a dot, e.g. \"docker.io/library/python:3.12-slim\"), use it as-is.\n\t// Otherwise treat it as relative to the test registry.\n\tsrcRef := src\n\tif !isAbsoluteImageRef(src) {\n\t\tsrcRef = testRegistry + \"/\" + src\n\t}\n\tdstRef := testRegistry + \"/\" + dst\n\n\tif err := crane.Copy(srcRef, dstRef, crane.Insecure); err != nil {\n\t\tts.Fatalf(\"registry-seed: failed to copy %s to %s: %v\", srcRef, dstRef, err)\n\t}\n\n\tts.Logf(\"registry-seed: copied %s to %s\", srcRef, dstRef)\n}\n\n// isAbsoluteImageRef returns true if ref looks like it contains an explicit\n// registry host (e.g. \"docker.io/library/python:3.12-slim\" or\n// \"ghcr.io/foo/bar:latest\"). It checks whether the part before the first\n// slash contains a dot or a colon (port), which distinguishes a registry\n// host from a simple repository name like \"alpine:latest\".\nfunc isAbsoluteImageRef(ref string) bool {\n\thost, _, ok := strings.Cut(ref, \"/\")\n\tif !ok {\n\t\treturn false\n\t}\n\treturn strings.Contains(host, \".\") || strings.Contains(host, \":\")\n}\n\n// cmdRegistryInspect inspects a registry manifest and outputs JSON.\n// Usage: registry-inspect \n// Outputs the manifest result as JSON to stdout.\nfunc (h *Harness) cmdRegistryInspect(ts *testscript.TestScript, neg bool, args []string) {\n\tif len(args) < 1 {\n\t\tts.Fatalf(\"registry-inspect: usage: registry-inspect \")\n\t}\n\n\timageRef := os.Expand(args[0], ts.Getenv)\n\n\tclient := registry.NewRegistryClient()\n\tresult, err := client.Inspect(context.Background(), imageRef, nil)\n\n\tif neg {\n\t\tif err == nil {\n\t\t\tts.Fatalf(\"registry-inspect: expected failure but succeeded\")\n\t\t}\n\t\treturn\n\t}\n\n\tif err != nil {\n\t\tts.Fatalf(\"registry-inspect: failed to inspect %s: %v\", imageRef, err)\n\t}\n\n\t// Output as JSON\n\toutput, err := json.MarshalIndent(result, \"\", \" \")\n\tif err != nil {\n\t\tts.Fatalf(\"registry-inspect: failed to marshal result: %v\", err)\n\t}\n\n\t_, _ = ts.Stdout().Write(output)\n\t_, _ = ts.Stdout().Write([]byte(\"\\n\"))\n}\n\n// cmdDockerPush tags and pushes a local image to the test registry.\n// Usage: docker-push \n// Example: docker-push $TEST_IMAGE test/mymodel:v1\n// The image is pushed to $TEST_REGISTRY/\nfunc (h *Harness) cmdDockerPush(ts *testscript.TestScript, neg bool, args []string) {\n\tif len(args) < 2 {\n\t\tts.Fatalf(\"docker-push: usage: docker-push \")\n\t}\n\n\tlocalImage := os.Expand(args[0], ts.Getenv)\n\trepoTag := os.Expand(args[1], ts.Getenv)\n\n\ttestRegistry := ts.Getenv(\"TEST_REGISTRY\")\n\tif testRegistry == \"\" {\n\t\tts.Fatalf(\"docker-push: TEST_REGISTRY not set (call registry-start first)\")\n\t}\n\n\tremoteRef := testRegistry + \"/\" + repoTag\n\n\t// Tag the image\n\ttagCmd := exec.Command(\"docker\", \"tag\", localImage, remoteRef)\n\ttagCmd.Stdout = ts.Stdout()\n\ttagCmd.Stderr = ts.Stderr()\n\tif err := tagCmd.Run(); err != nil {\n\t\tif neg {\n\t\t\treturn\n\t\t}\n\t\tts.Fatalf(\"docker-push: failed to tag image: %v\", err)\n\t}\n\n\t// Push the image\n\tpushCmd := exec.Command(\"docker\", \"push\", remoteRef)\n\tpushCmd.Stdout = ts.Stdout()\n\tpushCmd.Stderr = ts.Stderr()\n\terr := pushCmd.Run()\n\n\tif neg {\n\t\tif err == nil {\n\t\t\tts.Fatalf(\"docker-push: expected failure but succeeded\")\n\t\t}\n\t\treturn\n\t}\n\n\tif err != nil {\n\t\tts.Fatalf(\"docker-push: failed to push image: %v\", err)\n\t}\n\n\tts.Logf(\"docker-push: pushed %s to %s\", localImage, remoteRef)\n}\n\n// =============================================================================\n// Mock weights command\n// =============================================================================\n\n// mockWeightsLock mirrors the structure from pkg/model/weights_lock.go\n// SYNC: If pkg/model/WeightsLock changes, update this copy.\n// We duplicate it here to avoid importing pkg/model which transitively imports pkg/wheels.\ntype mockWeightsLock struct {\n\tVersion string `json:\"version\"`\n\tCreated time.Time `json:\"created\"`\n\tFiles []mockWeightFile `json:\"files\"`\n}\n\n// mockWeightFile mirrors WeightFile from pkg/model/weights.go\n// SYNC: If pkg/model/WeightFile changes, update this copy.\ntype mockWeightFile struct {\n\tName string `json:\"name\"`\n\tDest string `json:\"dest\"`\n\tDigestOriginal string `json:\"digestOriginal\"`\n\tDigest string `json:\"digest\"`\n\tSize int64 `json:\"size\"`\n\tSizeUncompressed int64 `json:\"sizeUncompressed\"`\n\tMediaType string `json:\"mediaType\"`\n\tContentType string `json:\"contentType,omitempty\"`\n}\n\n// cmdMockWeights generates mock weight files and a weights.lock file.\n// Usage: mock-weights [--count N] [--min-size S] [--max-size S]\n// Defaults:\n// - count: 2\n// - min-size: 1kb\n// - max-size: 10kb\n//\n// Creates files in $WORK/weights/ and writes $WORK/weights.lock\nfunc (h *Harness) cmdMockWeights(ts *testscript.TestScript, neg bool, args []string) {\n\tif neg {\n\t\tts.Fatalf(\"mock-weights: does not support negation\")\n\t}\n\n\t// Parse arguments\n\tcount := 2\n\tminSize := int64(1024) // 1KB\n\tmaxSize := int64(10 * 1024) // 10KB\n\n\tfor i := 0; i < len(args); i++ {\n\t\tswitch args[i] {\n\t\tcase \"--count\", \"-n\":\n\t\t\tif i+1 < len(args) {\n\t\t\t\tif n, err := strconv.Atoi(args[i+1]); err == nil {\n\t\t\t\t\tcount = n\n\t\t\t\t}\n\t\t\t\ti++\n\t\t\t}\n\t\tcase \"--min-size\":\n\t\t\tif i+1 < len(args) {\n\t\t\t\tif size, err := parseSize(args[i+1]); err == nil {\n\t\t\t\t\tminSize = size\n\t\t\t\t}\n\t\t\t\ti++\n\t\t\t}\n\t\tcase \"--max-size\":\n\t\t\tif i+1 < len(args) {\n\t\t\t\tif size, err := parseSize(args[i+1]); err == nil {\n\t\t\t\t\tmaxSize = size\n\t\t\t\t}\n\t\t\t\ti++\n\t\t\t}\n\t\t}\n\t}\n\n\tworkDir := ts.Getenv(\"WORK\")\n\tweightsDir := filepath.Join(workDir, \"weights\")\n\tlockPath := filepath.Join(workDir, \"weights.lock\")\n\n\t// Create weights directory\n\tif err := os.MkdirAll(weightsDir, 0o755); err != nil {\n\t\tts.Fatalf(\"mock-weights: failed to create weights dir: %v\", err)\n\t}\n\n\tvar files []mockWeightFile\n\n\tfor i := 1; i <= count; i++ {\n\t\t// Random size between min and max\n\t\tsize := minSize\n\t\tif maxSize > minSize {\n\t\t\tsize = minSize + mathrand.Int64N(maxSize-minSize+1) //nolint:gosec // test data, not security-sensitive\n\t\t}\n\n\t\t// Generate identifier (e.g., \"weights-001\")\n\t\tweightName := fmt.Sprintf(\"weights-%03d\", i)\n\t\tfilename := weightName + \".bin\"\n\t\tfilePath := filepath.Join(weightsDir, filename)\n\n\t\t// Generate random data\n\t\tdata := make([]byte, size)\n\t\tif _, err := cryptorand.Read(data); err != nil {\n\t\t\tts.Fatalf(\"mock-weights: failed to generate random data: %v\", err)\n\t\t}\n\n\t\t// Write file\n\t\tif err := os.WriteFile(filePath, data, 0o644); err != nil {\n\t\t\tts.Fatalf(\"mock-weights: failed to write %s: %v\", filename, err)\n\t\t}\n\n\t\t// Compute digest (uncompressed, since we're not actually compressing for tests)\n\t\thash := sha256.Sum256(data)\n\t\tdigest := \"sha256:\" + hex.EncodeToString(hash[:])\n\n\t\tfiles = append(files, mockWeightFile{\n\t\t\tName: weightName,\n\t\t\tDest: \"/cache/\" + filename,\n\t\t\tDigestOriginal: digest,\n\t\t\tDigest: digest, // Same as original since we're not compressing\n\t\t\tSize: size,\n\t\t\tSizeUncompressed: size,\n\t\t\t// MediaType matches production WeightBuilder output (uncompressed).\n\t\t\tMediaType: \"application/vnd.cog.weight.layer.v1\",\n\t\t\tContentType: \"application/octet-stream\",\n\t\t})\n\t}\n\n\t// Create weights.lock\n\tlock := mockWeightsLock{\n\t\tVersion: \"1.0\",\n\t\tCreated: time.Now().UTC(),\n\t\tFiles: files,\n\t}\n\n\tlockData, err := json.MarshalIndent(lock, \"\", \" \")\n\tif err != nil {\n\t\tts.Fatalf(\"mock-weights: failed to marshal weights.lock: %v\", err)\n\t}\n\n\tif err := os.WriteFile(lockPath, lockData, 0o644); err != nil {\n\t\tts.Fatalf(\"mock-weights: failed to write weights.lock: %v\", err)\n\t}\n\n\tts.Logf(\"mock-weights: created %d files in %s\", count, weightsDir)\n}\n\n// parseSize parses size strings like \"1kb\", \"10KB\", \"1mb\" into bytes.\nfunc parseSize(s string) (int64, error) {\n\ts = strings.TrimSpace(strings.ToLower(s))\n\tif s == \"\" {\n\t\treturn 0, fmt.Errorf(\"empty size string\")\n\t}\n\n\tvar multiplier int64 = 1\n\tvar numStr string\n\n\tswitch {\n\tcase strings.HasSuffix(s, \"gb\"):\n\t\tmultiplier = 1024 * 1024 * 1024\n\t\tnumStr = strings.TrimSuffix(s, \"gb\")\n\tcase strings.HasSuffix(s, \"mb\"):\n\t\tmultiplier = 1024 * 1024\n\t\tnumStr = strings.TrimSuffix(s, \"mb\")\n\tcase strings.HasSuffix(s, \"kb\"):\n\t\tmultiplier = 1024\n\t\tnumStr = strings.TrimSuffix(s, \"kb\")\n\tcase strings.HasSuffix(s, \"b\"):\n\t\tnumStr = strings.TrimSuffix(s, \"b\")\n\tdefault:\n\t\tnumStr = s\n\t}\n\n\tnum, err := strconv.ParseFloat(strings.TrimSpace(numStr), 64)\n\tif err != nil {\n\t\treturn 0, fmt.Errorf(\"invalid number: %s\", numStr)\n\t}\n\tif num < 0 {\n\t\treturn 0, fmt.Errorf(\"size cannot be negative\")\n\t}\n\n\treturn int64(num * float64(multiplier)), nil\n}\n\n// =============================================================================\n// Mock upload server commands\n// =============================================================================\n\n// cmdUploadServerStart starts a mock HTTP upload server on the host.\n// It accepts PUT requests, records them, and responds with a Location header.\n// Usage: upload-server-start\n// Exports $UPLOAD_SERVER_URL with the server's base URL.\nfunc (h *Harness) cmdUploadServerStart(ts *testscript.TestScript, neg bool, args []string) {\n\tif neg {\n\t\tts.Fatalf(\"upload-server-start: does not support negation\")\n\t}\n\n\tworkDir := ts.Getenv(\"WORK\")\n\n\th.uploadServersMu.Lock()\n\tif _, exists := h.uploadServers[workDir]; exists {\n\t\th.uploadServersMu.Unlock()\n\t\tts.Fatalf(\"upload-server-start: server already running for this test\")\n\t}\n\th.uploadServersMu.Unlock()\n\n\tmus := &mockUploadServer{}\n\n\tmux := http.NewServeMux()\n\tmux.HandleFunc(\"/\", func(w http.ResponseWriter, r *http.Request) {\n\t\tif r.Method != http.MethodPut {\n\t\t\thttp.Error(w, \"method not allowed\", http.StatusMethodNotAllowed)\n\t\t\treturn\n\t\t}\n\t\tbody, _ := io.ReadAll(r.Body)\n\t\trecord := mockUploadRecord{\n\t\t\tPath: r.URL.Path,\n\t\t\tContentType: r.Header.Get(\"Content-Type\"),\n\t\t\tSize: len(body),\n\t\t}\n\t\tmus.mu.Lock()\n\t\tmus.uploads = append(mus.uploads, record)\n\t\tmus.mu.Unlock()\n\n\t\t// Return a clean URL without query params (simulates a signed URL redirect)\n\t\tlocation := fmt.Sprintf(\"http://host.docker.internal:%d%s\", mus.port, r.URL.Path)\n\t\tw.Header().Set(\"Location\", location)\n\t\tw.WriteHeader(http.StatusOK)\n\t})\n\n\tmus.server = &http.Server{Handler: mux, ReadHeaderTimeout: 10 * time.Second} //nolint:gosec // test harness, not production\n\n\t// Bind to all interfaces so the container can reach us via host.docker.internal\n\tln, err := net.Listen(\"tcp\", \"0.0.0.0:0\") //nolint:gosec // must be reachable from Docker container\n\tif err != nil {\n\t\tts.Fatalf(\"upload-server-start: failed to listen: %v\", err)\n\t}\n\tmus.port = ln.Addr().(*net.TCPAddr).Port\n\n\tgo func() { _ = mus.server.Serve(ln) }()\n\n\th.uploadServersMu.Lock()\n\th.uploadServers[workDir] = mus\n\th.uploadServersMu.Unlock()\n\n\t// Advertise host.docker.internal so the container can reach the host server.\n\t// On Linux, cog serve adds --add-host=host.docker.internal:host-gateway.\n\t// On Mac, Docker Desktop resolves host.docker.internal automatically.\n\turl := fmt.Sprintf(\"http://host.docker.internal:%d/\", mus.port)\n\tts.Setenv(\"UPLOAD_SERVER_URL\", url)\n\tts.Logf(\"upload-server-start: listening on 0.0.0.0:%d, container URL: %s\", mus.port, url)\n}\n\n// cmdUploadServerCount verifies exactly N uploads were received.\n// Usage: upload-server-count N\nfunc (h *Harness) cmdUploadServerCount(ts *testscript.TestScript, neg bool, args []string) {\n\tif len(args) != 1 {\n\t\tts.Fatalf(\"upload-server-count: usage: upload-server-count N\")\n\t}\n\n\texpected, err := strconv.Atoi(args[0])\n\tif err != nil {\n\t\tts.Fatalf(\"upload-server-count: invalid count %q: %v\", args[0], err)\n\t}\n\n\tworkDir := ts.Getenv(\"WORK\")\n\th.uploadServersMu.Lock()\n\tmus, exists := h.uploadServers[workDir]\n\th.uploadServersMu.Unlock()\n\n\tif !exists {\n\t\tts.Fatalf(\"upload-server-count: no upload server running (call upload-server-start first)\")\n\t}\n\n\tmus.mu.Lock()\n\tgot := len(mus.uploads)\n\tmus.mu.Unlock()\n\n\tif neg {\n\t\tif got == expected {\n\t\t\tts.Fatalf(\"upload-server-count: expected NOT %d uploads but got %d\", expected, got)\n\t\t}\n\t\treturn\n\t}\n\n\tif got != expected {\n\t\tts.Fatalf(\"upload-server-count: expected %d uploads but got %d\", expected, got)\n\t}\n}\n\n// stopUploadServerByWorkDir shuts down the upload server for a work directory.\nfunc (h *Harness) stopUploadServerByWorkDir(workDir string) {\n\th.uploadServersMu.Lock()\n\tmus, exists := h.uploadServers[workDir]\n\tif !exists {\n\t\th.uploadServersMu.Unlock()\n\t\treturn\n\t}\n\tdelete(h.uploadServers, workDir)\n\th.uploadServersMu.Unlock()\n\n\tif mus.server != nil {\n\t\t_ = mus.server.Close()\n\t}\n}\n\n// =============================================================================\n// Webhook receiver commands\n// =============================================================================\n\n// cmdWebhookServerStart starts a webhook receiver that accepts prediction callbacks.\n// It parses the JSON payload to extract status and measure the output size, without\n// ever exposing the (potentially huge) output to testscript's log buffer.\n// Usage: webhook-server-start\n// Exports $WEBHOOK_URL with the server's callback URL.\nfunc (h *Harness) cmdWebhookServerStart(ts *testscript.TestScript, neg bool, args []string) {\n\tif neg {\n\t\tts.Fatalf(\"webhook-server-start: does not support negation\")\n\t}\n\n\tworkDir := ts.Getenv(\"WORK\")\n\n\th.webhookServersMu.Lock()\n\tif _, exists := h.webhookServers[workDir]; exists {\n\t\th.webhookServersMu.Unlock()\n\t\tts.Fatalf(\"webhook-server-start: server already running for this test\")\n\t}\n\th.webhookServersMu.Unlock()\n\n\tws := &webhookServer{\n\t\tdone: make(chan struct{}),\n\t}\n\n\tmux := http.NewServeMux()\n\tmux.HandleFunc(\"/\", func(w http.ResponseWriter, r *http.Request) {\n\t\tif r.Method != http.MethodPost {\n\t\t\thttp.Error(w, \"method not allowed\", http.StatusMethodNotAllowed)\n\t\t\treturn\n\t\t}\n\n\t\t// Stream-parse the JSON to extract status, measure output size, and\n\t\t// capture metrics without holding the entire output string in memory.\n\t\t// Output is json.RawMessage because it can be a string (single output)\n\t\t// or an array (iterator/streaming output).\n\t\tvar payload struct {\n\t\t\tStatus string `json:\"status\"`\n\t\t\tOutput json.RawMessage `json:\"output\"`\n\t\t\tError string `json:\"error\"`\n\t\t\tMetrics json.RawMessage `json:\"metrics\"`\n\t\t}\n\t\tif err := json.NewDecoder(r.Body).Decode(&payload); err != nil {\n\t\t\thttp.Error(w, \"bad json\", http.StatusBadRequest)\n\t\t\treturn\n\t\t}\n\n\t\tw.WriteHeader(http.StatusOK)\n\n\t\t// Only record terminal statuses\n\t\tswitch payload.Status {\n\t\tcase \"succeeded\", \"failed\", \"canceled\":\n\t\tdefault:\n\t\t\treturn\n\t\t}\n\n\t\tws.mu.Lock()\n\t\tdefer ws.mu.Unlock()\n\n\t\t// Only record the first terminal callback\n\t\tif ws.result != nil {\n\t\t\treturn\n\t\t}\n\t\t// Compute output size: for strings, use the unquoted length;\n\t\t// for arrays or other types, use the raw JSON byte length.\n\t\toutputSize := len(payload.Output)\n\t\tvar outputStr string\n\t\tif json.Unmarshal(payload.Output, &outputStr) == nil {\n\t\t\toutputSize = len(outputStr)\n\t\t}\n\n\t\tws.result = &webhookResult{\n\t\t\tStatus: payload.Status,\n\t\t\tOutputSize: outputSize,\n\t\t\tHasError: payload.Error != \"\",\n\t\t\tErrorMessage: payload.Error,\n\t\t\tMetrics: payload.Metrics,\n\t\t}\n\t\tclose(ws.done)\n\t})\n\n\tws.server = &http.Server{Handler: mux, ReadHeaderTimeout: 10 * time.Second} //nolint:gosec\n\n\t// Bind to all interfaces so the container can reach us via host.docker.internal\n\tln, err := net.Listen(\"tcp\", \"0.0.0.0:0\") //nolint:gosec\n\tif err != nil {\n\t\tts.Fatalf(\"webhook-server-start: failed to listen: %v\", err)\n\t}\n\tws.port = ln.Addr().(*net.TCPAddr).Port\n\n\tgo func() { _ = ws.server.Serve(ln) }()\n\n\th.webhookServersMu.Lock()\n\th.webhookServers[workDir] = ws\n\th.webhookServersMu.Unlock()\n\n\turl := fmt.Sprintf(\"http://host.docker.internal:%d/webhook\", ws.port)\n\tts.Setenv(\"WEBHOOK_URL\", url)\n\tts.Logf(\"webhook-server-start: listening on 0.0.0.0:%d, container URL: %s\", ws.port, url)\n}\n\n// cmdWebhookServerWait blocks until the webhook server receives a terminal prediction callback,\n// then writes a compact JSON summary to stdout for assertion with stdout/stderr matchers.\n// Usage: webhook-server-wait [timeout]\n// Default timeout: 120s\nfunc (h *Harness) cmdWebhookServerWait(ts *testscript.TestScript, neg bool, args []string) {\n\tif neg {\n\t\tts.Fatalf(\"webhook-server-wait: does not support negation\")\n\t}\n\n\ttimeout := 120 * time.Second\n\tif len(args) > 0 {\n\t\tif d, err := time.ParseDuration(args[0]); err == nil {\n\t\t\ttimeout = d\n\t\t}\n\t}\n\n\tworkDir := ts.Getenv(\"WORK\")\n\th.webhookServersMu.Lock()\n\tws, exists := h.webhookServers[workDir]\n\th.webhookServersMu.Unlock()\n\n\tif !exists {\n\t\tts.Fatalf(\"webhook-server-wait: no webhook server running (call webhook-server-start first)\")\n\t}\n\n\tselect {\n\tcase <-ws.done:\n\tcase <-time.After(timeout):\n\t\tts.Fatalf(\"webhook-server-wait: timed out after %s waiting for terminal webhook\", timeout)\n\t}\n\n\tws.mu.Lock()\n\tresult := ws.result\n\tws.mu.Unlock()\n\n\tout, _ := json.Marshal(result)\n\t_, _ = ts.Stdout().Write(out)\n\t_, _ = ts.Stdout().Write([]byte(\"\\n\"))\n}\n\n// stopWebhookServerByWorkDir shuts down the webhook server for a work directory.\nfunc (h *Harness) stopWebhookServerByWorkDir(workDir string) {\n\th.webhookServersMu.Lock()\n\tws, exists := h.webhookServers[workDir]\n\tif !exists {\n\t\th.webhookServersMu.Unlock()\n\t\treturn\n\t}\n\tdelete(h.webhookServers, workDir)\n\th.webhookServersMu.Unlock()\n\n\tif ws.server != nil {\n\t\t_ = ws.server.Close()\n\t}\n}\n" }, { "path": "integration-tests/login/login_test.go", "content": "//go:build integration\n\n// Package login provides integration tests for the cog login command.\n//\n// These tests verify:\n// - Generic registry login with username/password (PTY-based)\n// - Provider routing based on --registry flag\n// - Help text and CLI flags\n//\n// This test file is written in Go (not txtar) because:\n// - Login requires interactive input (PTY for generic provider)\n// - We need fine-grained control over stdin and stdout\n//\n// Note: Replicate provider token verification is tested in unit tests\n// (pkg/provider/replicate/replicate_test.go) since mocking the r8.im\n// hostname requires DNS-level changes not suitable for integration tests.\npackage login_test\n\nimport (\n\t\"bytes\"\n\t\"os\"\n\t\"os/exec\"\n\t\"runtime\"\n\t\"strings\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/creack/pty\"\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/integration-tests/harness\"\n)\n\n// TestLoginGenericRegistryPTY tests interactive login to a generic registry.\n// This test uses PTY to simulate interactive terminal input.\nfunc TestLoginGenericRegistryPTY(t *testing.T) {\n\tif testing.Short() {\n\t\tt.Skip(\"skipping in short mode\")\n\t}\n\n\t// PTY tests only work reliably on Unix-like systems\n\tif runtime.GOOS == \"windows\" {\n\t\tt.Skip(\"PTY tests not supported on Windows\")\n\t}\n\n\t// Get cog binary\n\tcogBinary, err := harness.ResolveCogBinary()\n\trequire.NoError(t, err, \"failed to resolve cog binary\")\n\n\t// Test login to a fake generic registry\n\t// Note: This will fail at the Docker credential save step, but we can verify\n\t// the interactive prompts work correctly up to that point\n\tt.Run(\"prompts for username and password\", func(t *testing.T) {\n\t\tcmd := exec.Command(cogBinary, \"login\", \"--registry\", \"fake-registry.example.com\")\n\t\tcmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\n\t\t// Start with a PTY\n\t\tptmx, err := pty.Start(cmd)\n\t\trequire.NoError(t, err, \"failed to start PTY\")\n\t\tdefer func() {\n\t\t\tptmx.Close()\n\t\t\tcmd.Process.Kill()\n\t\t\tcmd.Wait()\n\t\t}()\n\n\t\t// Set terminal size\n\t\tif err := pty.Setsize(ptmx, &pty.Winsize{Rows: 24, Cols: 80}); err != nil {\n\t\t\tt.Logf(\"failed to set terminal size: %v\", err)\n\t\t}\n\n\t\t// Use a single mutex-protected buffer for thread safety\n\t\t// This avoids the race condition of multiple goroutines reading from the PTY\n\t\tvar bufMu bytes.Buffer\n\t\tvar mu sync.Mutex\n\n\t\t// Start a single reader goroutine\n\t\tdone := make(chan struct{})\n\t\tgo func() {\n\t\t\ttmp := make([]byte, 1024)\n\t\t\tfor {\n\t\t\t\tselect {\n\t\t\t\tcase <-done:\n\t\t\t\t\treturn\n\t\t\t\tdefault:\n\t\t\t\t\tn, err := ptmx.Read(tmp)\n\t\t\t\t\tif n > 0 {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tbufMu.Write(tmp[:n])\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t}\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}()\n\t\tdefer close(done)\n\n\t\t// Helper to get current buffer contents\n\t\tgetOutput := func() string {\n\t\t\tmu.Lock()\n\t\t\tdefer mu.Unlock()\n\t\t\treturn bufMu.String()\n\t\t}\n\n\t\t// Helper to wait for a pattern in output with timeout\n\t\twaitForPattern := func(pattern string, timeout time.Duration) (string, bool) {\n\t\t\tdeadline := time.Now().Add(timeout)\n\t\t\tfor time.Now().Before(deadline) {\n\t\t\t\toutput := getOutput()\n\t\t\t\tif strings.Contains(strings.ToLower(output), strings.ToLower(pattern)) {\n\t\t\t\t\treturn output, true\n\t\t\t\t}\n\t\t\t\ttime.Sleep(100 * time.Millisecond)\n\t\t\t}\n\t\t\treturn getOutput(), false\n\t\t}\n\n\t\t// Wait for and verify username prompt\n\t\toutput, found := waitForPattern(\"username\", 5*time.Second)\n\t\tt.Logf(\"Output after start: %q\", output)\n\n\t\tassert.Contains(t, output, \"fake-registry.example.com\", \"expected output to mention registry host\")\n\t\tassert.True(t, found, \"expected username prompt, got: %q\", output)\n\n\t\t// Send username\n\t\t_, err = ptmx.Write([]byte(\"testuser\\n\"))\n\t\trequire.NoError(t, err, \"failed to write username\")\n\n\t\t// Wait for password prompt\n\t\toutput, found = waitForPattern(\"password\", 3*time.Second)\n\t\tt.Logf(\"Output after username: %q\", output)\n\n\t\tassert.True(t, found, \"expected password prompt, got: %q\", output)\n\n\t\t// Send password (will fail at Docker credential save, but we've verified the flow)\n\t\t_, err = ptmx.Write([]byte(\"testpass\\n\"))\n\t\trequire.NoError(t, err, \"failed to write password\")\n\n\t\t// Read final output briefly (expect failure since we can't actually save credentials)\n\t\ttime.Sleep(2 * time.Second)\n\t\toutput = getOutput()\n\t\tt.Logf(\"Final output: %q\", output)\n\t})\n\n\tt.Run(\"rejects empty username\", func(t *testing.T) {\n\t\tcmd := exec.Command(cogBinary, \"login\", \"--registry\", \"fake-registry.example.com\")\n\t\tcmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\n\t\tptmx, err := pty.Start(cmd)\n\t\trequire.NoError(t, err, \"failed to start PTY\")\n\t\tdefer func() {\n\t\t\tptmx.Close()\n\t\t\tcmd.Process.Kill()\n\t\t\tcmd.Wait()\n\t\t}()\n\n\t\t// Use a mutex-protected buffer for thread safety\n\t\tvar bufMu bytes.Buffer\n\t\tvar mu sync.Mutex\n\n\t\t// Start reader goroutine\n\t\tdone := make(chan struct{})\n\t\tgo func() {\n\t\t\ttmp := make([]byte, 1024)\n\t\t\tfor {\n\t\t\t\tselect {\n\t\t\t\tcase <-done:\n\t\t\t\t\treturn\n\t\t\t\tdefault:\n\t\t\t\t\tn, err := ptmx.Read(tmp)\n\t\t\t\t\tif n > 0 {\n\t\t\t\t\t\tmu.Lock()\n\t\t\t\t\t\tbufMu.Write(tmp[:n])\n\t\t\t\t\t\tmu.Unlock()\n\t\t\t\t\t}\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}()\n\t\tdefer close(done)\n\n\t\t// Helper to check buffer contents\n\t\tgetOutput := func() string {\n\t\t\tmu.Lock()\n\t\t\tdefer mu.Unlock()\n\t\t\treturn bufMu.String()\n\t\t}\n\n\t\t// Wait for username prompt\n\t\tdeadline := time.Now().Add(5 * time.Second)\n\t\tfor time.Now().Before(deadline) {\n\t\t\tif strings.Contains(strings.ToLower(getOutput()), \"username\") {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\ttime.Sleep(100 * time.Millisecond)\n\t\t}\n\n\t\toutput := getOutput()\n\t\trequire.Contains(t, strings.ToLower(output), \"username\", \"did not get username prompt: %q\", output)\n\n\t\t// Send empty username\n\t\t_, err = ptmx.Write([]byte(\"\\n\"))\n\t\trequire.NoError(t, err, \"failed to write empty username\")\n\n\t\t// Wait for error about empty username\n\t\tdeadline = time.Now().Add(5 * time.Second)\n\t\tfor time.Now().Before(deadline) {\n\t\t\toutput = getOutput()\n\t\t\tif strings.Contains(strings.ToLower(output), \"empty\") || strings.Contains(strings.ToLower(output), \"cannot\") {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\ttime.Sleep(100 * time.Millisecond)\n\t\t}\n\n\t\toutput = getOutput()\n\t\tt.Logf(\"Output: %q\", output)\n\n\t\t// Verify we got an error about empty username\n\t\tlowerOutput := strings.ToLower(output)\n\t\tassert.True(t, strings.Contains(lowerOutput, \"empty\") || strings.Contains(lowerOutput, \"cannot\"),\n\t\t\t\"expected error about empty username, got: %q\", output)\n\t})\n}\n\n// TestLoginProviderRouting tests that the --registry flag correctly routes to the appropriate provider.\n// This test verifies the routing behavior by checking error messages and prompts.\nfunc TestLoginProviderRouting(t *testing.T) {\n\tif testing.Short() {\n\t\tt.Skip(\"skipping in short mode\")\n\t}\n\n\t// PTY tests only work reliably on Unix-like systems\n\tif runtime.GOOS == \"windows\" {\n\t\tt.Skip(\"PTY tests not supported on Windows\")\n\t}\n\n\tcogBinary, err := harness.ResolveCogBinary()\n\trequire.NoError(t, err, \"failed to resolve cog binary\")\n\n\ttests := []struct {\n\t\tname string\n\t\tregistry string\n\t\texpectReplicate bool // True if we expect Replicate provider behavior\n\t}{\n\t\t{\n\t\t\tname: \"default registry uses Replicate\",\n\t\t\tregistry: \"\",\n\t\t\texpectReplicate: true,\n\t\t},\n\t\t{\n\t\t\tname: \"r8.im uses Replicate\",\n\t\t\tregistry: \"r8.im\",\n\t\t\texpectReplicate: true,\n\t\t},\n\t\t{\n\t\t\tname: \"custom registry uses generic\",\n\t\t\tregistry: \"ghcr.io\",\n\t\t\texpectReplicate: false,\n\t\t},\n\t\t{\n\t\t\tname: \"dockerhub uses generic\",\n\t\t\tregistry: \"docker.io\",\n\t\t\texpectReplicate: false,\n\t\t},\n\t\t{\n\t\t\tname: \"localhost uses generic\",\n\t\t\tregistry: \"localhost:5000\",\n\t\t\texpectReplicate: false,\n\t\t},\n\t}\n\n\tfor _, tc := range tests {\n\t\tt.Run(tc.name, func(t *testing.T) {\n\t\t\targs := []string{\"login\"}\n\t\t\tif tc.registry != \"\" {\n\t\t\t\targs = append(args, \"--registry\", tc.registry)\n\t\t\t}\n\n\t\t\tcmd := exec.Command(cogBinary, args...)\n\t\t\tcmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\n\t\t\t// Start with a PTY to handle interactive prompts\n\t\t\tptmx, err := pty.Start(cmd)\n\t\t\trequire.NoError(t, err, \"failed to start PTY\")\n\t\t\tdefer func() {\n\t\t\t\tptmx.Close()\n\t\t\t\tcmd.Process.Kill()\n\t\t\t\tcmd.Wait()\n\t\t\t}()\n\n\t\t\t// Read initial output\n\t\t\tvar buf bytes.Buffer\n\t\t\tdeadline := time.Now().Add(5 * time.Second)\n\t\t\ttmp := make([]byte, 1024)\n\t\t\tfor time.Now().Before(deadline) {\n\t\t\t\tptmx.SetReadDeadline(time.Now().Add(100 * time.Millisecond))\n\t\t\t\tn, _ := ptmx.Read(tmp)\n\t\t\t\tif n > 0 {\n\t\t\t\t\tbuf.Write(tmp[:n])\n\t\t\t\t\t// Check if we have enough output to determine provider\n\t\t\t\t\toutput := buf.String()\n\t\t\t\t\tif tc.expectReplicate {\n\t\t\t\t\t\t// Replicate provider shows \"Hit enter to get started\" message\n\t\t\t\t\t\tif strings.Contains(output, \"Hit enter\") || strings.Contains(output, \"browser\") {\n\t\t\t\t\t\t\tt.Logf(\"Confirmed Replicate provider: %q\", output)\n\t\t\t\t\t\t\treturn\n\t\t\t\t\t\t}\n\t\t\t\t\t} else {\n\t\t\t\t\t\t// Generic provider shows \"Username:\" prompt directly\n\t\t\t\t\t\tif strings.Contains(output, \"Username\") {\n\t\t\t\t\t\t\tt.Logf(\"Confirmed Generic provider: %q\", output)\n\t\t\t\t\t\t\treturn\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\n\t\t\toutput := buf.String()\n\t\t\tif tc.expectReplicate {\n\t\t\t\tassert.NotContains(t, output, \"Username:\", \"expected Replicate provider, but got Generic provider with Username prompt\")\n\t\t\t} else {\n\t\t\t\tassert.True(t, strings.Contains(output, \"Username\") || strings.Contains(strings.ToLower(output), \"logging in\"),\n\t\t\t\t\t\"expected Generic provider prompts, got: %q\", output)\n\t\t\t}\n\t\t})\n\t}\n}\n\n// TestLoginEnvironmentVariable tests that COG_REGISTRY_HOST environment variable works.\nfunc TestLoginEnvironmentVariable(t *testing.T) {\n\tif testing.Short() {\n\t\tt.Skip(\"skipping in short mode\")\n\t}\n\n\t// PTY tests only work reliably on Unix-like systems\n\tif runtime.GOOS == \"windows\" {\n\t\tt.Skip(\"PTY tests not supported on Windows\")\n\t}\n\n\tcogBinary, err := harness.ResolveCogBinary()\n\trequire.NoError(t, err, \"failed to resolve cog binary\")\n\n\tt.Run(\"COG_REGISTRY_HOST sets default registry\", func(t *testing.T) {\n\t\tcmd := exec.Command(cogBinary, \"login\")\n\t\tcmd.Env = append(os.Environ(),\n\t\t\t\"COG_NO_UPDATE_CHECK=1\",\n\t\t\t\"COG_REGISTRY_HOST=custom-registry.example.com\",\n\t\t)\n\n\t\t// Start with a PTY\n\t\tptmx, err := pty.Start(cmd)\n\t\trequire.NoError(t, err, \"failed to start PTY\")\n\t\tdefer func() {\n\t\t\tptmx.Close()\n\t\t\tcmd.Process.Kill()\n\t\t\tcmd.Wait()\n\t\t}()\n\n\t\t// Read output\n\t\tvar buf bytes.Buffer\n\t\tdeadline := time.Now().Add(5 * time.Second)\n\t\ttmp := make([]byte, 1024)\n\t\tfor time.Now().Before(deadline) {\n\t\t\tptmx.SetReadDeadline(time.Now().Add(100 * time.Millisecond))\n\t\t\tn, _ := ptmx.Read(tmp)\n\t\t\tif n > 0 {\n\t\t\t\tbuf.Write(tmp[:n])\n\t\t\t\t// Stop early if we see the expected registry\n\t\t\t\tif strings.Contains(buf.String(), \"custom-registry.example.com\") {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\toutput := buf.String()\n\t\tt.Logf(\"Output: %s\", output)\n\n\t\t// Verify the custom registry is mentioned in output\n\t\tassert.Contains(t, output, \"custom-registry.example.com\", \"expected custom registry in output\")\n\n\t\t// Since custom-registry.example.com is not r8.im, it should use generic provider\n\t\tif !strings.Contains(output, \"Username\") {\n\t\t\tt.Logf(\"Note: Generic provider should prompt for Username\")\n\t\t}\n\t})\n\n\tt.Run(\"--registry flag overrides COG_REGISTRY_HOST\", func(t *testing.T) {\n\t\tcmd := exec.Command(cogBinary, \"login\", \"--registry\", \"override-registry.example.com\")\n\t\tcmd.Env = append(os.Environ(),\n\t\t\t\"COG_NO_UPDATE_CHECK=1\",\n\t\t\t\"COG_REGISTRY_HOST=ignored-registry.example.com\",\n\t\t)\n\n\t\t// Start with a PTY\n\t\tptmx, err := pty.Start(cmd)\n\t\trequire.NoError(t, err, \"failed to start PTY\")\n\t\tdefer func() {\n\t\t\tptmx.Close()\n\t\t\tcmd.Process.Kill()\n\t\t\tcmd.Wait()\n\t\t}()\n\n\t\t// Read output\n\t\tvar buf bytes.Buffer\n\t\tdeadline := time.Now().Add(5 * time.Second)\n\t\ttmp := make([]byte, 1024)\n\t\tfor time.Now().Before(deadline) {\n\t\t\tptmx.SetReadDeadline(time.Now().Add(100 * time.Millisecond))\n\t\t\tn, _ := ptmx.Read(tmp)\n\t\t\tif n > 0 {\n\t\t\t\tbuf.Write(tmp[:n])\n\t\t\t\t// Stop early if we see the expected registry\n\t\t\t\tif strings.Contains(buf.String(), \"override-registry.example.com\") {\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\toutput := buf.String()\n\t\tt.Logf(\"Output: %s\", output)\n\n\t\t// Verify the override registry is used, not the env var one\n\t\tassert.Contains(t, output, \"override-registry.example.com\", \"expected override registry in output\")\n\t\tassert.NotContains(t, output, \"ignored-registry.example.com\", \"env var registry should have been overridden, but it appeared in output\")\n\t})\n}\n\n// TestLoginHelp tests that the login command shows appropriate help text.\nfunc TestLoginHelp(t *testing.T) {\n\tcogBinary, err := harness.ResolveCogBinary()\n\trequire.NoError(t, err, \"failed to resolve cog binary\")\n\n\tcmd := exec.Command(cogBinary, \"login\", \"--help\")\n\tcmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\n\toutput, err := cmd.CombinedOutput()\n\trequire.NoError(t, err, \"help command failed\")\n\n\thelpText := string(output)\n\tt.Logf(\"Help text:\\n%s\", helpText)\n\n\t// Verify help contains expected information\n\texpectedStrings := []string{\n\t\t\"login\",\n\t\t\"registry\",\n\t\t\"--token-stdin\",\n\t\t\"container registry\", // Updated description mentions \"container registry\"\n\t}\n\n\tfor _, expected := range expectedStrings {\n\t\tassert.True(t, strings.Contains(strings.ToLower(helpText), strings.ToLower(expected)),\n\t\t\t\"expected help to contain %q\", expected)\n\t}\n\n\t// Verify help mentions both Replicate and generic registry support\n\tassert.Contains(t, helpText, \"Replicate\", \"expected help to mention Replicate\")\n\tassert.True(t, strings.Contains(helpText, \"other registries\") && strings.Contains(helpText, \"username and password\"),\n\t\t\"expected help to mention generic registry login with username/password\")\n}\n\n// TestLoginSuggestFor tests that similar commands are suggested.\nfunc TestLoginSuggestFor(t *testing.T) {\n\tcogBinary, err := harness.ResolveCogBinary()\n\trequire.NoError(t, err, \"failed to resolve cog binary\")\n\n\t// Test that \"cog auth\" suggests \"cog login\"\n\tcmd := exec.Command(cogBinary, \"auth\")\n\tcmd.Env = append(os.Environ(), \"COG_NO_UPDATE_CHECK=1\")\n\n\toutput, err := cmd.CombinedOutput()\n\t// We expect an error since \"auth\" is not a valid command\n\tif err == nil {\n\t\tt.Logf(\"Unexpected success, output: %s\", output)\n\t}\n\n\toutputStr := string(output)\n\tt.Logf(\"Output for 'cog auth': %s\", outputStr)\n\n\t// Check if login is suggested\n\tif strings.Contains(outputStr, \"login\") {\n\t\tt.Logf(\"'login' suggested for 'auth' command (good)\")\n\t}\n}\n" }, { "path": "integration-tests/suite_test.go", "content": "//go:build integration\n\npackage integration_test\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\t\"os/signal\"\n\t\"path/filepath\"\n\t\"runtime\"\n\t\"sort\"\n\t\"strings\"\n\t\"syscall\"\n\t\"testing\"\n\n\t\"github.com/rogpeppe/go-internal/testscript\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/integration-tests/harness\"\n)\n\n// TestMain sets up signal handling to force exit on cancellation.\n// Without this, go test ignores SIGTERM and keeps running when CI cancels.\nfunc TestMain(m *testing.M) {\n\tsigCh := make(chan os.Signal, 1)\n\tsignal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)\n\n\tgo func() {\n\t\tsig := <-sigCh\n\t\tfmt.Fprintf(os.Stderr, \"\\nReceived %v, forcing exit...\\n\", sig)\n\t\tos.Exit(1)\n\t}()\n\n\tos.Exit(m.Run())\n}\n\nfunc TestIntegration(t *testing.T) {\n\tdir := \"tests\"\n\n\th, err := harness.New()\n\trequire.NoError(t, err, \"failed to create harness\")\n\n\tfiles, err := filepath.Glob(filepath.Join(dir, \"*.txtar\"))\n\trequire.NoError(t, err)\n\tsort.Strings(files)\n\tfor _, f := range files {\n\t\tname := strings.TrimSuffix(filepath.Base(f), filepath.Ext(f))\n\t\tt.Run(name, func(t *testing.T) {\n\t\t\tif !strings.HasSuffix(name, \"_serial\") {\n\t\t\t\tt.Parallel()\n\t\t\t}\n\t\t\ttestscript.Run(t, testscript.Params{\n\t\t\t\tFiles: []string{f},\n\t\t\t\tSetup: h.Setup,\n\t\t\t\tCmds: h.Commands(),\n\t\t\t\tCondition: condition,\n\t\t\t})\n\t\t})\n\t}\n\n}\n\n// condition provides custom conditions for testscript.\n// Supported conditions:\n// - linux/linux_amd64/amd64: platform guards for specialized tests.\n//\n// Note: testscript has built-in support for [short] which checks testing.Short().\nfunc condition(cond string) (bool, error) {\n\tnegated := false\n\tfor strings.HasPrefix(cond, \"!\") {\n\t\tnegated = !negated\n\t\tcond = cond[1:]\n\t}\n\n\tvar value bool\n\tswitch cond {\n\tcase \"linux\":\n\t\tvalue = runtime.GOOS == \"linux\"\n\tcase \"amd64\":\n\t\tvalue = runtime.GOARCH == \"amd64\"\n\tcase \"linux_amd64\":\n\t\tvalue = runtime.GOOS == \"linux\" && runtime.GOARCH == \"amd64\"\n\tdefault:\n\t\treturn false, fmt.Errorf(\"unknown condition: %s\", cond)\n\t}\n\n\tif negated {\n\t\tvalue = !value\n\t}\n\treturn value, nil\n}\n" }, { "path": "integration-tests/tests/apt_packages.txtar", "content": "# Skip for cog-dataclass and coglet (Rust) which require Python 3.10+\n# Test that system packages are installed correctly via system_packages in cog.yaml\n\n# Build the image (the run command verifies git is installed)\ncog build -t $TEST_IMAGE\n\n# Verify the predictor works\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n gpu: true\n python_version: \"3.10\"\n system_packages:\n - \"git\"\n run:\n - command: git --version\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n def predict(self, s: str) -> str:\n return \"hello \" + s\n" }, { "path": "integration-tests/tests/async_generator_precollect.txtar", "content": "# Test that async generator output is pre-collected before response.\n#\n# Coglet collects all async generator yields into a list before sending\n# the response. This test verifies all items arrive in the response and\n# that predict_time reflects the full generation duration.\n\ncog serve\n\ncurl POST /predictions '{\"input\":{}}'\nstdout '\"status\":\"succeeded\"'\n# All 5 items should be present in the output\nstdout '\"output\":\\[\"chunk-0\",\"chunk-1\",\"chunk-2\",\"chunk-3\",\"chunk-4\"\\]'\n# predict_time should reflect full generation time (at least ~0.5s)\nstdout '\"predict_time\":(0\\.[5-9][0-9]*|[1-9][0-9]*(\\.[0-9]+)?)'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\nconcurrency:\n max: 1\n\n-- predict.py --\nimport asyncio\nfrom typing import AsyncIterator\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n async def predict(self) -> AsyncIterator[str]:\n for i in range(5):\n await asyncio.sleep(0.1)\n yield f\"chunk-{i}\"\n" }, { "path": "integration-tests/tests/async_predictor.txtar", "content": "# Build the image\ncog build -t $TEST_IMAGE\n\n# Async prediction works\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n async def predict(self, s: str) -> str:\n return \"hello \" + s\n" }, { "path": "integration-tests/tests/async_sleep.txtar", "content": "# Test async predictor with sleep\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Async prediction with sleep works\ncog predict $TEST_IMAGE -i s=sleepyhead -i sleep=0.1\nstdout 'wake up sleepyhead'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.11\"\npredict: \"predict.py:Predictor\"\nconcurrency:\n max: 5\n\n-- predict.py --\nimport asyncio\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n async def predict(self, s: str, sleep: float) -> str:\n await asyncio.sleep(sleep)\n return f\"wake up {s}\"\n" }, { "path": "integration-tests/tests/bad_dockerignore.txtar", "content": "# Skip for cog-dataclass and coglet (Rust) which require Python 3.10+\n# Test that build fails with proper error when .cog is in .dockerignore\n! cog build -t $TEST_IMAGE\nstderr 'The .cog tmp path cannot be ignored by docker in .dockerignore'\n\n-- cog.yaml --\nbuild:\n gpu: true\n python_version: \"3.10\"\npredict: \"predict.py:Predictor\"\n\n-- .dockerignore --\n.cog\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n def predict(self, s: str) -> str:\n return \"hello \" + s\n" }, { "path": "integration-tests/tests/bool_input_output.txtar", "content": "# Test bool as a direct predict input and output type\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Bool input and output works via JSON (true -> false)\ncog predict $TEST_IMAGE --json '{\"flag\": true}'\nstdout '\"output\": false'\n\n# Bool input and output works via JSON (false -> true)\ncog predict $TEST_IMAGE --json '{\"flag\": false}'\nstdout '\"output\": true'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n def predict(self, flag: bool) -> bool:\n return not flag\n" }, { "path": "integration-tests/tests/build_base_image_sha.txtar", "content": "# Test that base image SHA is recorded in labels with --use-cog-base-image\n# Source: test_build.py::test_build_base_image_sha\n#\n# Uses a local test registry to avoid depending on the live r8.im registry.\n\n# Start local registry and seed a cog-base image from Docker Hub's python:3.12-slim\nregistry-start\nregistry-seed docker.io/library/python:3.12-slim cog-base:python3.12\n\n# Build with --use-cog-base-image flag against the local registry\nenv COG_REGISTRY_HOST=$TEST_REGISTRY\ncog build -t $TEST_IMAGE --use-cog-base-image\n\n# Verify the base image layer label matches one of the actual layers\nexec python3 -c 'import json,os,subprocess; image=os.environ[\"TEST_IMAGE\"]; base_layer=subprocess.check_output([\"docker\",\"inspect\",image,\"--format={{index .Config.Labels \\\"run.cog.cog-base-image-last-layer-sha\\\"}}\"], text=True).strip(); assert base_layer.startswith(\"sha256:\"), f\"Base layer label missing sha256 digest: {base_layer}\"; layers=json.loads(subprocess.check_output([\"docker\",\"inspect\",image,\"--format={{json .RootFS.Layers}}\"], text=True)); assert base_layer in layers, f\"Base layer {base_layer} not found in RootFS layers\"; print(base_layer)'\nstdout 'sha256:'\n\n-- cog.yaml --\nbuild:\n python_version: \"3.12\"\npredict: predict.py:Predictor\n\n-- predict.py --\nimport tempfile\n\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n def setup(self):\n self.foo = \"foo\"\n\n def predict(self, text: str, path: Path) -> Path:\n with open(path) as f:\n output = self.foo + text + f.read()\n tmpdir = Path(tempfile.mkdtemp())\n with open(tmpdir / \"output.txt\", \"w\") as fh:\n fh.write(output)\n return tmpdir / \"output.txt\"\n" }, { "path": "integration-tests/tests/build_cog_init.txtar", "content": "# Test that cog init creates a buildable project\n# Source: test_build.py::test_build_with_cog_init_templates\n\n# Initialize a new cog project\ncog init\n\n# Build the initialized project\ncog build -t $TEST_IMAGE\nstderr 'Image built as'\n\n# Verify the expected files were created\nexists cog.yaml\nexists predict.py\n" }, { "path": "integration-tests/tests/build_cog_version_match.txtar", "content": "# Note: This test can be flaky if the wheel build time and CI run time\n# cross day boundaries (version date mismatch). If that happens, the test\n# will fail with a date component mismatch.\n\n# Test that cog version in base image contains a version number\n# Source: test_build.py::test_cog_install_base_image\n#\n# This test verifies that when building with --use-cog-base-image,\n# the installed Python cog package has a valid version number.\n#\n# Uses a local test registry to avoid depending on the live r8.im registry.\n\n# Start local registry and seed a cog-base image from Docker Hub's python:3.12-slim\nregistry-start\nregistry-seed docker.io/library/python:3.12-slim cog-base:python3.12\n\n# Build using --use-cog-base-image against local registry\nenv COG_REGISTRY_HOST=$TEST_REGISTRY\ncog build -t $TEST_IMAGE --use-cog-base-image=true\n\n# Compare the embedded version label with the Python package version inside the image\nexec python3 check_versions.py\nstdout '[0-9]+\\.[0-9]+\\.[0-9]+'\n\n-- check_versions.py --\nimport os\nimport re\nimport subprocess\n\n# SemVer pattern from https://semver.org/\nSEMVER_PATTERN = r\"^(?P0|[1-9]\\d*)\\.(?P0|[1-9]\\d*)\\.(?P0|[1-9]\\d*)(?:-(?P(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\\.(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\\+(?P[0-9a-zA-Z-]+(?:\\.[0-9a-zA-Z-]+)*))?$\"\n\n# PEP 440 pattern from packaging.version.VERSION_PATTERN\nPEP440_PATTERN = r\"\"\"\n v?\n (?:\n (?:(?P[0-9]+)!)? # epoch\n (?P[0-9]+(?:\\.[0-9]+)*) # release segment\n (?P
                                           # pre-release\n            [-_\\.]?\n            (?P(a|alpha|b|beta|c|rc))\n            [-_\\.]?\n            (?P[0-9]+)?\n        )?\n        (?P                                          # post release\n            (?:-(?P[0-9]+))\n            |\n            (?:\n                [-_\\.]?\n                (?Ppost|rev|r)\n                [-_\\.]?\n                (?P[0-9]+)?\n            )\n        )?\n        (?P                                           # dev release\n            [-_\\.]?\n            (?Pdev)\n            [-_\\.]?\n            (?P[0-9]+)?\n        )?\n    )\n    (?:\\+(?P[a-z0-9]+(?:[-_\\.][a-z0-9]+)*))?       # local version\n\"\"\"\n\ndef assert_versions_match(semver_version, pep440_version):\n    semver_re = re.compile(SEMVER_PATTERN)\n    pep440_re = re.compile(PEP440_PATTERN, re.VERBOSE | re.IGNORECASE)\n\n    semver_match = semver_re.match(semver_version)\n    pep440_match = pep440_re.match(pep440_version)\n\n    assert semver_match, f\"Invalid semver version: {semver_version}\"\n    assert pep440_match, f\"Invalid PEP 440 version: {pep440_version}\"\n\n    semver_groups = semver_match.groupdict()\n    pep440_groups = pep440_match.groupdict()\n\n    semver_release = f\"{semver_groups['major']}.{semver_groups['minor']}.{semver_groups['patch']}\"\n\n    # Check base release version\n    assert semver_release == pep440_groups[\"release\"], (\n        f\"Release versions do not match: {semver_release} != {pep440_groups['release']}\"\n    )\n\n    # Check prerelease status\n    semver_pre = semver_groups[\"prerelease\"]\n    pep440_pre = pep440_groups[\"pre\"] or pep440_groups[\"dev\"]\n\n    assert bool(semver_pre) == bool(pep440_pre), \"Pre-release status does not match\"\n\n    if semver_pre:\n        if semver_pre.startswith(\"alpha\"):\n            assert pep440_groups[\"pre_l\"] == \"a\", \"Alpha pre-release status does not match\"\n            assert not pep440_groups[\"dev\"], \"Semver pre-release cannot also be a PEP440 dev build\"\n\n        if semver_pre.startswith(\"beta\"):\n            assert pep440_groups[\"pre_l\"] == \"b\", \"Beta pre-release status does not match\"\n            assert not pep440_groups[\"dev\"], \"Semver pre-release cannot also be a PEP440 dev build\"\n\n        if semver_pre.startswith(\"rc\"):\n            assert pep440_groups[\"pre_l\"] == \"rc\", \"Release candidate pre-release status does not match\"\n            assert not pep440_groups[\"dev\"], \"Semver pre-release cannot also be a PEP440 dev build\"\n\n        if semver_pre.startswith(\"dev\"):\n            assert pep440_groups[\"dev_l\"] == \"dev\", \"Dev build status does not match\"\n\n    if pep440_groups[\"local\"] is not None and semver_groups[\"buildmetadata\"] is not None:\n        # Both build metadata formats are: g.d\n        # The git short hash length can vary (typically 7-9 chars) depending on\n        # git settings and repo size, so we need to compare flexibly.\n        # Split by '.' and compare the git hash and date parts separately.\n        semver_parts = semver_groups[\"buildmetadata\"].split(\".\")\n        pep440_parts = pep440_groups[\"local\"].split(\".\")\n\n        # Compare git commit hash - one should be a prefix of the other\n        # (e.g., \"g5606e933\" and \"g5606e9331\" both refer to the same commit)\n        semver_hash = semver_parts[0] if semver_parts else \"\"\n        pep440_hash = pep440_parts[0] if pep440_parts else \"\"\n        hash_match = semver_hash.startswith(pep440_hash) or pep440_hash.startswith(semver_hash)\n        assert hash_match, (\n            f\"Git commit hash does not match: {semver_hash} vs {pep440_hash}\"\n        )\n\n        # Compare date parts if present (should be exact match)\n        if len(semver_parts) > 1 and len(pep440_parts) > 1:\n            assert semver_parts[1] == pep440_parts[1], (\n                f\"Date component does not match: {semver_parts[1]} != {pep440_parts[1]}\"\n            )\n\nimage = os.environ[\"TEST_IMAGE\"]\nlabel = subprocess.check_output(\n    [\"docker\", \"inspect\", image, \"--format={{index .Config.Labels \\\"run.cog.version\\\"}}\"],\n    text=True\n).strip()\npackage = subprocess.check_output(\n    [\"docker\", \"run\", \"--rm\", \"-t\", image, \"python\", \"-c\", \"import cog; print(cog.__version__)\"],\n    text=True\n).strip()\n\n# Validate package version format\npattern = re.compile(r\"^\\d+\\.\\d+\\.\\d+\")\nassert pattern.search(package), f\"Invalid package version: {package}\"\n\n# If label is \"dev\", skip version matching (special dev builds)\nif label != \"dev\":\n    assert pattern.search(label), f\"Invalid label version: {label}\"\n    assert_versions_match(label, package)\n\nprint(package)\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/build_gpu_labels.txtar",
    "content": "# Test GPU build labels\n# Source: test_build.py::test_build_gpu_model_on_cpu\n# Requires git init/tag for version labels\n\n[short] skip 'requires long GPU build time'\n\n# Setup git repo for version labels\nexec git init\nexec git config user.email noreply@replicate.com\nexec git config user.name 'Replicate Test Bot'\nexec git config commit.gpgsign false\nexec git commit --allow-empty -m initial\nexec git tag 0.0.1\n\n# Build the GPU image\ncog build -t $TEST_IMAGE\n\n# Check core labels exist\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"run.cog.version\"}}'\nstdout '.+'\n\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"run.cog.config\"}}'\nstdout '\"gpu\":true'\nstdout '\"cuda\":'\nstdout '\"cudnn\":'\nstdout '\"python_version\":\"3.12\"'\n\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"run.cog.openapi_schema\"}}'\nstdout 'openapi'\n\n# Check OCI labels for version info\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"org.opencontainers.image.version\"}}'\nstdout '.+'\n\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"org.opencontainers.image.revision\"}}'\nstdout '.+'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  gpu: true\npredict: predict.py:Predictor\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return text\n"
  },
  {
    "path": "integration-tests/tests/build_image_option.txtar",
    "content": "# Test that image: option in cog.yaml names the built image\n# Source: test_build.py::test_build_names_uses_image_option_in_cog_yaml\n\n# Build without explicit -t flag, should use image from cog.yaml\ncog build\n\n# Verify image exists with the name from cog.yaml\nexec docker images --format '{{.Repository}}'\nstdout 'cog-test-image-option'\n\n# Cleanup the custom image\nexec docker rmi cog-test-image-option\n\n-- cog.yaml --\nimage: cog-test-image-option\nbuild:\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return text\n"
  },
  {
    "path": "integration-tests/tests/build_openapi_schema.txtar",
    "content": "# Test that OpenAPI schema is embedded in image labels\n# Source: test_build.py::test_build_with_model\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Check the openapi_schema label exists and contains expected schema structure\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"run.cog.openapi_schema\"}}'\nstdout '\"title\":\"Input\"'\nstdout '\"required\":\\[\"text\",\"path\"\\]'\nstdout '\"text\":'\nstdout '\"path\":'\nstdout '\"type\":\"string\"'\nstdout '\"format\":\"uri\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n\n-- predict.py --\nimport tempfile\n\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def setup(self):\n        self.foo = \"foo\"\n\n    def predict(self, text: str, path: Path) -> Path:\n        with open(path) as f:\n            output = self.foo + text + f.read()\n        tmpdir = Path(tempfile.mkdtemp())\n        with open(tmpdir / \"output.txt\", \"w\") as fh:\n            fh.write(output)\n        return tmpdir / \"output.txt\"\n"
  },
  {
    "path": "integration-tests/tests/build_openapi_schema_complex.txtar",
    "content": "# Test that the OpenAPI schema for complex input/output types is correct.\n#\n# Verifies schema generation for:\n# - Multiple input types with constraints (ge, le, choices)\n# - Optional fields with defaults\n# - Secret type\n# - Structured BaseModel output with nested types\n\ncog build -t $TEST_IMAGE\n\n# Extract the schema from the image label\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"run.cog.openapi_schema\"}}'\n\n# Input schema checks\nstdout '\"title\":\"Input\"'\nstdout '\"temperature\"'\nstdout '\"prompt\"'\nstdout '\"style\"'\nstdout '\"api_key\"'\nstdout '\"image\"'\n\n# Constraints on temperature\nstdout '\"minimum\":0'\nstdout '\"maximum\":2'\n\n# Choices for style (enum)\nstdout '\"enum\":\\[\"fast\",\"balanced\",\"quality\"\\]'\n\n# Secret type renders as string with format\nstdout '\"x-cog-secret\":true'\n\n# Optional field has default\nstdout '\"default\":\"hello\"'\n\n# Output schema has structured type\nstdout '\"Output\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom typing import Optional\n\nfrom cog import BaseModel, BasePredictor, Input, Path, Secret\n\n\nclass Output(BaseModel):\n    text: str\n    score: float\n\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        prompt: str = Input(description=\"The prompt\", default=\"hello\"),\n        temperature: float = Input(description=\"Sampling temp\", ge=0, le=2, default=0.7),\n        style: str = Input(description=\"Style\", choices=[\"fast\", \"balanced\", \"quality\"], default=\"balanced\"),\n        api_key: Secret = Input(description=\"API key\"),\n        image: Optional[Path] = Input(description=\"Optional image\", default=None),\n    ) -> Output:\n        return Output(text=f\"generated from {prompt}\", score=temperature)\n"
  },
  {
    "path": "integration-tests/tests/build_pip_freeze.txtar",
    "content": "\n# Test that pip freeze label is embedded in image\n# Source: test_build.py::test_pip_freeze\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Check the pip_freeze label exists and contains expected packages\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"run.cog.pip_freeze\"}}'\nstdout 'structlog'\nstdout 'coglet'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n\n-- predict.py --\nimport tempfile\n\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def setup(self):\n        self.foo = \"foo\"\n\n    def predict(self, text: str, path: Path) -> Path:\n        with open(path) as f:\n            output = self.foo + text + f.read()\n        tmpdir = Path(tempfile.mkdtemp())\n        with open(tmpdir / \"output.txt\", \"w\") as fh:\n            fh.write(output)\n        return tmpdir / \"output.txt\"\n"
  },
  {
    "path": "integration-tests/tests/build_python313_base_image.txtar",
    "content": "# Test Python 3.13 works with --use-cog-base-image\n# Source: test_build.py::test_python_313_base_images\n#\n# Uses a local test registry to avoid depending on the live r8.im registry.\n\n# Start local registry and seed a cog-base image from Docker Hub's python:3.13-slim\nregistry-start\nregistry-seed docker.io/library/python:3.13-slim cog-base:python3.13\n\n# Build using Python 3.13 with --use-cog-base-image against local registry\nenv COG_REGISTRY_HOST=$TEST_REGISTRY\ncog build -t $TEST_IMAGE --use-cog-base-image\n\n# Verify build succeeded by running a prediction\ncog predict $TEST_IMAGE -i num=7\nstdout '14'\n\n-- cog.yaml --\nbuild:\n  gpu: false\n  python_version: \"3.13\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, num: int) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/build_torch_version_required.txtar",
    "content": "# Test that build fails with proper error when torch is specified without version and GPU is enabled\n# This validates the error message from cudasFromTorch when torch version is empty\n\n! cog build -t $TEST_IMAGE\nstderr 'torch version must be specified when using CUDA'\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  python_version: \"3.12\"\n  python_requirements: requirements.txt\npredict: predict.py:Predictor\n\n-- requirements.txt --\ntorch\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return text\n"
  },
  {
    "path": "integration-tests/tests/ca_cert.txtar",
    "content": "# Test CA certificate injection via COG_CA_CERT\n#\n# This test verifies that custom CA certificates can be injected into\n# cog-built images, allowing HTTPS connections to servers using those CAs.\n\n# Build the HTTPS test server helper\nexec go build -C $REPO_ROOT/test-helpers/https-server -o $WORK/https-server .\n\n# Start the HTTPS test server in background using the embedded certs\nexec $WORK/https-server --cert=$WORK/certs/server.crt --key=$WORK/certs/server.key --addr=:8443 &\n\n# Wait for the server to be ready\nexec sh -c 'for i in 1 2 3 4 5 6 7 8 9 10; do curl -ksf https://localhost:8443/ && exit 0; sleep 1; done; exit 1'\n\n# Build a minimal cog image\ncog build -t $TEST_IMAGE\n\n# ============================================\n# Test 1: Without CA cert, HTTPS should FAIL\n# ============================================\n# The container tries to curl the HTTPS server, which uses a self-signed cert.\n# Without the CA cert installed, this should fail with a certificate error.\n! docker-run $TEST_IMAGE curl --fail --max-time 5 https://host.docker.internal:8443/\n\n# ============================================\n# Test 2: With COG_CA_CERT, HTTPS should WORK\n# ============================================\n# Set the CA cert environment variable and rebuild\nenv COG_CA_CERT=$WORK/certs/ca.crt\n\ncog build -t $TEST_IMAGE-with-ca\n\n# Now the HTTPS request should succeed because the CA cert is installed\ndocker-run $TEST_IMAGE-with-ca curl --fail --max-time 5 https://host.docker.internal:8443/\nstdout 'OK'\n\n# Verify the environment variables are set correctly in the image\ndocker-run $TEST_IMAGE-with-ca printenv SSL_CERT_FILE\nstdout '/etc/ssl/certs/ca-certificates.crt'\n\ndocker-run $TEST_IMAGE-with-ca printenv REQUESTS_CA_BUNDLE\nstdout '/etc/ssl/certs/ca-certificates.crt'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return text\n\n-- certs/ca.crt --\n-----BEGIN CERTIFICATE-----\nMIIDDTCCAfWgAwIBAgIUayiqAjIWvavCGAlFwU4CtOlDyjgwDQYJKoZIhvcNAQEL\nBQAwFjEUMBIGA1UEAwwLQ29nIFRlc3QgQ0EwHhcNMjYwMTE0MjAxMjE0WhcNMzYw\nMTEyMjAxMjE0WjAWMRQwEgYDVQQDDAtDb2cgVGVzdCBDQTCCASIwDQYJKoZIhvcN\nAQEBBQADggEPADCCAQoCggEBAOanC1VJPLjZrfW+hLL6FmDsnNokMU1rI8KCWjrE\nG2BLWOODSOECd1TWuJ84leiwOKuqj9FXlHzf0wr/D6MnQq39R4yDHKdbHYVuwRBu\nuP3M3M3LWkqs7FDcXRz2htEoSoFAfoNo85Paj8rpFYwzLsuS/DtxX2yM5ja1UAZk\nSNjrWF7DY97cT9njLF2QYFLj1unWAlVKoR90cYZZ72S4QIWsTQXBNN3GR/GC80AJ\nvaxC83n4fCN94vJgO4reMAlojFNlXSgqQkEf8z+SMcuzHNcV/FkNOArTXHaYLeaB\nyKChtDIlHV9W0+Ifsr+qYkWCN2Aznw5Yz5bhXrFLd+BcHUcCAwEAAaNTMFEwHQYD\nVR0OBBYEFDAHM6Rgg47dX2D0h7gBw365IZ6kMB8GA1UdIwQYMBaAFDAHM6Rgg47d\nX2D0h7gBw365IZ6kMA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZIhvcNAQELBQADggEB\nAMd8xNa6BO7nZBfpwrPaijLMawIkv37ngM2dkNFKPMV/Vl9urAAnCDtjFq/pCXnC\nlXvja0vy6nFmRITmetabLdhBHeDe8lcV1OHgyksy6AnRT6zLu/Jw3cx5/U3zLwM7\nzLZkSylhYaNVpqaTRyzdFbP5V8h/QjpL+ffYTQgNMtRT0PphV4AvAqpfJJ18Jtcc\nK4jIXYMUKU4ZkAT9JUSTXa2aefudzjMMr9GwvZGn/6ZpU3Y8H/DmmizoHNQRuC97\nuGnxwufInGQ9W20UnUam9May0tsea654Ebtjw7QDzvMTFIsMFVOjBVOXMuB8PAfL\nATgc+2ToYm+V3Z1f46mm/4U=\n-----END CERTIFICATE-----\n\n-- certs/server.crt --\n-----BEGIN CERTIFICATE-----\nMIIDRzCCAi+gAwIBAgIUa+aDQvQpf7Sq+7foxZ7MNjJKcnIwDQYJKoZIhvcNAQEL\nBQAwFjEUMBIGA1UEAwwLQ29nIFRlc3QgQ0EwHhcNMjYwMTE0MjAxMjE5WhcNMzYw\nMTEyMjAxMjE5WjAfMR0wGwYDVQQDDBRob3N0LmRvY2tlci5pbnRlcm5hbDCCASIw\nDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAJuuQyShcfUyXWhVzYGZ2UTo2Do1\nTfHFrcHLsXB2UvuMnH67x3h1ylLM2gmMnMipVFVTL6OAhZQi9BHiRRoCepMyFJKU\nRKxkXpIF6knfsmB3FsAcTKale3PUYrzTj63BCGZmnS768Drv9e48A2ZvsPTBkyuN\nkjTgC8tIDt5b8xKMRDFG5ZhBxC2+nNBoLBB4ujBF32dWcxGWkFUWsIyN3oH3MsQp\nydO2FnOlc+bQI/hlanl+4CnL/TczOe5O906TL/oC8Wq5nYooGgfG2ZDfSpw7/Ver\nlE0nc4Sy8+d8XL2fNWgM9ISL2sCJ1DT8dZVufmda5MIo0UMCRzCmy0Q6aqcCAwEA\nAaOBgzCBgDA+BgNVHREENzA1ghRob3N0LmRvY2tlci5pbnRlcm5hbIIJbG9jYWxo\nb3N0ggxodHRwcy1zZXJ2ZXKHBH8AAAEwHQYDVR0OBBYEFK/uLPc1TpgjXzMTPCMJ\nQjfx52+6MB8GA1UdIwQYMBaAFDAHM6Rgg47dX2D0h7gBw365IZ6kMA0GCSqGSIb3\nDQEBCwUAA4IBAQBYvP1G1bJM6tOVixEqpxWkGd6Ghr3J/R4hzJDKtMfO8O4FdlZJ\nFG9OaAXJqWmlXszXF2cD2IcdOeayR1oTDTyaId464Y5Fi5WDhaIOAuwlepvAwQod\np1xXbbI6k2n60sSvaCOT9KWwM/zMda94awc2oBYWAaqAJWtRqR+sHnpIe5PmVQN/\nkMfugBZCt21v6tvOEyc94xU6XgqYbbAyZlrFL33KbDpOhnEQzq4F0fKcWit3STR1\noyzlcWaFNysaNOBjxBflBMIKMnpg3DvcK8SUHqCDYjGZzodtC/7f7/I/1zkBvouN\nAdvFtZREbY0xv0tnHl0Afx46Z+hR9rPtjzWh\n-----END CERTIFICATE-----\n\n-- certs/server.key --\n-----BEGIN PRIVATE KEY-----\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQCbrkMkoXH1Ml1o\nVc2BmdlE6Ng6NU3xxa3By7FwdlL7jJx+u8d4dcpSzNoJjJzIqVRVUy+jgIWUIvQR\n4kUaAnqTMhSSlESsZF6SBepJ37JgdxbAHEympXtz1GK804+twQhmZp0u+vA67/Xu\nPANmb7D0wZMrjZI04AvLSA7eW/MSjEQxRuWYQcQtvpzQaCwQeLowRd9nVnMRlpBV\nFrCMjd6B9zLEKcnTthZzpXPm0CP4ZWp5fuApy/03MznuTvdOky/6AvFquZ2KKBoH\nxtmQ30qcO/1Xq5RNJ3OEsvPnfFy9nzVoDPSEi9rAidQ0/HWVbn5nWuTCKNFDAkcw\npstEOmqnAgMBAAECggEAC+dIJP3fK8NdFwQwgW9VCIrRNaoruofF4GKFv7acY7V9\npccP2msPPEODjGVe+4zO8PM6WkMSc6A0j0WAyRtVafnTTt3dXl0SShH/twROrEeO\nysOfLMLMbK/ZmNyISN3QmZvQ+u2e/rKoWD3oeKWjnyNJ8HOTsU1MOY/Z6zCWpl1K\npAiVVvDVvUzrenMLVHlbyHXPYOS+oktctVd57bCNnipG4b/i4pqw08LAkP26jSCm\nyIOg0r+RocN4GhhbmzP4Plmv0JCXhQposIDhC7KfbXEWaa3nFF2nfR5QPVUF5tht\nxLPVrBMac8oT2owQoerhrgCQLi3b4Lmdloz7TwEcQQKBgQDVrwajIW1z5W+HdUm6\nVcJf6+I9v4vL0EXPMSK4+XwLzTkC/RCz0wPfrRGlhdrMZ1wU9zCqFw1LCp+ixsMN\ntC3MWGncFY5TOVf+jqom3PTESC5O9AySRAqI1jE9BwnmDkLpQ5zkYZaF/MAg07px\nCW5r43EKGL5AW4Umv0OEuAEExwKBgQC6grMIz4IuWnZhj03DPRqSx4AymyqrzMjp\nkNzCb1Wz6iKFMeIIPC8Mz0iju9QBVofh8Lxj1Bdqbh3kN4NyfT8nII3i/6Cvb9Ol\nagMgTq1q5BgcSl4jliDbn1gCMq9MCJf8y+xXDQCLIKNaJidXQjQZn1XK/bV1IkPx\nlkuG6znLIQKBgEVcV+Ix4o5hJi+pEbKLTdnG/pwehek1hMN5ZpT2Xp6SEfR3YqmM\nUFCVpAm/hkMdNdWUW1aKvwThwOmcbQoQt2ECPfJziMxY68g0VOTiig0AhQ+Zxk7g\nCS9bn4X4t+zWKj//c3jqeGqrnU3KjFVOw2n/3NxzJaZMTs9B/E+jTqlXAoGBALKc\nQanZVvjfBulM3BJxrMYNqZZNBGM8HNeYM+E7z54ZRW+6opRyVjh1NUIfuNqDLGPS\nMAeF79qrk5KfGxGEIfttcJOHbDE17UBGsrG4xthLkU9eZKK9vb+07ApG0ZsFy896\n1l1TBUc3PVgym5AzxUMYVIetyZ1f8CMmZDPThighAoGBAIaoc0LvNVW9O9wT74dr\ny1ThnmpgeIfg+4yyfY3Eel5wnfdOJ+J9g0vfV+PH/2B45Jdkr82L+XjPipzP3tqP\nfQnfvBCfWXqiCV5yTK4OcVCicKNRj25Oq0vUAgSua8tE6hFaIp5hnmbLYiLIyOAO\nng5irDkEC4tQl7D4WoX1Viye\n-----END PRIVATE KEY-----\n"
  },
  {
    "path": "integration-tests/tests/cancel_async_prediction.txtar",
    "content": "# Test cancellation of a running async prediction.\n#\n# Submits an async prediction (PUT with Prefer: respond-async) that sleeps\n# for 60s via asyncio.sleep, waits for it to start processing, cancels via\n# the /cancel endpoint, and verifies the webhook receives \"canceled\" status.\n\n[short] skip 'requires Docker build'\n\ncog build -t $TEST_IMAGE\n\nwebhook-server-start\ncog serve --upload-url http://unused/\n\n# Submit async prediction that would run for 60s\ncurl -H Prefer:respond-async PUT /predictions/cancel-async-test '{\"id\":\"cancel-async-test\",\"input\":{\"duration\":60},\"webhook\":\"$WEBHOOK_URL\",\"webhook_events_filter\":[\"completed\"]}'\n\n# Give the prediction time to start processing\nexec sleep 2\n\n# Cancel it\ncurl POST /predictions/cancel-async-test/cancel '{}'\n\n# Wait for webhook callback\nwebhook-server-wait 30s\n\n# Prediction should be canceled\nstdout '\"status\":\"canceled\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport asyncio\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def setup(self):\n        pass\n\n    async def predict(self, duration: float = 60.0) -> str:\n        await asyncio.sleep(duration)\n        return \"completed\"\n"
  },
  {
    "path": "integration-tests/tests/cancel_repeated.txtar",
    "content": "# Test that cancelling the same prediction multiple times doesn't panic or break.\n#\n# Uses time.sleep() (C-level nanosleep) — the harder cancellation case since\n# the thread is blocked in native code.  Fires 3 cancel requests in quick\n# succession and verifies the webhook still receives \"canceled\" status.\n\n[short] skip 'requires Docker build'\n\ncog build -t $TEST_IMAGE\n\nwebhook-server-start\ncog serve --upload-url http://unused/\n\n# Submit async prediction that sleeps for 5s (nanosleep blocks in C;\n# cancel fires once sleep returns, so keep it short for CI)\ncurl -H Prefer:respond-async PUT /predictions/cancel-repeat-test '{\"id\":\"cancel-repeat-test\",\"input\":{\"duration\":5},\"webhook\":\"$WEBHOOK_URL\",\"webhook_events_filter\":[\"completed\"]}'\n\n# Give the prediction time to start processing\nexec sleep 2\n\n# Cancel it 3 times in rapid succession — none should 500 or panic\ncurl POST /predictions/cancel-repeat-test/cancel '{}'\ncurl POST /predictions/cancel-repeat-test/cancel '{}'\ncurl POST /predictions/cancel-repeat-test/cancel '{}'\n\n# Wait for webhook callback\nwebhook-server-wait 30s\n\n# Prediction should be canceled\nstdout '\"status\":\"canceled\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport time\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def setup(self):\n        pass\n\n    def predict(self, duration: float = 60.0) -> str:\n        # C-level blocking sleep (nanosleep) — harder to cancel than a busy-loop\n        time.sleep(duration)\n        return \"completed\"\n"
  },
  {
    "path": "integration-tests/tests/cancel_sync_prediction.txtar",
    "content": "# Test cancellation of a running sync prediction.\n#\n# Submits an async prediction (PUT with Prefer: respond-async) that busy-loops\n# for 60s, waits for it to start processing, cancels via the /cancel endpoint,\n# and verifies the webhook receives \"canceled\" status.\n\n[short] skip 'requires Docker build'\n\ncog build -t $TEST_IMAGE\n\nwebhook-server-start\ncog serve --upload-url http://unused/\n\n# Submit async prediction that would run for 60s\ncurl -H Prefer:respond-async PUT /predictions/cancel-sync-test '{\"id\":\"cancel-sync-test\",\"input\":{\"duration\":60},\"webhook\":\"$WEBHOOK_URL\",\"webhook_events_filter\":[\"completed\"]}'\n\n# Give the prediction time to start processing\nexec sleep 2\n\n# Cancel it\ncurl POST /predictions/cancel-sync-test/cancel '{}'\n\n# Wait for webhook callback\nwebhook-server-wait 30s\n\n# Prediction should be canceled\nstdout '\"status\":\"canceled\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport time\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def setup(self):\n        pass\n\n    def predict(self, duration: float = 60.0) -> str:\n        # Busy-loop (hits bytecode boundaries, cancellable via PyThreadState_SetAsyncExc)\n        deadline = time.monotonic() + duration\n        while time.monotonic() < deadline:\n            pass\n        return \"completed\"\n"
  },
  {
    "path": "integration-tests/tests/coglet_iterator_path_output.txtar",
    "content": "# Test iterator prediction with Path outputs (no upload URL — files written to disk)\n\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE\n\n# cog predict writes file outputs to disk, not as base64 to stdout\nstderr 'Written output to: output.0.png'\nstderr 'Written output to: output.1.png'\nstderr 'Written output to: output.2.png'\n! stderr 'failed'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pillow==10.4.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os\nimport tempfile\nfrom typing import Iterator\n\nfrom cog import BasePredictor, Path\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> Iterator[Path]:\n        for color in [\"red\", \"blue\", \"green\"]:\n            d = tempfile.mkdtemp()\n            p = os.path.join(d, f\"{color}.png\")\n            Image.new(\"RGB\", (10, 10), color).save(p)\n            yield Path(p)\n"
  },
  {
    "path": "integration-tests/tests/coglet_iterator_upload_url.txtar",
    "content": "# Test that iterator Path outputs are uploaded per-yield to --upload-url.\n\ncog build -t $TEST_IMAGE\n\n# Start mock upload server on the host, sets $UPLOAD_SERVER_URL\nupload-server-start\n\ncog serve --upload-url $UPLOAD_SERVER_URL\n\n# Run a prediction — three Path outputs should be uploaded, not base64-encoded\ncurl POST /predictions '{\"input\":{}}'\nstdout '\"status\":\"succeeded\"'\n# Outputs should be URLs pointing at the mock server, not data URIs\nstdout '\"output\":\\[\"http://host.docker.internal'\n! stdout 'data:image/'\n\n# Verify the mock server received exactly 3 PUT uploads\nupload-server-count 3\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pillow==10.4.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os\nimport tempfile\nfrom typing import Iterator\n\nfrom cog import BasePredictor, Path\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> Iterator[Path]:\n        for color in [\"red\", \"blue\", \"green\"]:\n            d = tempfile.mkdtemp()\n            p = os.path.join(d, f\"{color}.png\")\n            Image.new(\"RGB\", (10, 10), color).save(p)\n            yield Path(p)\n"
  },
  {
    "path": "integration-tests/tests/coglet_large_file_upload_serial.txtar",
    "content": "[short] skip 'large file test - slow'\n\n# Test that a large binary file (50 MiB) is successfully uploaded via --upload-url.\n#\n# Verifies the upload pipeline handles large binary payloads — not just large\n# strings through the IPC spill path (which coglet_large_output.txtar covers).\n\ncog build -t $TEST_IMAGE\n\nupload-server-start\ncog serve --upload-url $UPLOAD_SERVER_URL\n\n# Prediction generates a 50 MiB binary file\ncurl POST /predictions '{\"input\":{}}'\nstdout '\"status\":\"succeeded\"'\n# Output should be a URL, not inline data\nstdout '\"output\":\"http://host.docker.internal'\n! stdout 'data:'\n\n# Verify upload server received exactly 1 file\nupload-server-count 1\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os\nimport tempfile\n\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> Path:\n        d = tempfile.mkdtemp()\n        p = os.path.join(d, \"large_output.bin\")\n        # Write 50 MiB of binary data\n        with open(p, \"wb\") as f:\n            chunk = b\"\\xde\\xad\\xbe\\xef\" * 256  # 1 KiB chunk\n            for _ in range(50 * 1024):  # 50 MiB\n                f.write(chunk)\n        return Path(p)\n"
  },
  {
    "path": "integration-tests/tests/coglet_large_input.txtar",
    "content": "# Test that inputs larger than the 6 MiB IPC inline threshold spill to disk\n# and are rehydrated correctly by the worker.\n# Without input spilling this would exceed the 8 MiB LengthDelimitedCodec\n# frame limit and break the bridge.\n#\n# Strategy: generate a JSON file with a 7 MiB padding string on the host,\n# then POST it via the harness curl @file syntax. The predictor echoes\n# back len(padding) to prove the full input survived the spill-rehydrate\n# round-trip.\n#\n# Uses async prediction + webhook so the output goes directly to our\n# Go webhook receiver — never through testscript's log buffer.\n\nwebhook-server-start\ncog serve --upload-url http://unused/\n\n# Generate a ~7 MiB JSON request body.\n# dd produces 7340032 bytes of 'A', wrapped in a JSON prediction request.\nexec sh -c 'printf \"{\\\"id\\\":\\\"large-input-test\\\",\\\"input\\\":{\\\"padding\\\":\\\"\" > large_input.json && dd if=/dev/zero bs=1024 count=7168 2>/dev/null | tr \"\\0\" \"A\" >> large_input.json && printf \"\\\"},\\\"webhook\\\":\\\"$WEBHOOK_URL\\\",\\\"webhook_events_filter\\\":[\\\"completed\\\"]}\" >> large_input.json'\n\n# POST the large payload via @file\ncurl -H Prefer:respond-async POST /predictions @large_input.json\n\n# Wait for the webhook callback (up to 120s)\nwebhook-server-wait\n\n# Prediction succeeded — input was spilled and rehydrated correctly\nstdout '\"status\":\"succeeded\"'\n\n# Output is the string \"7340032\" (len of 7 MiB padding), which is 7 bytes\nstdout '\"output_size\":7'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, padding: str = \"\") -> str:\n        return str(len(padding))\n"
  },
  {
    "path": "integration-tests/tests/coglet_large_output.txtar",
    "content": "# Test that outputs larger than the 8MiB IPC frame limit spill to disk\n# and are reconstructed correctly by the orchestrator.\n# Without spilling this would panic the bridge and poison the slot.\n#\n# Uses async prediction + webhook so the 9MiB output goes directly to our\n# Go webhook receiver — never through testscript's log buffer.\n#\n# --upload-url is set to a dummy value so cog serve adds\n# --add-host=host.docker.internal:host-gateway (needed on Linux for the\n# webhook callback to reach the host). Nothing is actually uploaded because\n# the output is a plain string, not a Path.\n\nwebhook-server-start\ncog serve --upload-url http://unused/\n\n# Async prediction — server returns 202 immediately, delivers result to webhook\ncurl -H Prefer:respond-async POST /predictions '{\"id\":\"large-output-test\",\"webhook\":\"$WEBHOOK_URL\",\"webhook_events_filter\":[\"completed\"]}'\n\n# Wait for the webhook callback (up to 120s)\nwebhook-server-wait\n\n# 1. Prediction succeeded\nstdout '\"status\":\"succeeded\"'\n\n# 2. Output is correct — 9 * 1024 * 1024 = 9437184 bytes\nstdout '\"output_size\":9437184'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        # 9MiB string — exceeds the 8MiB IPC frame limit\n        return \"x\" * (9 * 1024 * 1024)\n"
  },
  {
    "path": "integration-tests/tests/coglet_list_path_single_element.txtar",
    "content": "# Test that List[Path] with a single element returns an array, not a scalar.\n# Regression test: the orchestrator must not collapse [url] → url for\n# single-element list outputs. The schema declares Output as \"type\": \"array\",\n# so the response must always be an array regardless of item count.\n\ncog build -t $TEST_IMAGE\n\n# Start mock upload server on the host, sets $UPLOAD_SERVER_URL\nupload-server-start\n\ncog serve --upload-url $UPLOAD_SERVER_URL\n\n# Single element: output MUST be [\"url\"], not \"url\"\ncurl POST /predictions '{\"input\":{\"count\":1}}'\nstdout '\"status\":\"succeeded\"'\n# Match array with exactly one URL element (not a bare string)\nstdout '\"output\":\\[\"http://host.docker.internal[^\"]*\"\\]'\n! stdout 'data:image/'\n\nupload-server-count 1\n\n# Multiple elements: output must be [\"url1\",\"url2\"]\ncurl POST /predictions '{\"input\":{\"count\":2}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output\":\\[\"http://host.docker.internal[^\"]*\",\"http://host.docker.internal[^\"]*\"\\]'\n! stdout 'data:image/'\n\n# 1 from first prediction + 2 from second = 3 total uploads\nupload-server-count 3\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pillow==10.4.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os\nimport tempfile\nfrom typing import List\n\nfrom cog import BasePredictor, Input, Path\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(self, count: int = Input(description=\"Number of images\", default=1)) -> List[Path]:\n        outputs = []\n        colors = [\"red\", \"blue\", \"green\", \"yellow\"]\n        for i in range(count):\n            d = tempfile.mkdtemp()\n            p = os.path.join(d, f\"{colors[i % len(colors)]}.png\")\n            Image.new(\"RGB\", (10, 10), colors[i % len(colors)]).save(p)\n            outputs.append(Path(p))\n        return outputs\n"
  },
  {
    "path": "integration-tests/tests/coglet_list_path_upload_url.txtar",
    "content": "# Test that List[Path] outputs are uploaded to --upload-url.\n# This verifies that list returns (not just iterators) go through the\n# FileOutput IPC path for upload instead of being base64-encoded inline.\n\ncog build -t $TEST_IMAGE\n\n# Start mock upload server on the host, sets $UPLOAD_SERVER_URL\nupload-server-start\n\ncog serve --upload-url $UPLOAD_SERVER_URL\n\n# Run a prediction — three Path outputs should be uploaded, not base64-encoded\ncurl POST /predictions '{\"input\":{}}'\nstdout '\"status\":\"succeeded\"'\n# Outputs should be URLs pointing at the mock server, not data URIs\nstdout '\"output\":\\[\"http://host.docker.internal'\n! stdout 'data:image/'\n\n# Verify the mock server received exactly 3 PUT uploads\nupload-server-count 3\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pillow==10.4.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os\nimport tempfile\nfrom typing import List\n\nfrom cog import BasePredictor, Path\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> List[Path]:\n        outputs = []\n        for color in [\"red\", \"blue\", \"green\"]:\n            d = tempfile.mkdtemp()\n            p = os.path.join(d, f\"{color}.png\")\n            Image.new(\"RGB\", (10, 10), color).save(p)\n            outputs.append(Path(p))\n        return outputs\n"
  },
  {
    "path": "integration-tests/tests/coglet_metrics.txtar",
    "content": "# Test that user-emitted metrics appear in sync prediction responses.\n#\n# Verifies:\n# 1. record_metric() with default mode (replace) appears in response\n# 2. Increment and append accumulation modes work\n# 3. predict_time is always present (system metric)\n# 4. predict_time overrides any user-set predict_time\n# 5. Dict-style metrics access (scope.metrics[\"key\"] = value)\n# 6. Dot-path keys create nested objects\n\ncog serve\n\n# Prediction that records metrics via current_scope()\ncurl POST /predictions '{\"input\":{}}'\nstdout '\"status\":\"succeeded\"'\n\n# User metrics present\nstdout '\"temperature\":0.7'\nstdout '\"token_count\":3'\nstdout '\"tags\":\\[\"fast\",\"cached\"\\]'\nstdout '\"timing\":\\{.*\"preprocess\":0.1'\nstdout '\"dict_metric\":\"hello\"'\n\n# System predict_time overrides user value — should be a float, not 999\nstdout '\"predict_time\":'\n! stdout '\"predict_time\":999'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, current_scope\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        scope = current_scope()\n\n        # Replace mode (default)\n        scope.record_metric(\"temperature\", 0.7)\n\n        # Increment mode: 1 + 2 = 3\n        scope.record_metric(\"token_count\", 1, mode=\"incr\")\n        scope.record_metric(\"token_count\", 2, mode=\"incr\")\n\n        # Append mode: builds array\n        scope.record_metric(\"tags\", \"fast\", mode=\"append\")\n        scope.record_metric(\"tags\", \"cached\", mode=\"append\")\n\n        # Dot-path key creates nested object\n        scope.record_metric(\"timing.preprocess\", 0.1)\n\n        # Dict-style access\n        scope.metrics[\"dict_metric\"] = \"hello\"\n\n        # User-set predict_time should be overridden by system\n        scope.record_metric(\"predict_time\", 999)\n\n        return \"ok\"\n"
  },
  {
    "path": "integration-tests/tests/coglet_metrics_webhook.txtar",
    "content": "# Test that user-emitted metrics appear in webhook payloads.\n#\n# Uses async prediction with webhook to verify metrics flow through\n# the supervisor and webhook sender. The webhook server captures the\n# terminal payload including metrics for assertion.\n#\n# --upload-url is set to a dummy value so cog serve adds\n# --add-host=host.docker.internal:host-gateway (needed on Linux for the\n# webhook callback to reach the host).\n\nwebhook-server-start\ncog serve --upload-url http://unused/\n\n# Async prediction — server returns 202, delivers result to webhook\ncurl -H Prefer:respond-async POST /predictions '{\"id\":\"metrics-webhook-test\",\"webhook\":\"$WEBHOOK_URL\",\"webhook_events_filter\":[\"completed\"]}'\n\n# Wait for the webhook callback\nwebhook-server-wait\n\n# Prediction succeeded\nstdout '\"status\":\"succeeded\"'\n\n# User metrics present in webhook payload\nstdout '\"model_version\":\"v2.1\"'\nstdout '\"confidence\":0.95'\n\n# System predict_time always present\nstdout '\"predict_time\":'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, current_scope\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        scope = current_scope()\n        scope.record_metric(\"model_version\", \"v2.1\")\n        scope.record_metric(\"confidence\", 0.95)\n        return \"ok\"\n"
  },
  {
    "path": "integration-tests/tests/coglet_single_path_output.txtar",
    "content": "# Test that a single Path return (not List[Path]) returns a scalar string, not an array.\n# This complements coglet_list_path_single_element.txtar — verifying that\n# the schema-driven output wrapping correctly distinguishes:\n#   -> Path       →  \"url\"    (scalar)\n#   -> List[Path] →  [\"url\"]  (array, even with one element)\n\ncog build -t $TEST_IMAGE\n\n# Start mock upload server on the host, sets $UPLOAD_SERVER_URL\nupload-server-start\n\ncog serve --upload-url $UPLOAD_SERVER_URL\n\n# Single Path output: must be \"url\" (scalar string), NOT [\"url\"]\ncurl POST /predictions '{\"input\":{}}'\nstdout '\"status\":\"succeeded\"'\n# Output should be a bare string URL, not wrapped in an array\nstdout '\"output\":\"http://host.docker.internal'\n! stdout '\"output\":\\['\n! stdout 'data:image/'\n\nupload-server-count 1\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pillow==10.4.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os\nimport tempfile\n\nfrom cog import BasePredictor, Path\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> Path:\n        d = tempfile.mkdtemp()\n        p = os.path.join(d, \"output.png\")\n        Image.new(\"RGB\", (10, 10), \"red\").save(p)\n        return Path(p)\n"
  },
  {
    "path": "integration-tests/tests/complex_output.txtar",
    "content": "# Test complex/structured output type using cog.BaseModel (dataclass)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict returns structured output\ncog predict $TEST_IMAGE -i msg='test error'\nstdout '\"success\": false'\nstdout '\"error\": \"test error\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom typing import Optional\n\nfrom cog import BaseModel, BasePredictor, Path\n\n\nclass ModelOutput(BaseModel):\n    success: bool\n    error: Optional[str]\n    segmented_image: Optional[Path]\n\n\nclass Predictor(BasePredictor):\n    def predict(self, msg: str) -> ModelOutput:\n        return ModelOutput(success=False, error=msg, segmented_image=None)\n"
  },
  {
    "path": "integration-tests/tests/concatenate_iterator_output.txtar",
    "content": "# Test ConcatenateIterator[str] as predict output type\n#\n# ConcatenateIterator is the primary streaming text output type for LLMs.\n# cog predict renders each yielded token as an array element.\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Streaming output yields individual tokens\ncog predict $TEST_IMAGE -i prompt=hello\nstdout '\"hello\"'\nstdout '\" world\"'\nstdout '\" !\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, ConcatenateIterator\n\n\nclass Predictor(BasePredictor):\n    def predict(self, prompt: str) -> ConcatenateIterator[str]:\n        for token in [prompt, \" world\", \" !\"]:\n            yield token\n"
  },
  {
    "path": "integration-tests/tests/config_subdirectory.txtar",
    "content": "# Test that cog.yaml is discovered from subdirectories\n# Source: test_config.py::test_config\n\n# Create a subdirectory structure\nmkdir some\nmkdir some/sub\nmkdir some/sub/dir\n\n# Run cog from a subdirectory - it should find cog.yaml in parent\ncd some/sub/dir\ncog run echo hello world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n"
  },
  {
    "path": "integration-tests/tests/debug_secrets.txtar",
    "content": "# Test that cog debug shows secret mount syntax in Dockerfile output\n# Source: test_run.py::test_run_with_secret\n\n# Run cog debug to see the generated Dockerfile\ncog debug\nstdout 'RUN echo hello world'\nstdout 'RUN --mount=type=secret,id=foo,target=secret.txt echo shh'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  run:\n    - echo hello world\n    - command: >-\n        echo shh\n      mounts:\n        - type: secret\n          id: foo\n          target: secret.txt\npredict: \"predict.py:Predictor\"\n\n-- secret.txt --\nsecret content here\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/dict_output.txtar",
    "content": "# Test bare dict return type works for predict output\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict returns a dict\ncog predict $TEST_IMAGE -i name=alice\nstdout '\"greeting\": \"hello alice\"'\nstdout '\"length\": 5'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, name: str) -> dict:\n        return {\"greeting\": \"hello \" + name, \"length\": len(name)}\n"
  },
  {
    "path": "integration-tests/tests/emit_metric_deprecated.txtar",
    "content": "# Test that a predictor using the deprecated emit_metric() still builds,\n# runs, and records metrics correctly.\n#\n# emit_metric() was dropped without a compat shim in 0.17.0, causing a hard\n# ImportError on startup for models that use it. This test covers the two\n# call styles used across replicate/* models:\n#\n#   from cog import emit_metric; emit_metric(\"key\", value)\n#   cog.emit_metric(\"key\", value)\n\ncog serve\n\n# Prediction works and the metric appears in the response\ncurl POST /predictions '{\"input\":{}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output_tokens\":42'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport cog\nfrom cog import BasePredictor, emit_metric\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        # Both call styles should work\n        emit_metric(\"output_tokens\", 42)\n        cog.emit_metric(\"input_tokens\", 10)\n        return \"ok\"\n"
  },
  {
    "path": "integration-tests/tests/env_vars.txtar",
    "content": "# Build the image\ncog build -t $TEST_IMAGE\n\n# Environment variables are set\ncog predict $TEST_IMAGE -i name=TEST_VAR\nstdout 'test_value'\n\ncog predict $TEST_IMAGE -i name=NAME\nstdout 'michael'\n\n-- cog.yaml --\npredict: \"predict.py:Predictor\"\nbuild:\n  python_version: \"3.12\"\nenvironment:\n  - NAME=michael\n  - TEST_VAR=test_value\n\n-- predict.py --\nfrom cog import BasePredictor\nimport os\n\nclass Predictor(BasePredictor):\n    def predict(self, name: str) -> str:\n        return f\"ENV[{name}]={os.getenv(name)}\"\n"
  },
  {
    "path": "integration-tests/tests/experimental_feature_warning.txtar",
    "content": "# Test that a predictor importing the deprecated ExperimentalFeatureWarning\n# still builds and runs successfully.\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Prediction works despite the deprecated import\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport warnings\nfrom cog import BasePredictor, ExperimentalFeatureWarning\n\nwarnings.filterwarnings(\"ignore\", category=ExperimentalFeatureWarning)\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/ffmpeg_package.txtar",
    "content": "# Test that ffmpeg system package is installed (common ML dependency)\n\n# Build the image (the run command verifies ffmpeg is installed)\ncog build -t $TEST_IMAGE\n\n# Verify the predictor works\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  python_version: \"3.12\"\n  python_packages:\n    - \"torch==2.5.1\"\n  cuda: \"12.4\"\n  run:\n    - command: ffmpeg --help\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/file_input.txtar",
    "content": "# Build the image\ncog build -t $TEST_IMAGE\n\n# File input works\ncog predict $TEST_IMAGE -i file=@input.txt\nstdout 'hello from file'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Path\n\nclass Predictor(BasePredictor):\n    def predict(self, file: Path) -> str:\n        return file.read_text()\n\n-- input.txt --\nhello from file\n"
  },
  {
    "path": "integration-tests/tests/file_list_input.txtar",
    "content": "\n# Test list[File] input type (multiple file inputs using File type)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict with multiple file inputs\ncog predict $TEST_IMAGE -i files=@file1.txt -i files=@file2.txt\nstdout 'content one'\nstdout 'content two'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, File\n\n\nclass Predictor(BasePredictor):\n    def predict(self, files: list[File]) -> str:\n        output_parts = []\n        for f in files:\n            content = f.read()\n            if isinstance(content, bytes):\n                content = content.decode('utf-8')\n            output_parts.append(content)\n        return \"\\n\\n\".join(output_parts)\n\n-- file1.txt --\ncontent one\n-- file2.txt --\ncontent two\n"
  },
  {
    "path": "integration-tests/tests/float_input_output.txtar",
    "content": "# Test float input and output types work correctly\ncog build -t $TEST_IMAGE\n\n# Float input and output works\ncog predict $TEST_IMAGE -i num=10\nstdout '20'\n\n# Negative numbers work\ncog predict $TEST_IMAGE -i num=-10\nstdout '-20'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.11\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(\n        self, num: float = Input(description=\"Number of things\")\n    ) -> float:\n        return num * 2.0\n"
  },
  {
    "path": "integration-tests/tests/function_predictor.txtar",
    "content": "# Test function-based predictor (no class, just a function)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict using function-based predictor\ncog predict $TEST_IMAGE -i prompt=world\nstdout 'HELLO WORLD'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.13\"\npredict: \"predict.py:run\"\n\n-- predict.py --\nfrom cog import Input\n\n\ndef run(\n    prompt: str = Input(),\n) -> str:\n    return f\"HELLO {prompt.upper()}\"\n"
  },
  {
    "path": "integration-tests/tests/future_annotations.txtar",
    "content": "\n# Test from __future__ import annotations support\n# This tests that future annotations work correctly with Path types\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict - creates and returns an image\ncog predict $TEST_IMAGE\n\n# Verify output file was created\nexec test -f output.png\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - pillow==10.0.0\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom __future__ import annotations\n\nfrom cog import BasePredictor, Input, Path\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> Path:\n        \"\"\"Create and return a simple test image.\"\"\"\n        # Create a simple red image\n        img = Image.new(\"RGB\", (100, 100), color=\"red\")\n        output_path = Path(\"/tmp/output.png\")\n        img.save(output_path)\n        return output_path\n"
  },
  {
    "path": "integration-tests/tests/glb_project.txtar",
    "content": "# Test GLB (3D model) file output\n\n# Create a minimal GLB placeholder file (GLB files start with \"glTF\" magic bytes)\nexec sh -c 'printf \"glTF\" > mesh.glb'\n\n# Build and predict - verifies GLB file output works\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.13\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def setup(self) -> None:\n        if not Path(\"mesh.glb\").exists():\n            raise ValueError(\"Example file mesh.glb does not exist\")\n\n    def predict(self) -> Path:\n        return Path(\"mesh.glb\")\n"
  },
  {
    "path": "integration-tests/tests/granite_project.txtar",
    "content": "# Test that Pydantic 2 is not clobbered to a <2 version\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict and verify pydantic version is preserved\ncog predict $TEST_IMAGE\nstdout '2.11.9'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.11\"\n  python_requirements: requirements.txt\npredict: \"predict.py:Predictor\"\n\n-- requirements.txt --\npydantic==2.11.9\n\n-- predict.py --\nfrom cog import BasePredictor\nimport pydantic\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        return pydantic.__version__\n"
  },
  {
    "path": "integration-tests/tests/healthcheck.txtar",
    "content": "# Test custom healthcheck functionality\n# This tests the user-defined healthcheck() method in predictors\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Test 1: Healthy healthcheck returns READY status\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n! stdout 'user_healthcheck_error'\n\n# Test 2: Make a prediction to ensure predictor works\ncurl POST /predictions '{\"input\":{\"text\":\"world\"}}'\nstdout '\"output\":\"hello world\"'\n\n# Test 3: Health check still works after prediction\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return f\"hello {text}\"\n    \n    def healthcheck(self) -> bool:\n        \"\"\"Custom healthcheck that always returns healthy.\"\"\"\n        return True\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_async.txtar",
    "content": "# Test async healthcheck function\n# Ensures async healthchecks work correctly\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Test async healthcheck returns healthy\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n! stdout 'user_healthcheck_error'\n\n# Make a prediction\ncurl POST /predictions '{\"input\":{\"text\":\"world\"}}'\nstdout '\"output\":\"hello world\"'\n\n# Healthcheck should still work\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport asyncio\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return f\"hello {text}\"\n\n    async def healthcheck(self) -> bool:\n        \"\"\"Async healthcheck function.\"\"\"\n        # Simulate async operation\n        await asyncio.sleep(0.1)\n        return True\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_async_exception.txtar",
    "content": "# Test async healthcheck that raises an exception\n# Tests that async def healthcheck() exceptions are handled\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Test: Async healthcheck raising exception gives UNHEALTHY with error\ncurl GET /health-check\nstdout '\"status\":\"UNHEALTHY\"'\nstdout 'user_healthcheck_error'\nstdout 'Async healthcheck error'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\nimport asyncio\n\nclass Predictor(BasePredictor):\n    async def predict(self, text: str) -> str:\n        return f\"hello {text}\"\n    \n    async def healthcheck(self) -> bool:\n        \"\"\"Async healthcheck that raises an exception.\"\"\"\n        await asyncio.sleep(0.01)\n        raise RuntimeError(\"Async healthcheck error\")\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_async_timeout.txtar",
    "content": "# Test async healthcheck timeout behavior\n# Tests that async def healthcheck() timeouts are handled (5 second limit)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Server should be healthy initially\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n# Trigger slow healthcheck mode via prediction\ncurl POST /predictions '{\"input\":{\"text\":\"trigger_slow\"}}'\nstdout '\"status\":\"succeeded\"'\n\n# Now healthcheck should timeout and return UNHEALTHY\ncurl GET /health-check\nstdout '\"status\":\"UNHEALTHY\"'\nstdout 'user_healthcheck_error'\nstdout 'timed out after 5.0 seconds'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\nimport asyncio\n\nclass Predictor(BasePredictor):\n    def setup(self) -> None:\n        self._slow_mode = False\n\n    async def predict(self, text: str) -> str:\n        if text == \"trigger_slow\":\n            self._slow_mode = True\n        return f\"hello {text}\"\n    \n    async def healthcheck(self) -> bool:\n        \"\"\"Async healthcheck that times out when triggered.\"\"\"\n        if self._slow_mode:\n            await asyncio.sleep(10)  # Sleep longer than 5s timeout\n        return True\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_async_unhealthy.txtar",
    "content": "# Test async unhealthy healthcheck behavior\n# Tests that async def healthcheck() returning False works correctly\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Test: Async healthcheck returning False gives UNHEALTHY\ncurl GET /health-check\nstdout '\"status\":\"UNHEALTHY\"'\nstdout 'user_healthcheck_error'\nstdout 'returned False'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\nimport asyncio\n\nclass Predictor(BasePredictor):\n    async def predict(self, text: str) -> str:\n        return f\"hello {text}\"\n    \n    async def healthcheck(self) -> bool:\n        \"\"\"Async healthcheck that returns unhealthy.\"\"\"\n        await asyncio.sleep(0.01)\n        return False\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_during_prediction.txtar",
    "content": "# Test healthcheck called during active prediction\n# Ensures healthcheck pipe doesn't interfere with prediction pipe\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server with concurrency enabled\ncog serve\n\n# Start a long-running prediction in background (5 seconds)\nexec bash -c 'curl -s -X POST $SERVER_URL/predictions -H \"Content-Type: application/json\" -d \"{\\\"input\\\":{\\\"sleep_time\\\":5}}\" > /tmp/prediction.json &'\n\n# Wait for prediction to start (500ms)\nexec sleep 0.5\n\n# Call healthcheck while prediction is running\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n# Call healthcheck again (multiple times during prediction)\nexec sleep 1\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\nexec sleep 1\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n# Wait for prediction to complete\nexec sleep 3\n\n# Verify prediction succeeded\nexec bash -c 'cat /tmp/prediction.json | grep -q \"\\\"output\\\":\\\"slept for 5 seconds\\\"\"'\n\n# Healthcheck should still work after prediction completes\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport time\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, sleep_time: int) -> str:\n        \"\"\"Sleep for specified seconds.\"\"\"\n        time.sleep(sleep_time)\n        return f\"slept for {sleep_time} seconds\"\n\n    def healthcheck(self) -> bool:\n        \"\"\"Healthcheck should work during predictions.\"\"\"\n        return True\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_exception.txtar",
    "content": "# Test healthcheck that raises an exception\n# This tests error handling when healthcheck throws\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Exception in healthcheck should return UNHEALTHY with error message\ncurl GET /health-check\nstdout '\"status\":\"UNHEALTHY\"'\nstdout 'user_healthcheck_error'\nstdout 'Critical system error'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def setup(self) -> None:\n        self._healthcheck_calls = 0\n\n    def predict(self, text: str) -> str:\n        return f\"hello {text}\"\n    \n    def healthcheck(self) -> bool:\n        \"\"\"Healthcheck that raises an exception after startup.\"\"\"\n        self._healthcheck_calls += 1\n        if self._healthcheck_calls == 1:\n            return True\n        raise RuntimeError(\"Critical system error\")\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_immediately_after_prediction.txtar",
    "content": "# Test healthcheck immediately after prediction completes\n# Ensures healthcheck works correctly in quick succession with predictions\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Make a quick prediction\ncurl POST /predictions '{\"input\":{\"text\":\"world\"}}'\nstdout '\"output\":\"hello world\"'\n\n# Immediately call healthcheck after prediction (no delay)\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n# Do it again - rapid fire prediction + healthcheck\ncurl POST /predictions '{\"input\":{\"text\":\"again\"}}'\nstdout '\"output\":\"hello again\"'\n\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n# One more time to ensure pattern holds\ncurl POST /predictions '{\"input\":{\"text\":\"final\"}}'\nstdout '\"output\":\"hello final\"'\n\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return f\"hello {text}\"\n\n    def healthcheck(self) -> bool:\n        \"\"\"Healthcheck should work immediately after predictions.\"\"\"\n        return True\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_repeated_calls.txtar",
    "content": "# Test repeated healthcheck calls (simulates supervisor container polling pattern)\n# This ensures no resource leaks over many healthcheck calls\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Call healthcheck 50 times sequentially (simulates supervisor polling)\n# Each call should succeed and return READY status\nexec bash -c 'for i in {1..50}; do curl -s $SERVER_URL/health-check | grep -q \"\\\"status\\\":\\\"READY\\\"\" || exit 1; done'\n\n# Make a prediction to ensure system still works after many healthchecks\ncurl POST /predictions '{\"input\":{\"text\":\"world\"}}'\nstdout '\"output\":\"hello world\"'\n\n# Healthcheck should still work after prediction\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def __init__(self):\n        self.call_count = 0\n\n    def predict(self, text: str) -> str:\n        return f\"hello {text}\"\n\n    def healthcheck(self) -> bool:\n        \"\"\"Custom healthcheck that tracks call count.\"\"\"\n        self.call_count += 1\n        # Always return healthy\n        return True\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_timeout.txtar",
    "content": "# Test healthcheck timeout behavior\n# This tests when sync healthcheck takes too long (>5 seconds)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Server should be healthy initially\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\n\n# Trigger slow healthcheck mode via prediction\ncurl POST /predictions '{\"input\":{\"text\":\"trigger_slow\"}}'\nstdout '\"status\":\"succeeded\"'\n\n# Now healthcheck should timeout and return UNHEALTHY\ncurl GET /health-check\nstdout '\"status\":\"UNHEALTHY\"'\nstdout 'user_healthcheck_error'\nstdout 'timed out after 5.0 seconds'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport time\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def setup(self) -> None:\n        self._slow_mode = False\n\n    def predict(self, text: str) -> str:\n        if text == \"trigger_slow\":\n            self._slow_mode = True\n        return f\"hello {text}\"\n    \n    def healthcheck(self) -> bool:\n        \"\"\"Sync healthcheck that times out when triggered.\"\"\"\n        if self._slow_mode:\n            time.sleep(10)  # Sleep longer than 5s timeout\n        return True\n"
  },
  {
    "path": "integration-tests/tests/healthcheck_unhealthy.txtar",
    "content": "# Test unhealthy healthcheck behavior\n# This tests when healthcheck returns False\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Unhealthy healthcheck should return UNHEALTHY status\ncurl GET /health-check\nstdout '\"status\":\"UNHEALTHY\"'\nstdout 'user_healthcheck_error'\nstdout 'user-defined healthcheck returned False'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def setup(self) -> None:\n        self._healthcheck_calls = 0\n\n    def predict(self, text: str) -> str:\n        return f\"hello {text}\"\n    \n    def healthcheck(self) -> bool:\n        \"\"\"Unhealthy healthcheck after startup.\"\"\"\n        self._healthcheck_calls += 1\n        if self._healthcheck_calls == 1:\n            return True\n        return False\n"
  },
  {
    "path": "integration-tests/tests/int_input_output.txtar",
    "content": "# Test integer input and output types work correctly\ncog build -t $TEST_IMAGE\n\n# Integer input and output works\ncog predict $TEST_IMAGE -i num=10\nstdout '20'\n\n# Negative numbers work\ncog predict $TEST_IMAGE -i num=-10\nstdout '-20'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.11\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(\n        self, num: int = Input(description=\"Number of things\")\n    ) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/int_none_output.txtar",
    "content": "\n# Test int return type that returns None\n# This tests the handling of None values for typed outputs\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict returns None despite int type annotation\n# When None is returned, cog shows \"No output generated\" on stderr\ncog predict $TEST_IMAGE\nstderr 'No output generated'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> int:\n        return None\n"
  },
  {
    "path": "integration-tests/tests/int_predictor.txtar",
    "content": "# Build the image\ncog build -t $TEST_IMAGE\n\n# Integer input and output works\ncog predict $TEST_IMAGE -i num=5\nstdout '10'\n\n# Negative numbers work\ncog predict $TEST_IMAGE -i num=-3\nstdout '-6'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, num: int) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/invalid_int_validation.txtar",
    "content": "\n# Test input validation with ge/le constraints\n# Build should fail because default=1 violates ge=2 constraint\n\n# Build should fail at schema validation\n! cog build -t $TEST_IMAGE\nstderr 'invalid'\nstderr 'minimum'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input\n\n\nclass Predictor(BasePredictor):\n    def predict(\n        self, num: int = Input(description=\"Number of things\", default=1, ge=2, le=10)\n    ) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/iterator_error_midstream.txtar",
    "content": "# Test that a generator which yields items then raises an exception\n# correctly reports failure. This is a common real-world failure mode\n# (model produces partial output then hits an error).\n\nwebhook-server-start\ncog serve --upload-url http://unused/\n\n# Async prediction — generator yields 3 items then raises\ncurl -H Prefer:respond-async POST /predictions '{\"id\":\"iter-error-test\",\"webhook\":\"$WEBHOOK_URL\",\"webhook_events_filter\":[\"completed\"]}'\n\nwebhook-server-wait\n\n# Prediction should fail\nstdout '\"status\":\"failed\"'\nstdout '\"has_error\":true'\nstdout '\"error_message\":\".*generator exploded\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom typing import Iterator\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> Iterator[str]:\n        yield \"chunk-1\"\n        yield \"chunk-2\"\n        yield \"chunk-3\"\n        raise RuntimeError(\"generator exploded\")\n"
  },
  {
    "path": "integration-tests/tests/iterator_string_output.txtar",
    "content": "# Test Iterator[str] as predict output type\n#\n# Iterator[str] yields individual string items as an array.\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Iterator output returns items\ncog predict $TEST_IMAGE -i count=3\nstdout 'item-0'\nstdout 'item-1'\nstdout 'item-2'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom typing import Iterator\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, count: int) -> Iterator[str]:\n        for i in range(count):\n            yield f\"item-{i}\"\n"
  },
  {
    "path": "integration-tests/tests/legacy_sdk_schema.txtar",
    "content": "# Test that building with a legacy SDK (< 0.17.0) falls back to runtime\n# schema generation instead of using the static Go tree-sitter parser.\n#\n# SDK 0.16.12 uses pydantic-based runtime introspection for schema generation\n# and predates coglet. Coglet is not installed because:\n#   1. No explicit COGLET_WHEEL is set\n#   2. The SDK dependency handles it — and 0.16.12 doesn't depend on coglet\n\n# Override SDK wheel to use PyPI 0.16.12 (legacy, pre-coglet)\nenv COG_SDK_WHEEL=pypi:0.16.12\n\n# Build should succeed — without COG_STATIC_SCHEMA=1 the build\n# automatically uses the legacy runtime schema generation path.\ncog build -t $TEST_IMAGE\n\n# Predict should work with the legacy SDK's built-in Python HTTP server\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/list_int_input_output.txtar",
    "content": "# Test list[int] as predict input and output types\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# List of ints works as input and output\ncog predict $TEST_IMAGE --json '{\"numbers\": [1, 2, 3]}'\nstdout '\"status\": \"succeeded\"'\nstdout '\"output\":'\nstdout '2'\nstdout '4'\nstdout '6'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, numbers: list[int]) -> list[int]:\n        return [n * 2 for n in numbers]\n"
  },
  {
    "path": "integration-tests/tests/list_string_output.txtar",
    "content": "# Test list[str] as predict output type\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# List output returns items\ncog predict $TEST_IMAGE -i text='hello world foo'\nstdout 'hello'\nstdout 'world'\nstdout 'foo'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> list[str]:\n        return text.split()\n"
  },
  {
    "path": "integration-tests/tests/many_inputs.txtar",
    "content": "\n# Test predictor with many different input types\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict with various input types\ncog predict $TEST_IMAGE -i no_default=hello -i path=@path.txt -i image=@image.jpg -i choices=foo -i int_choices=3\nstdout 'hello default 20 world jpg foo 6'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input, Path\n\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        no_default: str,\n        default_without_input: str = \"default\",\n        input_with_default: int = Input(default=10),\n        path: Path = Input(description=\"Some path\"),\n        image: Path = Input(description=\"Some path\"),\n        choices: str = Input(choices=[\"foo\", \"bar\"]),\n        int_choices: int = Input(description=\"hello\", choices=[3, 4, 5]),\n    ) -> str:\n        with path.open() as f:\n            path_contents = f.read().strip()\n        image_extension = str(image).split(\".\")[-1]\n        return (\n            no_default\n            + \" \"\n            + default_without_input\n            + \" \"\n            + str(input_with_default * 2)\n            + \" \"\n            + path_contents\n            + \" \"\n            + image_extension\n            + \" \"\n            + choices\n            + \" \"\n            + str(int_choices * 2)\n        )\n\n-- path.txt --\nworld\n-- image.jpg --\nfake image content\n"
  },
  {
    "path": "integration-tests/tests/multi_file_schema.txtar",
    "content": "# Test that schema generation works when the output type is defined in a\n# separate Python module. This exercises the cross-file model resolution:\n# the predictor imports Output from output_types.py, and the static Go parser\n# finds and parses output_types.py to resolve the BaseModel fields.\n#\n# Covers:\n#   - Output type imported from a sibling module (from output_types import Output)\n#   - BaseModel fields appear correctly in the OpenAPI schema label\n#   - Prediction works end-to-end with the multi-file setup\n\n# Opt in to static schema generation\nenv COG_STATIC_SCHEMA=1\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Verify schema is in Docker label with correct structure\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"run.cog.openapi_schema\"}}'\n\n# Output schema should be an object with fields from output_types.py\nstdout '\"type\":\"object\"'\nstdout '\"text\":'\nstdout '\"score\":'\nstdout '\"tags\":'\n\n# Input should have the prompt field\nstdout '\"prompt\":'\nstdout '\"required\":\\[\"prompt\"\\]'\n\n# Predict should work end-to-end\ncog predict $TEST_IMAGE -i prompt=hello\nstdout '\"text\": \"hello\"'\nstdout '\"score\": 1'\nstdout '\"tags\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pydantic>2\"\npredict: \"predict.py:Predictor\"\n\n-- output_types.py --\nfrom pydantic import BaseModel\n\n\nclass Output(BaseModel):\n    text: str\n    score: float\n    tags: list[str]\n\n-- predict.py --\nfrom cog import BasePredictor\nfrom output_types import Output\n\n\nclass Predictor(BasePredictor):\n    def predict(self, prompt: str) -> Output:\n        return Output(text=prompt, score=1.0, tags=[\"default\"])\n"
  },
  {
    "path": "integration-tests/tests/nested_output_types.txtar",
    "content": "# Test structured output type with multiple field types.\n#\n# Verifies that coglet correctly serializes a BaseModel output containing:\n# - Multiple primitive types (str, int, float, bool)\n# - Optional fields (with value and without)\n# - List of primitive types\n#\n# Note: nested BaseModel fields (e.g., a field typed as another BaseModel)\n# are NOT supported by cog's type system. This test covers the supported\n# complex output patterns.\n\ncog serve\n\ncurl POST /predictions '{\"input\":{\"name\":\"test\"}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"name\":\"test\"'\nstdout '\"count\":42'\nstdout '\"score\":0.95'\nstdout '\"passed\":true'\nstdout '\"tags\":\\[\"fast\",\"cached\"\\]'\nstdout '\"note\":\"extra info\"'\nstdout '\"extra\":null'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom typing import List, Optional\n\nfrom cog import BaseModel, BasePredictor\n\n\nclass Output(BaseModel):\n    name: str\n    count: int\n    score: float\n    passed: bool\n    tags: List[str]\n    note: Optional[str]\n    extra: Optional[str]\n\n\nclass Predictor(BasePredictor):\n    def predict(self, name: str) -> Output:\n        return Output(\n            name=name,\n            count=42,\n            score=0.95,\n            passed=True,\n            tags=[\"fast\", \"cached\"],\n            note=\"extra info\",\n            extra=None,\n        )\n"
  },
  {
    "path": "integration-tests/tests/no_predictor.txtar",
    "content": "# Test error when no predictor is defined\n# Build should fail when cog.yaml has no predict field\n\n# Build should fail with error about missing predictor\n! cog build -t $TEST_IMAGE\nstderr 'predict'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n"
  },
  {
    "path": "integration-tests/tests/non_base_predictor_class.txtar",
    "content": "\n# Build image\ncog build -t $TEST_IMAGE\n\n# Predict using class without BasePredictor\ncog predict $TEST_IMAGE -i text=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n\n-- predict.py --\nfrom cog import Input\n\n\nclass Predictor:\n    def predict(self, text: str = Input(default=\"world\")) -> str:\n        return f\"hello {text}\"\n"
  },
  {
    "path": "integration-tests/tests/non_base_predictor_function.txtar",
    "content": "\n# Build image\ncog build -t $TEST_IMAGE\n\n# Predict using standalone function\ncog predict $TEST_IMAGE -i text=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: predict.py:predict\n\n-- predict.py --\nfrom cog import Input\n\n\ndef predict(text: str = Input(default=\"world\")) -> str:\n    return f\"hello {text}\"\n"
  },
  {
    "path": "integration-tests/tests/oci_bundle_build.txtar",
    "content": "# Test building an OCI bundle with declarative weights in cog.yaml.\n# Verifies: cog.yaml weights declaration -> cog weights build -> cog build (COG_OCI_INDEX=1)\n# The image should build successfully and predictions should work.\n\n# Create weight files (small, deterministic)\nmkdir weights\nexec sh -c 'dd if=/dev/zero bs=1024 count=1 2>/dev/null | tr \"\\0\" \"A\" > weights/model-a.bin'\nexec sh -c 'dd if=/dev/zero bs=1024 count=1 2>/dev/null | tr \"\\0\" \"B\" > weights/model-b.bin'\n\n# Step 1: Build weights.lock from cog.yaml declarations\ncog weights build\nstderr 'Generated weights.lock'\nstderr '2 file'\nexists weights.lock\n\n# Step 2: Build with OCI index mode enabled\nenv COG_OCI_INDEX=1\ncog build -t $TEST_IMAGE\nstderr 'Image built as'\n\n# Verify image was built\nexec docker image inspect $TEST_IMAGE\nstdout 'run.cog.config'\n\n# Verify prediction works\ncog predict $TEST_IMAGE -i text=hello\nstdout 'processed: hello'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\nweights:\n  - name: alpha\n    source: weights/model-a.bin\n    target: /weights/model-a.bin\n  - name: beta\n    source: weights/model-b.bin\n    target: /weights/model-b.bin\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return f\"processed: {text}\"\n"
  },
  {
    "path": "integration-tests/tests/oci_bundle_inspect.txtar",
    "content": "# Test cog inspect on a pushed OCI bundle with declarative weights.\n# Verifies: push bundle -> cog inspect --remote --json shows correct structure.\n# The inspect output should show an OCI index with image + weight manifests.\n\n[short] skip 'requires local registry'\n\n# Start test registry\nregistry-start\n\n# Create weight files (small, deterministic)\nmkdir weights\nexec sh -c 'dd if=/dev/zero bs=1024 count=1 2>/dev/null | tr \"\\0\" \"A\" > weights/model-a.bin'\nexec sh -c 'dd if=/dev/zero bs=1024 count=1 2>/dev/null | tr \"\\0\" \"B\" > weights/model-b.bin'\n\n# Build weights.lock\ncog weights build\nstderr 'Generated weights.lock'\nexists weights.lock\n\n# Build and push with OCI index mode\nenv COG_OCI_INDEX=1\ncog push $TEST_REGISTRY/test/inspect-model:v1\n\n# Inspect the pushed bundle\ncog inspect --remote --json $TEST_REGISTRY/test/inspect-model:v1\n\n# Verify it's an OCI index\nstdout '\"type\": \"index\"'\n\n# Verify image manifest is present\nstdout '\"type\": \"image\"'\n\n# Verify weight manifests are present with correct names\nstdout '\"type\": \"weights\"'\nstdout '\"name\": \"alpha\"'\nstdout '\"name\": \"beta\"'\n\n# Verify weight targets\nstdout '\"target\": \"/weights/model-a.bin\"'\nstdout '\"target\": \"/weights/model-b.bin\"'\n\n# Verify layers are populated\nstdout '\"layers\"'\nstdout '\"digest\": \"sha256:'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\nweights:\n  - name: alpha\n    source: weights/model-a.bin\n    target: /weights/model-a.bin\n  - name: beta\n    source: weights/model-b.bin\n    target: /weights/model-b.bin\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return f\"processed: {text}\"\n"
  },
  {
    "path": "integration-tests/tests/oci_bundle_push.txtar",
    "content": "# Test pushing an OCI bundle with declarative weights via cog push.\n# Verifies: cog.yaml weights -> cog weights build -> cog push (BundlePusher path)\n# The push should create an OCI index with image + weight manifests.\n\n[short] skip 'requires local registry'\n\n# Start test registry\nregistry-start\n\n# Create weight files (small, deterministic)\nmkdir weights\nexec sh -c 'dd if=/dev/zero bs=1024 count=1 2>/dev/null | tr \"\\0\" \"A\" > weights/model-a.bin'\nexec sh -c 'dd if=/dev/zero bs=1024 count=1 2>/dev/null | tr \"\\0\" \"B\" > weights/model-b.bin'\n\n# Step 1: Build weights.lock\ncog weights build\nstderr 'Generated weights.lock'\nexists weights.lock\n\n# Step 2: Build and push with OCI index mode\nenv COG_OCI_INDEX=1\ncog push $TEST_REGISTRY/test/bundle-model:v1\n\n# Verify push succeeded — should mention pushing\nstderr -count=1 'Pushing'\n\n# Step 3: Verify the pushed artifact is an OCI index with image + weight manifests\nregistry-inspect $TEST_REGISTRY/test/bundle-model:v1\n# Verify it's an OCI index\nstdout 'application/vnd.oci.image.index.v1\\+json'\n# Verify weight annotations are present\nstdout 'vnd.cog.reference.type.*weights'\nstdout 'vnd.cog.weight.name.*alpha'\nstdout 'vnd.cog.weight.name.*beta'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\nweights:\n  - name: alpha\n    source: weights/model-a.bin\n    target: /weights/model-a.bin\n  - name: beta\n    source: weights/model-b.bin\n    target: /weights/model-b.bin\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str) -> str:\n        return f\"processed: {text}\"\n"
  },
  {
    "path": "integration-tests/tests/optional_path_input.txtar",
    "content": "# Test optional Path input with None default (Pydantic 1 compatibility)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict without providing optional input - should return red image\ncog predict $TEST_IMAGE\n\n# Verify output file was created\nexec test -f output.webp\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pillow\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Path, Input\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        test_image: Path | None = Input(description=\"Test image\", default=None)\n    ) -> Path:\n        \"\"\"Run a single prediction on the model\"\"\"\n        im = Image.new(\"RGB\", (100, 100), color=\"red\")\n        im.save(Path(\"./hello.webp\"))\n        return Path(\"./hello.webp\")\n"
  },
  {
    "path": "integration-tests/tests/path_input.txtar",
    "content": "# Test Path input type (read file content)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict reads content from file\ncog predict $TEST_IMAGE -i path=@input.txt\nstdout 'hello from file'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def predict(self, path: Path) -> str:\n        with open(path) as f:\n            return f.read()\n\n-- input.txt --\nhello from file\n"
  },
  {
    "path": "integration-tests/tests/path_input_output.txtar",
    "content": "# Test Path input and output with setup method\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict with text and path input, returns path output\ncog predict $TEST_IMAGE -i text=bar -i path=@input.txt\n\n# Verify output file was created and contains expected content\nexec test -f output.txt\nexec grep -q foobarbaz output.txt\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n\n-- predict.py --\nimport tempfile\n\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def setup(self):\n        self.foo = \"foo\"\n\n    def predict(self, text: str, path: Path) -> Path:\n        with open(path) as f:\n            output = self.foo + text + f.read()\n        tmpdir = Path(tempfile.mkdtemp())\n        with open(tmpdir / \"output.txt\", \"w\") as fh:\n            fh.write(output)\n        return tmpdir / \"output.txt\"\n\n-- input.txt --\nbaz\n"
  },
  {
    "path": "integration-tests/tests/path_list_input.txtar",
    "content": "# Test list[Path] input type (multiple file inputs)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict with multiple file inputs\ncog predict $TEST_IMAGE -i paths=@1.txt -i paths=@2.txt\nstdout 'test1'\nstdout 'test2'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def predict(self, paths: list[Path]) -> str:\n        output_parts = []\n        for path in paths:\n            with open(path) as f:\n                output_parts.append(f.read())\n        return \"\".join(output_parts)\n\n-- 1.txt --\ntest1\n-- 2.txt --\ntest2\n"
  },
  {
    "path": "integration-tests/tests/path_list_output.txtar",
    "content": "# Test List[Path] output (multiple file outputs)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict writes multiple files\ncog predict $TEST_IMAGE\n\n# Verify files were created with expected content\nexec cat output.0.txt\nstdout 'foo'\nexec cat output.1.txt\nstdout 'bar'\nexec cat output.2.txt\nstdout 'baz'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom typing import List\n\nfrom cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> List[Path]:\n        predictions = [\"foo\", \"bar\", \"baz\"]\n        output = []\n        for i, prediction in enumerate(predictions):\n            out_path = Path(f\"/tmp/out-{i}.txt\")\n            with out_path.open(\"w\") as f:\n                f.write(prediction)\n            output.append(out_path)\n        return output\n"
  },
  {
    "path": "integration-tests/tests/path_output.txtar",
    "content": "# Test Path output (file output)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict writes file to output.bmp\ncog predict $TEST_IMAGE\n\n# Verify file was created and has expected size (255x255 RGB BMP = ~195KB)\nexec test -f output.bmp\nexec test -s output.bmp\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pillow==10.4.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os\nimport tempfile\n\nfrom cog import BasePredictor, Path\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> Path:\n        temp_dir = tempfile.mkdtemp()\n        temp_path = os.path.join(temp_dir, \"prediction.bmp\")\n        img = Image.new(\"RGB\", (255, 255), \"red\")\n        img.save(temp_path)\n        return Path(temp_path)\n"
  },
  {
    "path": "integration-tests/tests/predict_existing_image.txtar",
    "content": "# Test predict works with pre-built image from different directory\n# Source: test_predict.py::test_predict_runs_an_existing_image\n\n# Build the image first\ncog build -t $TEST_IMAGE\n\n# Create a different directory and run predict from there\nmkdir another_dir\ncd another_dir\n\n# Run predict on the pre-built image (no cog.yaml in current dir)\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/predict_json_file.txtar",
    "content": "# Test --json @file reads JSON from file\n# Source: test_predict.py::test_predict_json_input_filename\n\n# Build and run prediction with JSON from file\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE --json @input.json\nstdout '\"status\": \"succeeded\"'\nstdout '\"output\": \"hello sackfield\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n\n-- input.json --\n{\n    \"s\": \"sackfield\"\n}\n"
  },
  {
    "path": "integration-tests/tests/predict_json_input.txtar",
    "content": "# Test --json flag with inline JSON input\n# Source: test_predict.py::test_predict_json_input\n\n# Build and run prediction with inline JSON\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE --json '{\"s\": \"sackfield\"}'\nstdout '\"status\": \"succeeded\"'\nstdout '\"output\": \"hello sackfield\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/predict_json_output_file.txtar",
    "content": "# Test --json with --output writes JSON to file\n# Source: test_predict.py::test_predict_json_output\n\n# Build and run prediction with JSON output to file\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE --json '{\"s\": \"sackfield\"}' --output output.json\n\n# Verify file was created with correct content\nexists output.json\nexec cat output.json\nstdout '\"status\": \"succeeded\"'\nstdout '\"output\": \"hello sackfield\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/predict_json_stdin.txtar",
    "content": "# Test --json @- reads JSON from stdin\n# Source: test_predict.py::test_predict_json_input_stdin\n\n# Build and run prediction with JSON from stdin\ncog build -t $TEST_IMAGE\nstdin input.json\ncog predict $TEST_IMAGE --json @-\nstdout '\"status\": \"succeeded\"'\nstdout '\"output\": \"hello sackfield\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n\n-- input.json --\n{\"s\": \"sackfield\"}\n"
  },
  {
    "path": "integration-tests/tests/predict_json_stdin_dash.txtar",
    "content": "# Test --json - reads JSON directly from stdin\n# Source: test_predict.py::test_predict_json_input_stdin_dash\n\n# Build and run prediction with JSON from stdin using literal '-'\ncog build -t $TEST_IMAGE\nstdin input.json\ncog predict $TEST_IMAGE --json -\nstdout '\"status\": \"succeeded\"'\nstdout '\"output\": \"hello sackfield\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n\n-- input.json --\n{\"s\": \"sackfield\"}\n"
  },
  {
    "path": "integration-tests/tests/predict_many_inputs_image.txtar",
    "content": "\n# Test predict with many input types against pre-built image\n# Source: test_predict.py::test_predict_many_inputs_with_existing_image\n#\n# This test builds an image first, then runs predictions against that image\n# from a different directory (simulating using a pre-built image).\n\n# Build the image first\ncog build -t $TEST_IMAGE\n\n# Run prediction against the built image with various input types\n# Using @ syntax for file inputs\ncog predict $TEST_IMAGE -i no_default=hello -i path=@path.txt -i image=@image.jpg -i choices=foo -i int_choices=3\nstdout 'hello default 20 world jpg foo 6'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input, Path\n\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        no_default: str,\n        default_without_input: str = \"default\",\n        input_with_default: int = Input(default=10),\n        path: Path = Input(description=\"Some path\"),\n        image: Path = Input(description=\"Some path\"),\n        choices: str = Input(choices=[\"foo\", \"bar\"]),\n        int_choices: int = Input(description=\"hello\", choices=[3, 4, 5]),\n    ) -> str:\n        with path.open() as f:\n            path_contents = f.read().strip()\n        image_extension = str(image).split(\".\")[-1]\n        return (\n            no_default\n            + \" \"\n            + default_without_input\n            + \" \"\n            + str(input_with_default * 2)\n            + \" \"\n            + path_contents\n            + \" \"\n            + image_extension\n            + \" \"\n            + choices\n            + \" \"\n            + str(int_choices * 2)\n        )\n\n-- path.txt --\nworld\n-- image.jpg --\nfake image content\n"
  },
  {
    "path": "integration-tests/tests/predict_output_file.txtar",
    "content": "# Test -o flag writes Path output to specified file\n# Source: test_predict.py::test_predict_writes_files_to_files_with_custom_name\n\n# Build and run prediction with custom output filename\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE -o myoutput.bmp\n\n# Verify file exists and has non-zero size\nexists myoutput.bmp\nexec test -s myoutput.bmp\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pillow==10.4.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os\nimport tempfile\n\nfrom cog import BasePredictor, Path\nfrom PIL import Image\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> Path:\n        temp_dir = tempfile.mkdtemp()\n        temp_path = os.path.join(temp_dir, \"prediction.bmp\")\n        img = Image.new(\"RGB\", (255, 255), \"red\")\n        img.save(temp_path)\n        return Path(temp_path)\n"
  },
  {
    "path": "integration-tests/tests/predict_output_string.txtar",
    "content": "# Test -o flag writes string output to file\n# Source: test_predict.py::test_predict_writes_strings_to_files\n\n# Build image\ncog build -t $TEST_IMAGE\n\n# Run prediction with -o to write output to file\ncog predict $TEST_IMAGE -i s=world -o out.txt\n\n# Verify file contents\nexec cat out.txt\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/predict_sys_exit.txtar",
    "content": "# Test that sys.exit() in predict() fails the prediction but does NOT kill\n# the worker. A subsequent prediction should still succeed.\n#\n# sys.exit() raises SystemExit, which is caught by the PyO3 boundary.\n# The prediction fails, but the worker subprocess stays alive.\n\ncog serve\n\n# First prediction calls sys.exit(1) — should fail\ncurl POST /predictions '{\"input\":{\"do_exit\":true}}'\nstdout '\"status\":\"failed\"'\nstdout '\"error\":\".*SystemExit'\n\n# Second prediction — worker should still be alive and accept it\ncurl POST /predictions '{\"input\":{\"do_exit\":false}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output\":\"still alive\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport sys\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, do_exit: bool = False) -> str:\n        if do_exit:\n            sys.exit(1)\n        return \"still alive\"\n"
  },
  {
    "path": "integration-tests/tests/prediction_error_response.txtar",
    "content": "# Test that a runtime exception in predict() returns a well-formed error response.\n#\n# When predict() raises an exception, coglet returns HTTP 200 with\n# status \"failed\", the error message, and predict_time in metrics.\n# This test verifies the response shape — not just that it fails.\n\ncog serve\n\n# ValueError in predict()\ncurl POST /predictions '{\"input\":{\"mode\":\"value_error\"}}'\nstdout '\"status\":\"failed\"'\nstdout '\"error\":\".*this is a value error\"'\nstdout '\"predict_time\":'\n\n# RuntimeError in predict()\ncurl POST /predictions '{\"input\":{\"mode\":\"runtime_error\"}}'\nstdout '\"status\":\"failed\"'\nstdout '\"error\":\".*runtime problem\"'\nstdout '\"predict_time\":'\n\n# Generic Exception in predict()\ncurl POST /predictions '{\"input\":{\"mode\":\"generic\"}}'\nstdout '\"status\":\"failed\"'\nstdout '\"error\":\".*something went wrong\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, mode: str) -> str:\n        if mode == \"value_error\":\n            raise ValueError(\"this is a value error\")\n        elif mode == \"runtime_error\":\n            raise RuntimeError(\"runtime problem\")\n        elif mode == \"generic\":\n            raise Exception(\"something went wrong\")\n        return \"ok\"\n"
  },
  {
    "path": "integration-tests/tests/pty_echo.txtar",
    "content": "[short] skip 'slow test - requires Docker build'\n\n# Test that cog run works with PTY for simple commands\n# This is a simpler variant that just tests echo works\n\n# Run echo command with PTY (no input needed)\n# cog run builds from current directory and runs the command\npty-run /dev/null cog run echo \"hello from cog run\"\nstdout 'hello from cog run'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str = \"world\") -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/pty_interactive.txtar",
    "content": "[short] skip 'slow test - requires Docker build'\n\n# Test that cog run /bin/bash works with interactive PTY\n# This verifies bidirectional PTY interaction (send input, receive output)\n\n# Run bash with PTY input file - send commands and verify output\n# cog run builds from current directory and runs the command\npty-run pty_input.txt cog run /bin/bash\nstdout 'SENTINEL_12345'\nstdout 'HELLO_WORLD'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str = \"world\") -> str:\n        return \"hello \" + s\n\n-- pty_input.txt --\necho SENTINEL_12345\necho HELLO_WORLD\nexit\n"
  },
  {
    "path": "integration-tests/tests/pydantic2.txtar",
    "content": "# Test explicit Pydantic 2 dependency\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict with Pydantic 2\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  gpu: false\n  python_version: \"3.12\"\n  python_packages:\n    - \"pydantic>2\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/pydantic2_output.txtar",
    "content": "# Test that Pydantic v2 BaseModel works as prediction output type.\n# Coglet's make_encodeable() must call model_dump() to serialize.\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict returns structured Pydantic output\ncog predict $TEST_IMAGE -i name=alice -i score=0.95\nstdout '\"name\": \"alice\"'\nstdout '\"score\": 0.95'\nstdout '\"tags\"'\nstdout 'default'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"pydantic>2\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom typing import List\n\nfrom pydantic import BaseModel as PydanticBaseModel\n\nfrom cog import BasePredictor\n\n\nclass Result(PydanticBaseModel):\n    name: str\n    score: float\n    tags: List[str]\n\n\nclass Predictor(BasePredictor):\n    def predict(self, name: str, score: float = 0.5) -> Result:\n        return Result(name=name, score=score, tags=[\"default\"])\n"
  },
  {
    "path": "integration-tests/tests/python313.txtar",
    "content": "# Test Python 3.13 support\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict with Python 3.13\ncog predict $TEST_IMAGE -i num=5\nstdout '10'\n\n-- cog.yaml --\nbuild:\n  gpu: false\n  python_version: \"3.13\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, num: int) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/python37_deprecated.txtar",
    "content": "# Test that Python 3.7 is deprecated and build fails with appropriate error\n# Build should fail with deprecation error\n! cog build -t $TEST_IMAGE\nstderr 'invalid build.python_version \"3.7\": minimum supported Python version is 3.10'\n\n-- cog.yaml --\nbuild:\n  gpu: false\n  python_version: \"3.7\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, num: int) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/python38_deprecated.txtar",
    "content": "# Test that Python 3.8 is deprecated and build fails with appropriate error\n# Build should fail with deprecation error\n! cog build -t $TEST_IMAGE\nstderr 'invalid build.python_version \"3.8\": minimum supported Python version is 3.10'\n\n-- cog.yaml --\nbuild:\n  gpu: false\n  python_version: \"3.8\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, num: int) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/python39_deprecated.txtar",
    "content": "# Test that Python 3.9 is deprecated and build fails with appropriate error\n# Build should fail with deprecation error\n! cog build -t $TEST_IMAGE\nstderr 'invalid build.python_version \"3.9\": minimum supported Python version is 3.10'\n\n-- cog.yaml --\nbuild:\n  gpu: false\n  python_version: \"3.9\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, num: int) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/run_basic.txtar",
    "content": "# Test basic cog run functionality\n# Source: test_run.py::test_run\n\ncog run echo hello world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n"
  },
  {
    "path": "integration-tests/tests/run_stdin_cat.txtar",
    "content": "# Test stdin piped through cat and returned to stdout\n# Source: test_run.py::test_run_with_piped_stdin_returned_to_stdout\n\nstdin input.txt\ncog run cat\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.13\"\n\n-- input.txt --\nhello world\n"
  },
  {
    "path": "integration-tests/tests/run_stdin_unconsumed.txtar",
    "content": "# Test stdin handling when piped but not consumed by the command\n# Source: test_run.py::test_run_with_unconsumed_piped_stdin\n\nstdin input.txt\ncog run echo hello-from-echo\nstdout 'hello-from-echo'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.13\"\n\n-- input.txt --\nhello-from-stdin\n"
  },
  {
    "path": "integration-tests/tests/scope_context.txtar",
    "content": "# Test that per-prediction context is available via current_scope().context.\n#\n# Verifies:\n# 1. context dict from request body is accessible in the predictor\n# 2. Empty context (default) returns an empty dict\n# 3. Multiple key-value pairs are preserved\n\ncog serve\n\n# Prediction with context — predictor returns sorted key:value pairs\ncurl POST /predictions '{\"input\":{},\"context\":{\"api_token\":\"secret123\",\"region\":\"us-east-1\"}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output\":\"api_token:secret123, region:us-east-1\"'\n\n# Prediction without context — should get empty dict\ncurl POST /predictions '{\"input\":{\"expect_empty\":\"true\"}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output\":\"context_is_empty=True\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input, current_scope\n\n\nclass Predictor(BasePredictor):\n    def predict(self, expect_empty: str = Input(default=\"false\")) -> str:\n        ctx = current_scope().context\n\n        if expect_empty == \"true\":\n            return f\"context_is_empty={len(ctx) == 0}\"\n\n        # Return the context as a formatted string so we can assert on individual keys\n        return \", \".join(f\"{k}:{v}\" for k, v in sorted(ctx.items()))\n"
  },
  {
    "path": "integration-tests/tests/secrets.txtar",
    "content": "# Test that build secrets can be mounted during Docker build\n\n# Set environment variable for the env-secret\nenv ENV_SECRET=env_secret_value\n\n# Build with secrets (file-based and env-based)\ncog build -t $TEST_IMAGE --secret id=file-secret,src=file-secret.txt --secret id=env-secret,env=ENV_SECRET\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.13\"\n  run:\n    - command: >-\n        ID=\"file-secret\";\n        EXPECTED_VALUE=\"file_secret_value\";\n        EXPECTED_PATH=\"/etc/file_secret.txt\";\n        [ \"$(cat \"$EXPECTED_PATH\")\" = \"$EXPECTED_VALUE\" ] || ( echo \"Assertion failed\"; exit 1; )\n      mounts:\n        - type: secret\n          id: file-secret\n          target: /etc/file_secret.txt\n    - command: >-\n        ID=\"env-secret\";\n        EXPECTED_VALUE=\"env_secret_value\";\n        EXPECTED_PATH=\"/var/env-secret.txt\";\n        [ \"$(cat \"$EXPECTED_PATH\")\" = \"$EXPECTED_VALUE\" ] || ( echo \"Assertion failed\"; exit 1; )\n      mounts:\n        - type: secret\n          id: env-secret\n          target: /var/env-secret.txt\npredict: \"predict.py:Predictor\"\n\n-- file-secret.txt --\nfile_secret_value\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, num: int) -> int:\n        return num * 2\n"
  },
  {
    "path": "integration-tests/tests/sequential_state_leak.txtar",
    "content": "# Test that module-level state persists across sequential predictions.\n#\n# This is intentional behavior — the worker process is long-lived and\n# module-level state IS shared. This test documents that behavior:\n# a module-level list that gets appended to on each call accumulates.\n\ncog serve\n\n# First prediction — list starts empty, appends \"a\"\n# Note: \".\" in regexes below matches Python's single-quote characters\n# (testscript single-quoted args cannot embed literal single quotes)\ncurl POST /predictions '{\"input\":{\"item\":\"a\"}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output\":\"\\[.a.\\]\"'\n\n# Second prediction — list now has [\"a\"], appends \"b\"\ncurl POST /predictions '{\"input\":{\"item\":\"b\"}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output\":\"\\[.a., .b.\\]\"'\n\n# Third prediction — list now has [\"a\", \"b\"], appends \"c\"\ncurl POST /predictions '{\"input\":{\"item\":\"c\"}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output\":\"\\[.a., .b., .c.\\]\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n# Module-level state — persists across predictions in the same worker\n_state: list = []\n\n\nclass Predictor(BasePredictor):\n    def predict(self, item: str) -> str:\n        _state.append(item)\n        return str(_state)\n"
  },
  {
    "path": "integration-tests/tests/setup_slow_serial.txtar",
    "content": "[short] skip 'slow test - skip in short mode'\n\n# Test that a slow setup() completes successfully without being killed by\n# an internal timeout. Coglet should not impose its own setup timeout —\n# the external orchestrator (director) is the authority on setup timeouts.\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server — setup takes ~15s but should complete fine\ncog serve\n\n# Verify the server is healthy and predictions work\ncurl POST /predictions '{\"input\":{\"s\":\"hello\"}}'\nstdout '\"output\":\"hello hello\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport time\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def setup(self) -> None:\n        print(\"Starting slow setup...\")\n        time.sleep(15)\n        print(\"Slow setup complete.\")\n\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/setup_subprocess_double_fork.txtar",
    "content": "[short] skip 'slow test - skip in short mode'\n\n# Test double fork subprocess spawned during setup\n# This ensures stream redirection works correctly with daemonized processes\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Make predictions that communicate via file system with forked process\ncurl POST /predictions '{\"input\":{\"s\":\"friendo1\"}}'\nstdout '\"output\":\"hello friendo1\"'\n\ncurl POST /predictions '{\"input\":{\"s\":\"friendo2\"}}'\nstdout '\"output\":\"hello friendo2\"'\n\ncurl POST /predictions '{\"input\":{\"s\":\"friendo3\"}}'\nstdout '\"output\":\"hello friendo3\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport os.path\nimport signal\nimport subprocess\nimport sys\nimport time\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    \"\"\"\n    This predictor checks the case where a process is spawned during setup and then each\n    prediction depends on being able to communicate with that process. In the event that\n    stream redirection is not working correctly, the forked process will not be able to\n    write to stdout/stderr and will likely exit. Any state other than \"running\" is\n    considered an error condition and raises SystemExit to interrupt any more prediction\n    serving.\n\n    This variant runs a forked python process via a shell wrapper to which a \"message\" is\n    sent via file for each call to `predict`.\n    \"\"\"\n\n    def setup(self) -> None:\n        print(\"---> starting background process\")\n\n        self.bg = subprocess.Popen([\"bash\", \"run-forker.sh\"])\n\n        print(f\"---> started background process pid={self.bg.pid}\")\n\n    def predict(self, s: str) -> str:\n        status = self.bg.poll()\n\n        print(f\"---> background job status={status}\")\n\n        if status is not None:\n            raise SystemExit\n\n        print(f\"---> sending message to background job pid={self.bg.pid}\")\n\n        with open(\".inbox\", \"w\") as inbox:\n            inbox.write(s)\n\n        print(f\"---> sent message to background job pid={self.bg.pid}\")\n\n        now = time.time()\n\n        print(f\"---> waiting for outbox message from background job pid={self.bg.pid}\")\n\n        while not os.path.exists(\".outbox\"):\n            if time.time() - now > 5:\n                raise TimeoutError\n\n            time.sleep(0.01)\n\n        try:\n            with open(\".outbox\", \"r\") as outbox:\n                print(f\"---> relaying message from background job pid={self.bg.pid}\")\n\n                return outbox.read()\n\n        finally:\n            os.unlink(\".outbox\")\n\n-- run-forker.sh --\n#!/usr/bin/env bash\npython ./forker.py &\nwait\n\n-- forker.py --\nimport os\nimport signal\nimport time\n\n\ndef main():\n    child_pid = os.fork()\n    is_child = child_pid == 0\n\n    pid = os.getpid()\n    was_pinged = False\n\n    while True:\n        if os.path.exists(\".inbox\") and is_child:\n            s = \"\"\n\n            with open(\".inbox\", \"r\") as inbox:\n                print(f\"---> CHILD ({pid}) reading request\")\n\n                s = inbox.read()\n\n            os.unlink(\".inbox\")\n\n            with open(\".outbox\", \"w\") as outbox:\n                print(f\"---> CHILD ({pid}) sending response\")\n\n                outbox.write(\"hello \" + s)\n\n        if time.time() % 10 == 0:\n            if is_child:\n                print(f\"---> CHILD ({pid}) \" + (\"here \" * 20))\n            else:\n                print(f\"===> PARENT ({pid})\")\n\n        time.sleep(0.01)\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "integration-tests/tests/setup_subprocess_double_fork_http.txtar",
    "content": "[short] skip 'slow test - skip in short mode'\n\n# Test double fork subprocess with HTTP server spawned during setup\n# This ensures stream redirection works correctly with daemonized HTTP servers\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Make predictions that communicate with forked HTTP server\ncurl POST /predictions '{\"input\":{\"s\":\"friendo1\"}}'\nstdout '\"output\":\"hello friendo1\"'\n\ncurl POST /predictions '{\"input\":{\"s\":\"friendo2\"}}'\nstdout '\"output\":\"hello friendo2\"'\n\ncurl POST /predictions '{\"input\":{\"s\":\"friendo3\"}}'\nstdout '\"output\":\"hello friendo3\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n  - requests\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport signal\nimport subprocess\nimport sys\n\nimport requests\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    \"\"\"\n    This predictor checks the case where a process is spawned during setup and then each\n    prediction depends on being able to communicate with that process. In the event that\n    stream redirection is not working correctly, the forked process will not be able to\n    write to stdout/stderr and will likely exit. Any state other than \"running\" is\n    considered an error condition and raises SystemExit to interrupt any more prediction\n    serving.\n\n    This variant runs a forked python HTTP server via a shell wrapper to which a request\n    is made during each call to `predict`.\n    \"\"\"\n\n    def setup(self) -> None:\n        print(\"---> starting background process\")\n\n        self.bg = subprocess.Popen([\"bash\", \"run-pong.sh\"])\n\n        print(f\"---> started background process pid={self.bg.pid}\")\n\n        # Wait for HTTP server to be ready\n        import time\n        for i in range(30):\n            try:\n                requests.get(\"http://127.0.0.1:7777/ping\", timeout=1)\n                print(\"---> background HTTP server is ready\")\n                break\n            except Exception:\n                print(f\"---> waiting for HTTP server ({i+1}/30)\")\n                time.sleep(0.5)\n        else:\n            raise RuntimeError(\"Background HTTP server failed to start\")\n\n    def predict(self, s: str) -> str:\n        status = self.bg.poll()\n\n        print(f\"---> background job status={status}\")\n\n        if status is None:\n            print(f\"---> sending request to background job pid={self.bg.pid}\")\n\n            print(requests.get(\"http://127.0.0.1:7777/ping\"))\n\n            print(f\"---> sent request to background job pid={self.bg.pid}\")\n        else:\n            raise SystemExit\n\n        return \"hello \" + s\n\n-- run-pong.sh --\n#!/usr/bin/env bash\npython ./pong.py &\nwait\n\n-- pong.py --\nimport os\nimport signal\nimport time\nfrom random import randint\nfrom wsgiref.simple_server import make_server\n\n\ndef main():\n    child_pid = os.fork()\n    is_child = child_pid == 0\n\n    pid = os.getpid()\n\n    if is_child:\n        make_server(\"127.0.0.1\", 7777, app).serve_forever()\n    else:\n        while True:\n            print(f\"===> PARENT ({pid})\")\n\n            time.sleep(10)\n\n\ndef app(environ, start_response):\n    print(f\"---> CHILD ({os.getpid()})\")\n\n    if environ[\"PATH_INFO\"] == \"/ping\":\n        start_response(\"200 OK\", [(\"content-type\", \"text/plain\")])\n        return [b\"PONG\\n\" for n in range(100 + randint(2, 32))]\n\n    start_response(\"404 Not Found\", [(\"content-type\", \"text/plain\")])\n    return [b\"NO\\n\"]\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "integration-tests/tests/setup_subprocess_multiprocessing.txtar",
    "content": "[short] skip 'slow test - skip in short mode'\n\n# Test multiprocessing.Process spawned during setup\n# This ensures stream redirection works correctly with Python multiprocessing\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Make a prediction that communicates via multiprocessing.Pipe\n# Note: The background process closes the connection after first use,\n# so we only test one prediction\ncurl POST /predictions '{\"input\":{\"s\":\"friendo1\"}}'\nstdout '\"output\":'\nstdout '\"status\":\"succeeded\"'\n# Check the logs show the background process communication worked\nstdout 'sending ping to background job'\nstdout 'received .* from background job'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.10\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport atexit\nimport multiprocessing\nimport pathlib\nimport signal\nimport subprocess\nimport sys\nimport time\n\nfrom cog.types import Path\nfrom cog import BasePredictor\n\nfrom bg import ponger\n\ndef cleanup():\n    for tmp in pathlib.Path(\"./\").glob(\"*.tmp\"):\n        if tmp.is_file():\n            tmp.unlink(missing_ok=True)\n\n\natexit.register(cleanup)\n\n\nclass Predictor(BasePredictor):\n    \"\"\"\n    This predictor checks the case where a process is spawned during setup via\n    multiprocessing and then each prediction causes that process to write to stdout.\n    \"\"\"\n\n    def setup(self) -> None:\n        print(\"---> starting background process\")\n\n        cleanup()\n\n        self.parent_conn, self.child_conn = multiprocessing.Pipe()\n        self.lock = multiprocessing.Lock()\n        self.bg = multiprocessing.Process(\n            target=ponger, args=(self.child_conn, self.lock)\n        )\n        self.bg.start()\n\n        print(f\"---> started background process pid={self.bg.pid}\")\n\n    def predict(self, s: str) -> Path:\n        if self.bg.is_alive():\n            print(f\"---> sending ping to background job pid={self.bg.pid}\")\n\n            self.child_conn.send(\"ping\")\n\n            print(f\"---> sent ping to background job pid={self.bg.pid}\")\n\n            pong = self.parent_conn.recv()\n\n            print(f\"---> received {pong} from background job pid={self.bg.pid}\")\n        else:\n            print(f\"---> background job died\")\n\n            raise SystemExit\n\n        out = Path(f\"cog-test-integration-out.{time.time_ns()}.tmp\")\n        out.write_text(\"hello \" + s)\n\n        print(f\"---> wrote output file {out}\")\n\n        return out\n\n-- bg.py --\nimport multiprocessing.connection\nimport multiprocessing.synchronize\nimport os\nimport time\n\n\ndef ponger(\n    conn: multiprocessing.connection.Connection, lock: multiprocessing.synchronize.Lock\n):\n    for i in range(100):\n        print(f\"Getting ready for some serious ponginggg ({i+1}%)\")\n        time.sleep(0.001 + (0.001 * (i + 1)))\n\n    print(\"ITS PONGIN TIME\")\n\n    pid = os.getpid()\n\n    while True:\n        try:\n            ping = conn.recv()\n            print(f\"received {ping} in {pid}\")\n\n            with lock:\n                print(f\"ponging from {pid}\")\n\n                conn.send(\"pong\")\n                conn.close()\n\n        except EOFError:\n            pass\n"
  },
  {
    "path": "integration-tests/tests/setup_subprocess_simple.txtar",
    "content": "[short] skip 'slow test - skip in short mode'\n\n# Test subprocess spawned during setup that writes to stdout\n# This ensures stream redirection works correctly when a background process\n# writes output during prediction serving.\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Make predictions - the subprocess writes to stdout when it receives SIGUSR1\ncurl POST /predictions '{\"input\":{\"s\":\"friendo1\"}}'\nstdout '\"output\":\"hello friendo1\"'\n\ncurl POST /predictions '{\"input\":{\"s\":\"friendo2\"}}'\nstdout '\"output\":\"hello friendo2\"'\n\ncurl POST /predictions '{\"input\":{\"s\":\"friendo3\"}}'\nstdout '\"output\":\"hello friendo3\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport signal\nimport subprocess\nimport sys\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    \"\"\"\n    This predictor checks the case where a process is spawned during setup and then each\n    prediction causes that process to write to stdout. In the event that stream\n    redirection is not working correctly, the forked process will not be able to write to\n    stdout/stderr and will likely exit. Any state other than \"running\" is considered an\n    error condition and raises SystemExit to interrupt any more prediction serving.\n\n    This variant runs a simple subprocess to which SIGUSR1 is sent during each call to\n    `predict`.\n    \"\"\"\n\n    def setup(self) -> None:\n        print(\"---> starting background process\")\n\n        self.bg = subprocess.Popen([\"bash\", \"child.sh\"])\n\n        print(f\"---> started background process pid={self.bg.pid}\")\n\n    def predict(self, s: str) -> str:\n        status = self.bg.poll()\n\n        if status is None:\n            print(f\"---> sending signal to background job pid={self.bg.pid}\")\n\n            self.bg.send_signal(signal.SIGUSR1)\n\n            print(f\"---> sent signal to background job pid={self.bg.pid}\")\n        else:\n            print(f\"---> background job died status={status}\")\n\n            raise SystemExit\n\n        return \"hello \" + s\n\n-- child.sh --\n#!/usr/bin/env bash\nset -euo pipefail\n\n# This _pong function and associated trap ensures that any SIGUSR1 sent during `predict`\n# will cause this process to write a decent amount of text to stdout. In the event that\n# stream redirection is not working correctly, this process will likely be in a defunct\n# state before the first SIGUSR1 can be sent.\n_pong() {\n  for i in $(seq 100); do\n    echo \"${0} (${$}) PONG (${i}/100)\"\n  done\n}\n\ntrap _pong USR1\n\n# This loop simulates a setup period for filling up any stdout buffer.\nfor i in $(seq 100); do\n  echo \"${0} ($$) SETTING UP (${i}/100)\"\n  sleep 0.01\ndone\n\n# This loop simulates periodic writes to stdout while the background process is running\n# for the purpose of ensuring the file descriptor is still usable.\nwhile true; do\n  now=\"$(date +%s)\"\n  now_mod=$((now % 10))\n\n  if [[ \"${now_mod}\" == 0 ]]; then\n    echo \"${0} (${$}) STILL HERE\"\n    sleep 1\n  fi\n\n  sleep 0.1\ndone\n"
  },
  {
    "path": "integration-tests/tests/setup_timeout_serial.txtar",
    "content": "[short] skip 'slow test - skip in short mode'\n\n# Test that COG_SETUP_TIMEOUT env var is respected. When set to a value\n# shorter than the actual setup duration, setup should fail with SETUP_FAILED.\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server — setup takes ~15s but timeout is 10s, so it should fail\n! cog serve\n\n# Verify the server reports SETUP_FAILED\ncurl GET /health-check\nstdout 'SETUP_FAILED'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\nenvironment:\n  - COG_SETUP_TIMEOUT=10\n\n-- predict.py --\nimport time\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def setup(self) -> None:\n        print(\"Starting slow setup...\")\n        time.sleep(15)\n        print(\"Slow setup complete.\")\n\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/setup_worker_tracing_logs.txtar",
    "content": "[short] skip 'slow test - skip in short mode'\n\n# Test worker tracing logs appear in setuplog\n# This ensures Rust tracing from the worker subprocess is properly captured during setup\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Start the server\ncog serve\n\n# Check setuplog contains orchestrator and worker tracing logs\ncurl GET /health-check\nstdout '\"status\":\"READY\"'\nstdout 'Spawning worker subprocess'\nstdout 'File descriptor redirection complete'\nstdout 'Connected to slot transport'\nstdout 'Server ready'\n\n# Verify logs after accumulation stopped are NOT included\n! stdout 'Setup complete, now accepting requests'\n\n# Make a prediction to verify it works\ncurl POST /predictions '{\"input\":{\"s\":\"test\"}}'\nstdout '\"output\":\"hello test\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n  - requests\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nimport signal\nimport subprocess\nimport sys\n\nimport requests\n\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    \"\"\"\n    This predictor spawns a double-forked HTTP server during setup to test\n    that worker tracing logs are properly captured in setuplog.\n    \"\"\"\n\n    def setup(self) -> None:\n        print(\"---> starting background process\")\n\n        self.bg = subprocess.Popen([\"bash\", \"run-pong.sh\"])\n\n        print(f\"---> started background process pid={self.bg.pid}\")\n\n        # Wait for HTTP server to be ready\n        import time\n        for i in range(30):\n            try:\n                requests.get(\"http://127.0.0.1:7777/ping\", timeout=1)\n                print(\"---> background HTTP server is ready\")\n                break\n            except Exception:\n                print(f\"---> waiting for HTTP server ({i+1}/30)\")\n                time.sleep(0.5)\n        else:\n            raise RuntimeError(\"Background HTTP server failed to start\")\n\n    def predict(self, s: str) -> str:\n        status = self.bg.poll()\n\n        if status is None:\n            requests.get(\"http://127.0.0.1:7777/ping\")\n        else:\n            raise SystemExit\n\n        return \"hello \" + s\n\n-- run-pong.sh --\n#!/usr/bin/env bash\npython ./pong.py &\nwait\n\n-- pong.py --\nimport os\nimport signal\nimport time\nfrom random import randint\nfrom wsgiref.simple_server import make_server\n\n\ndef main():\n    child_pid = os.fork()\n    is_child = child_pid == 0\n\n    pid = os.getpid()\n\n    if is_child:\n        make_server(\"127.0.0.1\", 7777, app).serve_forever()\n    else:\n        while True:\n            time.sleep(10)\n\n\ndef app(environ, start_response):\n    if environ[\"PATH_INFO\"] == \"/ping\":\n        start_response(\"200 OK\", [(\"content-type\", \"text/plain\")])\n        return [b\"PONG\\n\" for n in range(100 + randint(2, 32))]\n\n    start_response(\"404 Not Found\", [(\"content-type\", \"text/plain\")])\n    return [b\"NO\\n\"]\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "integration-tests/tests/static_schema_fallback.txtar",
    "content": "# Test that when static schema generation is opted in but encounters an\n# unresolvable output type, the build falls back to legacy runtime schema\n# generation instead of failing.\n#\n# The predictor returns a BaseModel subclass imported from a local package\n# (mypackage/__init__.py). The static parser's moduleToFilePath() converts\n# \"mypackage\" to \"mypackage.py\" instead of \"mypackage/__init__.py\", so the\n# file isn't found and the name stays unresolved (ErrUnresolvableType).\n# The legacy Python inspector imports modules normally via Python's import\n# system, which handles __init__.py packages correctly, so schema generation\n# succeeds at runtime.\n\n# Opt in to static schema generation\nenv COG_STATIC_SCHEMA=1\n\n# Build should succeed — static fails, legacy takes over\ncog build -t $TEST_IMAGE\nstderr 'Static schema generation failed'\nstderr 'Falling back to legacy runtime schema generation'\n\n# Predict should still work end-to-end via the legacy schema\ncog predict $TEST_IMAGE -i prompt=hello\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- mypackage/__init__.py --\nfrom cog import BaseModel\n\n\nclass Output(BaseModel):\n    text: str\n\n-- predict.py --\nfrom cog import BasePredictor\nfrom mypackage import Output\n\n\nclass Predictor(BasePredictor):\n    def predict(self, prompt: str) -> Output:\n        return Output(text=prompt + \" world\")\n"
  },
  {
    "path": "integration-tests/tests/static_schema_gen.txtar",
    "content": "# Test that the static schema generator (Go tree-sitter) produces a correct\n# OpenAPI schema that is embedded in the Docker image label and served by coglet.\n#\n# This exercises the full pipeline:\n#   1. cog build runs the Go parser on predict.py\n#   2. Schema is written to .cog/openapi_schema.json inside the image\n#   3. Schema is embedded as the run.cog.openapi_schema Docker label\n#   4. coglet loads the schema from disk and serves it at /openapi.json\n\n# Opt in to static schema generation\nenv COG_STATIC_SCHEMA=1\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Verify schema is in Docker label with correct structure\nexec docker inspect $TEST_IMAGE --format '{{index .Config.Labels \"run.cog.openapi_schema\"}}'\n\n# Top-level OpenAPI structure\nstdout '\"openapi\":\"3.0.2\"'\n\n# Input schema: required string field\nstdout '\"text\":'\nstdout '\"type\":\"string\"'\n\n# Input schema: optional int with default\nstdout '\"count\":'\nstdout '\"type\":\"integer\"'\nstdout '\"default\":5'\n\n# Input schema: choices generate enum\nstdout '\"enum\":\\[\"fast\",\"balanced\",\"quality\"\\]'\n\n# Input schema: Path type → uri format\nstdout '\"format\":\"uri\"'\n\n# Input schema: required array\nstdout '\"required\":\\[\"text\",\"image\"\\]'\n\n# Output type\nstdout '\"type\":\"string\"'\n\n# Predict should work end-to-end\ncog predict $TEST_IMAGE -i text=hello -i image=@test.jpg\nstdout 'hello-5-fast-jpg'\n\n# Prediction with overrides\ncog predict $TEST_IMAGE -i text=world -i count=3 -i mode=quality -i image=@test.jpg\nstdout 'world-3-quality-jpg'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input, Path\n\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        text: str,\n        image: Path,\n        count: int = Input(description=\"Number of iterations\", default=5, ge=1, le=100),\n        mode: str = Input(description=\"Quality mode\", default=\"fast\", choices=[\"fast\", \"balanced\", \"quality\"]),\n    ) -> str:\n        ext = str(image).split(\".\")[-1]\n        return f\"{text}-{count}-{mode}-{ext}\"\n\n-- test.jpg --\nfake image content\n"
  },
  {
    "path": "integration-tests/tests/string_list_input.txtar",
    "content": "# Test list[str] input type\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict with list input\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: list[str] = Input(description=\"A list of strings to print.\")) -> str:\n        return \"hello \" + \"|\".join(s)\n"
  },
  {
    "path": "integration-tests/tests/string_none_output.txtar",
    "content": "\n# Test str return type that returns None\n# This tests the handling of None values for typed outputs\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict returns None despite str type annotation\n# When None is returned, cog shows \"No output generated\" on stderr\ncog predict $TEST_IMAGE\nstderr 'No output generated'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        return None\n"
  },
  {
    "path": "integration-tests/tests/string_predictor.txtar",
    "content": "\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Basic prediction works\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n# Missing required input fails\n! cog predict $TEST_IMAGE -i wrong=value\nstderr 's: Field required'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/subdirectory_predictor.txtar",
    "content": "\n# Test predictor in subdirectory with imports from parent\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict using predictor in subdirectory\ncog predict $TEST_IMAGE -i s=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"my-subdir/predict.py:Predictor\"\n\n-- mylib.py --\ndef concat(a, b):\n    return a + \" \" + b\n\n-- my-subdir/predict.py --\nfrom cog import BasePredictor\nfrom mylib import concat\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return concat(\"hello\", s)\n"
  },
  {
    "path": "integration-tests/tests/tensorflow.txtar",
    "content": "[short] skip 'slow test - run without -short flag'\n\n# Test TensorFlow build and predict\ncog build -t $TEST_IMAGE\ncog predict $TEST_IMAGE\nstdout '2.11.1'\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  cuda: \"11.8\"\n  python_version: \"3.10\"\n  system_packages:\n    - \"libgl1-mesa-glx\"\n    - \"libglib2.0-0\"\n    - \"xvfb\"\n  python_requirements: \"requirements.txt\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nimport tensorflow\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        return tensorflow.__version__\n\n-- requirements.txt --\ncompel==2.0.3\ndiffusers>=0.27.1\ngputil==1.4.0\nloguru==0.7.2\nopencv-python>=4.9.0.80\npillow>=10.2.0\npsutil==6.1.1\nreplicate>=1.0.4\nsentry-sdk[fastapi,loguru]>=2.16.0\nantialiased_cnns==0.3\nbeautifulsoup4==4.13.4\nimageio==2.37.0\nipdb==0.13.13\nkornia==0.8.1\nmatplotlib==3.10.3\nnumpy==1.23.5\nopencv_python==4.11.0.86\nPillow==11.2.1\npytorch_lightning==2.3.3\nPyYAML==6.0.2\nRequests==2.32.3\nscipy==1.15.3\nscikit-image==0.24.0\ntensorflow==2.11.1\ntensorlayer==2.2.5\ntf_slim==1.1.0\ntimm==1.0.15\ntorch==2.0.1\ntorchvision==0.15.2\ntqdm==4.67.1\n"
  },
  {
    "path": "integration-tests/tests/torch_270_cuda_126.txtar",
    "content": "[short] skip 'slow test - run without -short flag'\n\n# Test Torch 2.7.0 + CUDA 12.6 base image build\ncog build -t $TEST_IMAGE\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  cuda: \"12.6\"\n  python_version: \"3.11\"\n  python_packages:\n    - \"torch==2.7.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/torch_271_cuda_128.txtar",
    "content": "[short] skip 'slow test - run without -short flag'\n\n# Test Torch 2.7.1 + CUDA 12.8 base image build\ncog build -t $TEST_IMAGE\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  cuda: \"12.8\"\n  python_version: \"3.12\"\n  python_packages:\n    - \"torch==2.7.1+cu128\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/torch_baseimage_fallback.txtar",
    "content": "# Test that cog falls back to a CUDA base image (not cog-base) when the\n# torch version is below the minimum supported for cog-base images.\n# Torch 1.13.0 < MinimumTorchVersion (1.13.1), so no cog-base match exists.\n#\n# Uses a local registry (empty, no images seeded) to ensure the base image\n# lookup fails deterministically, and cog debug to verify the generated\n# Dockerfile without doing a real Docker build.\n\n# Start an empty local registry so the cog-base lookup fails deterministically\nregistry-start\nenv COG_REGISTRY_HOST=$TEST_REGISTRY\n\n# Generate Dockerfile — should fall back to nvidia/cuda, not cog-base\ncog debug\nstdout 'FROM nvidia/cuda'\n! stdout 'cog-base'\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  python_version: \"3.10\"\n  python_packages:\n    - \"torch==1.13.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/torch_baseimage_no_cog_base.txtar",
    "content": "[short] skip 'slow test - run without -short flag'\nskip 'temporarily disabled - takes ~10min, blocking CI'\n\n# Test Torch 1.13.0 base image with --use-cog-base-image=false\ncog build -t $TEST_IMAGE --openapi-schema openapi.json --use-cog-base-image=false\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  python_version: \"3.10\"\n  python_packages:\n    - \"torch==1.13.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n\n-- openapi.json --\n{\n    \"info\": {\n        \"title\": \"Cog\",\n        \"version\": \"0.1.0\"\n    },\n    \"paths\": {\n        \"/\": {\n            \"get\": {\n                \"summary\": \"Root\",\n                \"responses\": {\n                    \"200\": {\n                        \"content\": {\n                            \"application/json\": {\n                                \"schema\": {\n                                    \"title\": \"Response Root  Get\"\n                                }\n                            }\n                        },\n                        \"description\": \"Successful Response\"\n                    }\n                },\n                \"operationId\": \"root__get\"\n            }\n        },\n        \"/predictions\": {\n            \"post\": {\n                \"summary\": \"Predict\",\n                \"responses\": {\n                    \"200\": {\n                        \"content\": {\n                            \"application/json\": {\n                                \"schema\": {\n                                    \"$ref\": \"#/components/schemas/PredictionResponse\"\n                                }\n                            }\n                        },\n                        \"description\": \"Successful Response\"\n                    }\n                },\n                \"description\": \"Run a single prediction on the model\",\n                \"operationId\": \"predict_predictions_post\"\n            }\n        }\n    },\n    \"openapi\": \"3.1.0\",\n    \"components\": {\n        \"schemas\": {\n            \"Input\": {\n                \"type\": \"object\",\n                \"title\": \"Input\",\n                \"properties\": {\n                    \"s\": {\n                        \"type\": \"string\",\n                        \"title\": \"S\"\n                    }\n                }\n            },\n            \"Output\": {\n                \"type\": \"string\",\n                \"title\": \"Output\"\n            },\n            \"PredictionResponse\": {\n                \"type\": \"object\",\n                \"title\": \"PredictionResponse\",\n                \"properties\": {\n                    \"output\": {\n                        \"$ref\": \"#/components/schemas/Output\"\n                    }\n                }\n            }\n        }\n    }\n}\n"
  },
  {
    "path": "integration-tests/tests/torch_baseimage_precompile.txtar",
    "content": "[short] skip 'slow test - run without -short flag'\nskip 'temporarily disabled - takes ~10min, blocking CI'\n\n# Test Torch 1.13.0 base image with --precompile flag\ncog build -t $TEST_IMAGE --openapi-schema openapi.json --use-cog-base-image=false --precompile\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  python_version: \"3.10\"\n  python_packages:\n    - \"torch==1.13.0\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n\n-- openapi.json --\n{\n    \"info\": {\n        \"title\": \"Cog\",\n        \"version\": \"0.1.0\"\n    },\n    \"paths\": {\n        \"/\": {\n            \"get\": {\n                \"summary\": \"Root\",\n                \"responses\": {\n                    \"200\": {\n                        \"content\": {\n                            \"application/json\": {\n                                \"schema\": {\n                                    \"title\": \"Response Root  Get\"\n                                }\n                            }\n                        },\n                        \"description\": \"Successful Response\"\n                    }\n                },\n                \"operationId\": \"root__get\"\n            }\n        },\n        \"/predictions\": {\n            \"post\": {\n                \"summary\": \"Predict\",\n                \"responses\": {\n                    \"200\": {\n                        \"content\": {\n                            \"application/json\": {\n                                \"schema\": {\n                                    \"$ref\": \"#/components/schemas/PredictionResponse\"\n                                }\n                            }\n                        },\n                        \"description\": \"Successful Response\"\n                    }\n                },\n                \"description\": \"Run a single prediction on the model\",\n                \"operationId\": \"predict_predictions_post\"\n            }\n        }\n    },\n    \"openapi\": \"3.1.0\",\n    \"components\": {\n        \"schemas\": {\n            \"Input\": {\n                \"type\": \"object\",\n                \"title\": \"Input\",\n                \"properties\": {\n                    \"s\": {\n                        \"type\": \"string\",\n                        \"title\": \"S\"\n                    }\n                }\n            },\n            \"Output\": {\n                \"type\": \"string\",\n                \"title\": \"Output\"\n            },\n            \"PredictionResponse\": {\n                \"type\": \"object\",\n                \"title\": \"PredictionResponse\",\n                \"properties\": {\n                    \"output\": {\n                        \"$ref\": \"#/components/schemas/Output\"\n                    }\n                }\n            }\n        }\n    }\n}\n"
  },
  {
    "path": "integration-tests/tests/torch_cuda_baseimage.txtar",
    "content": "# Test that cog correctly resolves a cog-base image for Torch 2.0.1+cu118\n# and generates the expected Dockerfile with --use-cog-base-image.\n#\n# Uses a local test registry seeded with a dummy image to avoid depending\n# on the live r8.im registry, which can be flaky in CI.\n\n# Start local registry and seed a dummy cog-base image\nregistry-start\nregistry-seed alpine:latest cog-base:cuda11.8-python3.10-torch2.0.1\n\n# Generate Dockerfile using the local registry as the cog-base source\nenv COG_REGISTRY_HOST=$TEST_REGISTRY\ncog debug --use-cog-base-image\nstdout 'FROM.*cog-base:cuda11.8-python3.10-torch2.0.1'\n\n-- cog.yaml --\nbuild:\n  gpu: true\n  python_version: \"3.10\"\n  python_packages:\n    - \"torch==2.0.1+cu118\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/train_basic.txtar",
    "content": "skip 'cog train requires static schema gen which is gated behind COG_STATIC_SCHEMA=1'\n\n# Test basic training functionality\n\n# Train with input (no pre-built image, runs from cog.yaml)\ncog train -i n=42\n\n# Verify weights file was created with correct size\nexec test -f weights.bin\nexec sh -c 'wc -c < weights.bin'\nstdout '42'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\ntrain: \"train.py:train\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n\n-- train.py --\nfrom cog import BaseModel, Input, Path\n\n\nclass TrainingOutput(BaseModel):\n    weights: Path\n\n\ndef train(\n    n: int,\n) -> TrainingOutput:\n    with open(\"weights.bin\", \"w\") as fh:\n        for _ in range(n):\n            fh.write(\"a\")\n\n    return TrainingOutput(\n        weights=Path(\"weights.bin\"),\n    )\n"
  },
  {
    "path": "integration-tests/tests/train_deprecated.txtar",
    "content": "# Test that the train command shows a deprecation warning\n\n# Train command should fail (no cog.yaml) but still show deprecation warning\n! cog train\nstderr 'Command \"train\" is deprecated'\nstderr 'will be removed in a future version of Cog'\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/training_setup.txtar",
    "content": "skip 'cog train requires static schema gen which is gated behind COG_STATIC_SCHEMA=1'\n\n# Test that training with setup method works correctly\n\n# Train with input\ncog train -i s=world\nstderr 'Trainer is setting up.'\n\n# Verify weights file was created with correct content\nexec cat weights\nstdout 'hello train world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\ntrain: \"train.py:Trainer\"\npredict: \"predict.py:Predictor\"\n\n-- train.py --\nfrom cog import BasePredictor\n\nclass Trainer(BasePredictor):\n    def setup(self) -> None:\n        print(\"Trainer is setting up.\")\n\n    def train(self, s: str) -> str:\n        print(\"Trainer.train called.\")\n        return \"hello train \" + s\n\n    def predict(self, s: str) -> str:\n        print(\"Trainer.predict called.\")\n        return \"hello predict \" + s\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/union_type.txtar",
    "content": "# Test new-style union type annotation (str | None)\n\n# Build the image\ncog build -t $TEST_IMAGE\n\n# Predict with input\ncog predict $TEST_IMAGE -i text=world\nstdout 'hello world'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor, Input\n\n\nclass Predictor(BasePredictor):\n    def setup(self):\n        self.prefix = \"hello\"\n\n    def predict(\n        self,\n        text: str | None = Input(\n            description=\"Text to prefix with 'hello '\", default=None\n        ),\n    ) -> str:\n        return self.prefix + \" \" + text\n"
  },
  {
    "path": "integration-tests/tests/webhook_delivery_failure.txtar",
    "content": "# Test that webhook delivery failure does not crash the server.\n#\n# When the webhook URL is unreachable, coglet retries with backoff\n# but eventually gives up. The server should remain healthy.\n#\n# --upload-url is set to a dummy so cog serve adds host networking.\n\ncog serve --upload-url http://unused/\n\n# Async prediction with a bogus webhook URL — delivery will fail\ncurl -H Prefer:respond-async POST /predictions '{\"id\":\"webhook-fail-test\",\"webhook\":\"http://host.docker.internal:1/nonexistent\",\"webhook_events_filter\":[\"completed\"]}'\n\n# Server should remain healthy while webhook delivery retries/fails.\n# Poll repeatedly instead of sleeping a fixed duration.\nexec bash -c 'for i in {1..10}; do curl -sf $SERVER_URL/health-check | grep -q \"\\\"status\\\":\\\"READY\\\"\" || exit 1; sleep 0.5; done'\n\n# A new sync prediction should still work\ncurl POST /predictions '{\"input\":{}}'\nstdout '\"status\":\"succeeded\"'\nstdout '\"output\":\"ok\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        return \"ok\"\n"
  },
  {
    "path": "integration-tests/tests/webhook_prediction_error.txtar",
    "content": "# Test that a failed prediction delivers the correct webhook payload.\n#\n# When predict() raises an exception during an async prediction, the webhook\n# should receive status \"failed\" with the error message populated.\n\nwebhook-server-start\ncog serve --upload-url http://unused/\n\n# Async prediction that raises an exception\ncurl -H Prefer:respond-async POST /predictions '{\"id\":\"webhook-error-test\",\"webhook\":\"$WEBHOOK_URL\",\"webhook_events_filter\":[\"completed\"]}'\n\nwebhook-server-wait\n\n# Webhook should report failure with error\nstdout '\"status\":\"failed\"'\nstdout '\"has_error\":true'\nstdout '\"error_message\":\".*prediction went wrong\"'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        raise RuntimeError(\"prediction went wrong\")\n"
  },
  {
    "path": "integration-tests/tests/weights_build.txtar",
    "content": "# Test that cog weights build generates weights.lock\n\n# Build should fail without weights section\n! cog weights build\nstderr 'no weights defined'\n\n# Add weights section and create weight file\ncp cog-with-weights.yaml cog.yaml\nmkdir models\nexec sh -c 'echo \"test model content\" > models/model.bin'\n\n# Build weights.lock\ncog weights build\nstderr 'Generated weights.lock'\nstderr '1 file'\n\n# Verify weights.lock was created\nexists weights.lock\n\n# Verify weights.lock contains expected content\nexec grep -q '\"name\": \"model\"' weights.lock\nexec grep -q '\"dest\": \"/cache/model.bin\"' weights.lock\nexec grep -q '\"digestOriginal\": \"sha256:' weights.lock\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- cog-with-weights.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\nweights:\n  - name: model\n    source: models/model.bin\n    target: /cache/model.bin\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/weights_push_inspect.txtar",
    "content": "# Test weights push and inspect lifecycle against a local registry.\n# Verifies: cog weights build -> cog weights push -> cog weights inspect (synced).\n\n[short] skip 'requires local registry'\n\n# Start test registry\nregistry-start\n\n# Create weight files (small, deterministic)\nmkdir weights\nexec sh -c 'dd if=/dev/zero bs=512 count=1 2>/dev/null | tr \"\\0\" \"A\" > weights/model-a.bin'\nexec sh -c 'dd if=/dev/zero bs=512 count=1 2>/dev/null | tr \"\\0\" \"B\" > weights/model-b.bin'\n\n# Step 1: Build weights.lock\ncog weights build\nstderr 'Generated weights.lock'\nstderr '2 file'\nexists weights.lock\n\n# Verify lock file structure\nexec grep -q '\"name\": \"alpha\"' weights.lock\nexec grep -q '\"name\": \"beta\"' weights.lock\nexec grep -q '\"digest\": \"sha256:' weights.lock\n\n# Step 2: Push weights to local registry (repo only, no tag)\ncog weights push $TEST_REGISTRY/test/weights-model\nstderr 'Pushed 2 weight artifact'\n# Push output should show the full ref for each weight\nstderr 'weights-alpha-'\nstderr 'weights-beta-'\n\n# Verify tags with :tag are rejected\n! cog weights push $TEST_REGISTRY/test/weights-model:v1\nstderr 'includes a tag or digest'\n\n# Step 3: Inspect — both weights should be synced\ncog weights inspect $TEST_REGISTRY/test/weights-model --json\nstdout '\"status\": \"synced\"'\n! stdout '\"status\": \"local-only\"'\n! stdout '\"status\": \"digest-mismatch\"'\n# Inspect should show the ref and layers for each weight\nstdout '\"ref\":'\nstdout '\"layers\":'\n\n# Verify tags with :tag are rejected for inspect too\n! cog weights inspect $TEST_REGISTRY/test/weights-model:v1\nstderr 'includes a tag or digest'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\nweights:\n  - name: alpha\n    source: weights/model-a.bin\n    target: /weights/model-a.bin\n  - name: beta\n    source: weights/model-b.bin\n    target: /weights/model-b.bin\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/wheel_coglet_missing.txtar",
    "content": "# Test COGLET_WHEEL with non-existent path gives clear error\n#\n# This test verifies that when COGLET_WHEEL points to a non-existent file,\n# a clear error message is shown.\n\nenv COG_SDK_WHEEL=$REPO_ROOT/dist\nenv COGLET_WHEEL=/nonexistent/path/coglet.whl\n! cog build -t $TEST_IMAGE\nstderr 'path not found'\nstderr '/nonexistent/path/coglet.whl'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/wheel_resolution.txtar",
    "content": "# Test wheel resolution from environment variables\n#\n# This test verifies that COG_SDK_WHEEL and COGLET_WHEEL environment variables\n# correctly resolve wheel paths, including:\n# - COG_SDK_WHEEL pointing to a directory resolves wheels inside it\n# - Clear errors when wheel not found\n#\n# Note: Tests run from a temp directory, so we use $REPO_ROOT/dist\n# (an absolute path exported by the test harness) to find wheels.\n\n# Test 1: COG_SDK_WHEEL pointing to repo dist/ directory finds wheel\nenv COG_SDK_WHEEL=$REPO_ROOT/dist\ncog build --debug -t $TEST_IMAGE\nstderr 'Using local cog wheel:'\n\n# Test 2: Relative path that doesn't exist gives clear error  \nenv COG_SDK_WHEEL=./nonexistent/wheel.whl\n! cog build -t $TEST_IMAGE\nstderr 'path not found'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n"
  },
  {
    "path": "integration-tests/tests/zsh_package.txtar",
    "content": "# Test that zsh system package is installed and available in /bin\ncog build -t $TEST_IMAGE\n\ncog predict $TEST_IMAGE\nstdout ',sh,'\nstdout ',zsh,'\n\n-- cog.yaml --\nbuild:\n  python_version: \"3.12\"\n  system_packages:\n    - \"zsh\"\npredict: \"predict.py:Predictor\"\n\n-- predict.py --\nfrom cog import BasePredictor\nimport os\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        return \"hello \" + \",\".join(os.listdir(\"/bin\"))\n"
  },
  {
    "path": "mise.toml",
    "content": "# =============================================================================\n# Cog Development Tasks\n# =============================================================================\n#\n# Run `mise tasks` to see all available tasks.\n#\n# ## Task Caching\n#\n# Some build tasks use `sources` and `outputs` for caching. When sources haven't\n# changed since the last run, mise skips the task. To force a rebuild:\n#\n#     mise run build:sdk --force\n#     mise run build:coglet:wheel --force\n#\n# Cached tasks:\n#   - build:sdk              - skips if python/**/*.py unchanged\n#   - build:coglet:wheel*    - skips if crates/**/*.rs unchanged\n#   - generate:stubs         - skips if coglet-python source unchanged\n#   - generate:compat        - skips if tools/compatgen unchanged\n#   - docs                   - skips if docs/**/*.md unchanged\n#\n# Tasks that always run (no caching):\n#   - fmt:*, lint:*, test:*  - always check current state\n#   - clean:*                - always destructive\n#\n# =============================================================================\n\nexperimental_monorepo_root = true\n\n[monorepo]\nconfig_roots = [\".\"]\n\n[tools]\n\ngo = \"latest\"\nuv = \"0.9.26\"\n\"pipx:nox\" = { version = \"2025.11.12\", uvx = true, uvx_args = \"--python-preference=managed -p 3.13\" }\n\"aqua:ziglang/zig\" = \"0.15.2\"\n\"rust\" = { version = \"1.93.0\", components = \"rustfmt,clippy\", targets =\"x86_64-unknown-linux-gnu,aarch64-unknown-linux-gnu,aarch64-apple-darwin\" }\ncargo-binstall = \"1.16.6\"\n# Cargo tools - use aqua backend where available for faster binary downloads\n# and better security (cosign/SLSA verification). Remaining cargo: tools use binstall.\n\"aqua:EmbarkStudios/cargo-deny\" = \"0.19.0\"\n\"aqua:mitsuhiko/insta\" = \"1.46.0\"\n\"cargo:cargo-nextest\" = \"0.9.120\"\n\"cargo:maturin\" = \"1.11.5\"\n\"aqua:rust-lang/rustup\" = \"latest\"\n\"aqua:rust-lang/rustup/rustup-init\" = \"latest\"\n\"aqua:rust-cross/cargo-zigbuild\" = \"0.20.1\"\n\"aqua:gotestyourself/gotestsum\" = \"1.13.0\"\n\"aqua:golangci/golangci-lint\" = \"2.10.1\"\nruff = \"0.14.13\"\nty = \"0.0.10\"\n\n[env]\n_.path = \"./bin\"\n_.file = [\".env\"]\n_.python.venv = \".venv\"\n# Set REPO_ROOT only if not already set (e.g., by CI)\nREPO_ROOT = \"{{env.REPO_ROOT | default(value=config_root)}}\"\n# CGo required for go-tree-sitter (static Python schema parser)\nCGO_ENABLED = \"1\"\n\n[settings]\nlockfile = true\nexperimental = true\n\n# =============================================================================\n# Helper tasks (hidden)\n# =============================================================================\n\n[tasks._setup_dist]\nhide = true\nsilent = true\ndescription = \"Create dist directory\"\nrun = \"mkdir -p dist\"\n\n[tasks._setup_venv]\nhide = true\nsilent = true\ndescription = \"Ensure root .venv exists with Python\"\nrun = \"test -d .venv || uv venv --quiet\"\n\n[tasks._clean_dist]\nhide = true\nsilent = true\ndescription = \"Clean dist directory\"\nrun = \"rm -f dist/cog-*.whl dist/cog-*.tar.gz dist/coglet-*.whl\"\n\n# =============================================================================\n# Build tasks\n# =============================================================================\n\n[tasks.build]\nalias = \"build:all\"\ndescription = \"Build all components\"\nrun = [\n  { tasks = [\"build:cog\", \"build:coglet:wheel:linux-x64\", \"build:sdk\"] },\n  { task = \"_build_summary\" },\n]\n\n[tasks._build_summary]\nhide = true\ndescription = \"Print build artifacts summary\"\nrun = \"\"\"\n#!/usr/bin/env bash\necho \"\"\necho \"=== Build Artifacts ===\"\nif BINARY=$(ls dist/go/*/cog 2>/dev/null | head -1); then\n  VERSION=$(\"$BINARY\" --version 2>/dev/null || echo \"unknown\")\n  echo \"  cli:        $BINARY ($VERSION)\"\nfi\nfor whl in dist/coglet-*.whl; do\n  [ -f \"$whl\" ] && echo \"  coglet:     $whl\"\ndone\nfor whl in dist/cog-*.whl; do\n  [ -f \"$whl\" ] && echo \"  python-sdk: $whl\"\ndone\necho \"\"\n\"\"\"\n\n[tasks.install]\ndepends = [\"build:cog\"]\ndescription = \"Build and symlink cog CLI\"\nusage = 'arg \"[dest]\" help=\"Directory to symlink into (e.g. ~/.local/bin)\" default=\"~/.local/bin\"'\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\nDEST=\"${usage_dest/#\\\\~/$HOME}\"\nBINARY=$(ls dist/go/*/cog 2>/dev/null | head -1)\nif [ -z \"$BINARY\" ]; then\n    echo \"Error: no cog binary found in dist/go/. Run 'mise run build:cog' first.\" >&2\n    exit 1\nfi\nBINARY=\"$(cd \"$(dirname \"$BINARY\")\" && pwd)/$(basename \"$BINARY\")\"\nmkdir -p \"$DEST\"\nln -sf \"$BINARY\" \"$DEST/cog\"\necho \"Installed $DEST/cog -> $BINARY\"\n\"\"\"\n\n[tasks.\"build:cog\"]\ndescription = \"Build cog CLI (development)\"\nsources = [\"cmd/**/*.go\", \"pkg/**/*.go\", \"go.mod\", \"go.sum\", \"crates/Cargo.toml\"]\noutputs = [\"dist/go/*/cog\"]\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\n# Don't set COG_VERSION — let goreleaser's snapshot template produce a dev version\n# (e.g. 0.17.1-dev+gabcdef) so local wheel auto-detection works.\nGOFLAGS=-buildvcs=false go run github.com/goreleaser/goreleaser/v2@latest build --clean --snapshot --single-target --id cog --output cog\n\"\"\"\n\n[tasks.\"build:cog:release\"]\ndescription = \"Build cog CLI (release)\"\nrun = \"go run github.com/goreleaser/goreleaser/v2@latest build --clean --single-target --id cog --output cog\"\n\n[tasks.\"build:rust\"]\ndescription = \"Build Rust workspace\"\nrun = \"cargo build --manifest-path crates/Cargo.toml --workspace\"\n\n[tasks.\"build:rust:release\"]\ndescription = \"Build Rust workspace (release)\"\nrun = \"cargo build --manifest-path crates/Cargo.toml --workspace --release\"\n\n[tasks.\"build:coglet\"]\ndescription = \"Build coglet Python wheel (development, local install)\"\nrun = [\n  { task = \"_setup_venv\" },\n  \"maturin develop --manifest-path crates/coglet-python/Cargo.toml\",\n]\n\n[tasks.\"build:coglet:wheel\"]\ndescription = \"Build coglet Python wheel (native platform)\"\n# No sources/outputs caching: the output glob dist/coglet-*.whl is too broad\n# and falsely matches cross-compiled wheels (e.g. linux-x64), causing skips.\n# Use --force if you need to bypass mise's staleness check.\nrun = [\n  { tasks = [\"_setup_dist\", \"_setup_venv\"] },\n  \"maturin build --release --out dist --manifest-path crates/coglet-python/Cargo.toml\",\n]\n\n[tasks.\"build:coglet:wheel:linux-x64\"]\ndescription = \"Build coglet Python wheel for Linux x86_64\"\nsources = [\"crates/**/*.rs\", \"crates/**/Cargo.toml\", \"Cargo.lock\"]\noutputs = [\"dist/coglet-*manylinux*x86_64*.whl\"]\nrun = [\n  { tasks = [\"_setup_dist\", \"_setup_venv\"] },\n  \"maturin build --release --out dist --manifest-path crates/coglet-python/Cargo.toml --target x86_64-unknown-linux-gnu --zig\",\n]\n\n[tasks.\"build:coglet:wheel:linux-arm64\"]\ndescription = \"Build coglet Python wheel for Linux ARM64\"\nsources = [\"crates/**/*.rs\", \"crates/**/Cargo.toml\", \"Cargo.lock\"]\noutputs = [\"dist/coglet-*manylinux*aarch64*.whl\"]\nrun = [\n  { tasks = [\"_setup_dist\", \"_setup_venv\"] },\n  \"maturin build --release --out dist --manifest-path crates/coglet-python/Cargo.toml --target aarch64-unknown-linux-gnu --zig\",\n]\n\n[tasks.\"build:coglet:wheel:darwin-arm64\"]\ndescription = \"Build coglet Python wheel for macOS ARM64\"\nsources = [\"crates/**/*.rs\", \"crates/**/Cargo.toml\", \"Cargo.lock\"]\noutputs = [\"dist/coglet-*-macosx_*_arm64.whl\"]\nrun = [\n  { tasks = [\"_setup_dist\", \"_setup_venv\"] },\n  \"maturin build --release --out dist --manifest-path crates/coglet-python/Cargo.toml --target aarch64-apple-darwin\",\n]\n\n[tasks.\"build:coglet:wheel:darwin-x64\"]\ndescription = \"Build coglet Python wheel for macOS x86_64\"\nsources = [\"crates/**/*.rs\", \"crates/**/Cargo.toml\", \"Cargo.lock\"]\noutputs = [\"dist/coglet-*-macosx_*_x86_64.whl\"]\nrun = [\n  { tasks = [\"_setup_dist\", \"_setup_venv\"] },\n  \"maturin build --release --out dist --manifest-path crates/coglet-python/Cargo.toml --target x86_64-apple-darwin\",\n]\n\n[tasks.\"build:coglet:wheel:all\"]\ndescription = \"Build coglet Python wheels for all platforms\"\nrun = [\n  { task = \"_setup_dist\" },\n  { tasks = [\"build:coglet:wheel:linux-x64\", \"build:coglet:wheel:linux-arm64\"] },\n  { task = \"build:coglet:wheel\" },\n]\n\n[tasks.\"build:sdk\"]\ndescription = \"Build cog SDK wheel\"\nsources = [\"python/**/*.py\", \"pyproject.toml\", \"crates/Cargo.toml\"]\noutputs = [\"dist/cog-*.whl\", \"dist/cog-*.tar.gz\"]\nrun = [\n  { tasks = [\"_setup_dist\", \"_setup_venv\"] },\n  \"\"\"\n#!/usr/bin/env bash\nset -euo pipefail\n# Version from Cargo.toml, converted to PEP 440\nRAW=$(grep '^version' crates/Cargo.toml | head -1 | sed 's/.*\"\\\\(.*\\\\)\"/\\\\1/')\nexport SETUPTOOLS_SCM_PRETEND_VERSION=$(echo \"$RAW\" | sed -E 's/-alpha/a/; s/-beta/b/; s/-rc/rc/; s/-dev/.dev/')\necho \"Building SDK wheel: $SETUPTOOLS_SCM_PRETEND_VERSION\"\nuv build --out-dir=dist .\n\"\"\",\n]\n\n[tasks.\"build:wheels\"]\ndescription = \"Build all wheels (coglet + sdk)\"\nrun = [\n  { task = \"_clean_dist\" },\n  { task = \"_setup_dist\" },\n  { tasks = [\"build:coglet:wheel:all\", \"build:sdk\"] },\n]\n\n# =============================================================================\n# Test tasks\n# =============================================================================\n\n[tasks.test]\ndescription = \"Run all unit tests (set INTEGRATION_TESTS=1 to include integration)\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\nmise run test:go\nmise run test:rust\nmise run test:python\nif [ \"${INTEGRATION_TESTS:-}\" = \"1\" ]; then\n  mise run test:integration\nfi\n\"\"\"\n\n[tasks.\"test:go\"]\ndescription = \"Run Go tests\"\nrun = \"gotestsum -- -short -timeout 1200s -parallel 5 ./...\"\n\n[tasks.\"test:rust\"]\ndescription = \"Run Rust workspace tests\"\nrun = \"cargo nextest run --manifest-path crates/Cargo.toml --workspace --exclude coglet-python --no-tests=pass\"\n\n[tasks.\"test:python\"]\ndescription = \"Run Python SDK tests (latest supported Python)\"\ndepends = [\"build:coglet:wheel\"]\nrun = \"nox -s tests -p 3.13\"\n\n[tasks.\"test:python:all\"]\ndescription = \"Run Python SDK tests on all supported Python versions\"\ndepends = [\"build:coglet:wheel\"]\nrun = \"nox -s tests\"\n\n[tasks.\"test:coglet:python\"]\ndescription = \"Run coglet Python binding tests (latest Python)\"\ndepends = [\"build:coglet:wheel\"]\nrun = \"nox -s coglet -p 3.13\"\n\n[tasks.\"test:coglet:python:all\"]\ndescription = \"Run coglet Python binding tests on all supported Python versions\"\ndepends = [\"build:coglet:wheel\"]\nrun = \"nox -s coglet\"\n\n[tasks.\"test:fuzz\"]\ndescription = \"Run Go fuzz tests (FUZZTIME=30s per target by default)\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\nFUZZTIME=\"${FUZZTIME:-30s}\"\necho \"Fuzzing schema type resolution ($FUZZTIME)...\"\ngo test ./pkg/schema/ -run='^$' -fuzz=FuzzResolveSchemaType -fuzztime=\"$FUZZTIME\"\necho \"Fuzzing JSON schema generation ($FUZZTIME)...\"\ngo test ./pkg/schema/ -run='^$' -fuzz=FuzzJSONSchema -fuzztime=\"$FUZZTIME\"\necho \"Fuzzing Python parser ($FUZZTIME)...\"\ngo test ./pkg/schema/python/ -run='^$' -fuzz=FuzzParsePredictor -fuzztime=\"$FUZZTIME\"\necho \"Fuzzing type annotation parsing ($FUZZTIME)...\"\ngo test ./pkg/schema/python/ -run='^$' -fuzz=FuzzParseTypeAnnotation -fuzztime=\"$FUZZTIME\"\necho \"All fuzz targets passed.\"\n\"\"\"\n\n[tasks.\"test:integration\"]\ndescription = \"Run integration tests (skips slow tests by default, set SHORT=0 for full suite)\"\ndepends = [\"clean:integration\", \"build:cog\", \"build:sdk\", \"build:coglet:wheel:linux-x64\"]\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\nSHORT_FLAG=\"-short\"\nif [ \"${SHORT:-1}\" = \"0\" ]; then\n  SHORT_FLAG=\"\"\nfi\n# If first arg is a bare name (no dash), treat as test name filter;\n# remaining args are passed through to go test.\n# e.g. mise run test:integration coglet_large_output -count=4\nif [ $# -gt 0 ] && [[ \"$1\" != -* ]]; then\n  gotestsum -- -tags integration -v $SHORT_FLAG -run \"TestIntegration/$1\" \"${@:2}\" -timeout 30m ./integration-tests/...\nelse\n  gotestsum -- -tags integration -v $SHORT_FLAG -parallel ${TEST_PARALLEL:-4} \"$@\" -timeout 30m ./integration-tests/...\nfi\n\"\"\"\n\n# =============================================================================\n# Format tasks\n# =============================================================================\n\n[tasks.fmt]\nalias = [\"format\", \"fmt:check\", \"format:check\"]\ndescription = \"Check formatting for all languages (non-destructive)\"\nrun = [\n  { tasks = [\"fmt:go\", \"fmt:rust\", \"fmt:python\", \"fmt:docs\"] },\n]\n\n[tasks.\"fmt:fix\"]\nalias = \"format:fix\"\ndescription = \"Fix formatting for all languages\"\nrun = [\n  { tasks = [\"fmt:go:fix\", \"fmt:rust:fix\", \"fmt:python:fix\", \"fmt:docs:fix\"] },\n]\n\n# Go formatting\n[tasks.\"fmt:go\"]\nalias = \"fmt:go:check\"\ndescription = \"Check Go formatting\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\ngo tool goimports -d .\ntest -z \"$(go tool goimports -l .)\"\n\"\"\"\n\n[tasks.\"fmt:go:fix\"]\ndescription = \"Fix Go formatting\"\nrun = \"go tool goimports -w -d .\"\n\n# Rust formatting\n[tasks.\"fmt:rust\"]\nalias = \"fmt:rust:check\"\ndescription = \"Check Rust formatting\"\nrun = \"cargo fmt --manifest-path crates/Cargo.toml --all -- --check\"\n\n[tasks.\"fmt:rust:fix\"]\ndescription = \"Fix Rust formatting\"\nrun = \"cargo fmt --manifest-path crates/Cargo.toml --all\"\n\n# Python formatting\n[tasks.\"fmt:python\"]\nalias = \"fmt:python:check\"\ndescription = \"Check Python formatting\"\nrun = \"ruff format --check .\"\n\n[tasks.\"fmt:python:fix\"]\ndescription = \"Fix Python formatting\"\nrun = \"ruff format .\"\n\n# Docs formatting\n[tasks.\"fmt:docs\"]\nalias = \"fmt:docs:check\"\ndescription = \"Check docs formatting\"\nrun = \"npx prettier --check 'docs/**/*.md' README.md\"\n\n[tasks.\"fmt:docs:fix\"]\ndescription = \"Fix docs formatting\"\nrun = \"npx prettier -w 'docs/**/*.md' README.md\"\n\n# =============================================================================\n# Lint tasks\n# =============================================================================\n\n[tasks.lint]\nalias = \"lint:check\"\ndescription = \"Run linters for all languages (non-destructive)\"\nrun = [\n  { tasks = [\"lint:go\", \"lint:rust\", \"lint:python\"] },\n]\n\n[tasks.\"lint:fix\"]\ndescription = \"Fix lint issues for all languages\"\nrun = [\n  { tasks = [\"lint:go:fix\", \"lint:rust:fix\", \"lint:python:fix\"] },\n]\n\n# Go linting\n[tasks.\"lint:go\"]\nalias = \"lint:go:check\"\ndescription = \"Lint Go code\"\nrun = \"golangci-lint run ./...\"\n\n[tasks.\"lint:go:fix\"]\ndescription = \"Fix Go lint issues (limited auto-fix)\"\nrun = \"golangci-lint run --fix ./...\"\n\n# Rust linting\n[tasks.\"lint:rust\"]\nalias = [\"lint:rust:check\", \"lint:rust:clippy\"]\ndescription = \"Lint Rust code (clippy)\"\nrun = \"cargo clippy --manifest-path crates/Cargo.toml --workspace -- -D warnings\"\n\n[tasks.\"lint:rust:deny\"]\ndescription = \"Check Rust licenses and advisories\"\nrun = \"cargo deny --manifest-path crates/Cargo.toml check\"\n\n[tasks.\"lint:rust:fix\"]\ndescription = \"Fix Rust lint issues\"\nrun = \"cargo clippy --manifest-path crates/Cargo.toml --workspace --fix --allow-dirty -- -D warnings\"\n\n# Python linting\n[tasks.\"lint:python\"]\nalias = \"lint:python:check\"\ndescription = \"Lint Python code\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\nruff check .\nmise run typecheck:python\n\"\"\"\n\n[tasks.\"lint:python:fix\"]\ndescription = \"Fix Python lint issues\"\nrun = \"ruff check --fix .\"\n\n# =============================================================================\n# Typecheck tasks\n# =============================================================================\n\n[tasks.typecheck]\ndescription = \"Run type checking for all languages\"\nrun = [\n  { tasks = [\"typecheck:rust\", \"typecheck:python\"] },\n]\n\n[tasks.\"typecheck:rust\"]\ndescription = \"Type check Rust code (cargo check)\"\nrun = \"cargo check --manifest-path crates/Cargo.toml --workspace\"\n\n[tasks.\"typecheck:python\"]\ndescription = \"Type check Python code\"\nrun = \"nox -s typecheck\"\n\n# =============================================================================\n# Generate tasks\n# =============================================================================\n\n[tasks.generate]\ndescription = \"Run all code generation\"\nrun = [\n  { tasks = [\"generate:stubs\"] },\n]\n\n[tasks.\"generate:stubs\"]\nalias = \"stub:generate\"\ndescription = \"Generate Python type stubs for coglet\"\nsources = [\"crates/coglet-python/src/**/*.rs\", \"crates/coglet-python/Cargo.toml\", \"crates/coglet-python/coglet/__init__.py\"]\noutputs = [\"crates/coglet-python/coglet/_sdk/__init__.pyi\", \"crates/coglet-python/coglet/__init__.pyi\"]\ndir = \"crates/coglet-python\"\nrun = [\n  { task = \"_setup_venv\" },\n  \"uv run --active cargo run --bin stub_gen\",\n]\n\n[tasks.\"generate:compat\"]\ndescription = \"Regenerate CUDA/PyTorch/TensorFlow compatibility matrices\"\nsources = [\"tools/compatgen/**/*.go\"]\noutputs = [\"pkg/config/cuda_base_images.json\", \"pkg/config/torch_compatibility_matrix.json\", \"pkg/config/tf_compatibility_matrix.json\"]\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\ntarget=\"${1:-all}\"\ncase \"$target\" in\n  cuda)\n    echo \"Generating CUDA base images...\"\n    go run ./tools/compatgen/main.go cuda -o pkg/config/cuda_base_images.json\n    ;;\n  torch)\n    echo \"Generating PyTorch compatibility matrix...\"\n    go run ./tools/compatgen/main.go torch -o pkg/config/torch_compatibility_matrix.json\n    ;;\n  tensorflow|tf)\n    echo \"Generating TensorFlow compatibility matrix...\"\n    go run ./tools/compatgen/main.go tensorflow -o pkg/config/tf_compatibility_matrix.json\n    ;;\n  all)\n    echo \"Generating CUDA base images...\"\n    go run ./tools/compatgen/main.go cuda -o pkg/config/cuda_base_images.json\n    echo \"Generating PyTorch compatibility matrix...\"\n    go run ./tools/compatgen/main.go torch -o pkg/config/torch_compatibility_matrix.json\n    echo \"Generating TensorFlow compatibility matrix...\"\n    go run ./tools/compatgen/main.go tensorflow -o pkg/config/tf_compatibility_matrix.json\n    ;;\n  *)\n    echo \"Unknown target: $target\"\n    echo \"Usage: mise run generate:compat [cuda|torch|tensorflow|all]\"\n    exit 1\n    ;;\nesac\necho \"Done.\"\n\"\"\"\n\n# =============================================================================\n# Stub tasks\n# =============================================================================\n\n[tasks.\"stub:check\"]\ndescription = \"Check that coglet Python stubs are up to date\"\ndir = \"crates/coglet-python\"\nrun = [\n  { task = \"generate:stubs\" },\n  '''\n#!/usr/bin/env bash\nset -e\nif ! git diff --quiet -- '**/*.pyi'; then\n  echo \"ERROR: Stubs are out of date:\"\n  git diff -- '**/*.pyi'\n  echo \"\"\n  echo \"Run 'mise run generate:stubs' to update.\"\n  exit 1\nfi\necho \"Stubs are up to date.\"\n''',\n]\n\n[tasks.\"stub:typecheck\"]\ndescription = \"Type-check coglet stubs with ty\"\nrun = \"ty check crates/coglet-python/coglet/__init__.pyi\"\n\n# =============================================================================\n# Clean tasks\n# =============================================================================\n\n[tasks.clean]\ndescription = \"Clean all build artifacts\"\nrun = [\n  { tasks = [\"clean:go\", \"clean:rust\", \"clean:python\", \"clean:integration\"] },\n  { task = \"_clean_dist\" },\n]\n\n[tasks.\"clean:go\"]\ndescription = \"Clean Go build artifacts\"\nrun = \"rm -rf cog base-image dist/go\"\n\n[tasks.\"clean:rust\"]\ndescription = \"Clean Rust build artifacts\"\nrun = \"cd crates && cargo clean\"\n\n[tasks.\"clean:python\"]\ndescription = \"Clean Python build artifacts\"\nrun = \"rm -rf .tox build python/cog.egg-info .venv crates/coglet-python/.venv crates/coglet-python/coglet/*.so\"\n\n[tasks.\"clean:integration\"]\ndescription = \"Clean cached integration test binary and embedded wheels\"\nrun = \"rm -f integration-tests/.bin/cog pkg/wheels/cog-*.whl pkg/wheels/coglet-*.whl\"\n\n# =============================================================================\n# Docs tasks\n# =============================================================================\n\n[tasks.docs]\ndescription = \"Build documentation\"\nsources = [\"docs/**/*.md\", \"README.md\", \"CONTRIBUTING.md\", \"mkdocs.yml\"]\noutputs = [\"site/**\"]\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\nuv pip install mkdocs-material\nsed 's/docs\\\\///g' README.md > ./docs/README.md\ncp CONTRIBUTING.md ./docs/\nmkdocs build\n\"\"\"\n\n[tasks.\"docs:serve\"]\ndescription = \"Serve documentation locally\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\nuv pip install mkdocs-material\nsed 's/docs\\\\///g' README.md > ./docs/README.md\ncp CONTRIBUTING.md ./docs/\nmkdocs serve\n\"\"\"\n\n[tasks.\"docs:llm\"]\ndescription = \"Update LLM documentation (llms.txt)\"\ndepends = [\"docs:cli\"]\nsources = [\"README.md\", \"docs/*.md\"]\noutputs = [\"docs/llms.txt\"]\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\n# Concatenate README (minus contributors section) + all docs into llms.txt\n# Use awk instead of sed for cross-platform compatibility (BSD sed vs GNU sed)\n# Only include git-tracked files (docs/ may contain mkdocs-generated copies of CONTRIBUTING.md, README.md)\n(awk '/^## Contributors/{exit} {print}' README.md; for file in $(git ls-files 'docs/*.md'); do printf '\\n\\n---\\n\\n' && cat \"$file\"; done) > docs/llms.txt\necho \"Updated docs/llms.txt\"\n\"\"\"\n\n[tasks.\"docs:llm:check\"]\ndescription = \"Check that llms.txt is up to date\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\ntmpfile=$(mktemp)\ntrap 'rm -f \"$tmpfile\"' EXIT\n# Generate to temp file and compare\n# Only include git-tracked files (docs/ may contain mkdocs-generated copies of CONTRIBUTING.md, README.md)\n(awk '/^## Contributors/{exit} {print}' README.md; for file in $(git ls-files 'docs/*.md'); do printf '\\n\\n---\\n\\n' && cat \"$file\"; done) > \"$tmpfile\"\nif ! diff -q \"$tmpfile\" docs/llms.txt > /dev/null 2>&1; then\n  echo \"ERROR: docs/llms.txt is out of date. Run 'mise run docs:llm' to update.\"\n  exit 1\nfi\necho \"docs/llms.txt is up to date\"\n\"\"\"\n\n[tasks.\"docs:cli\"]\ndescription = \"Generate CLI reference documentation\"\nsources = [\"pkg/cli/*.go\", \"cmd/cog/*.go\"]\noutputs = [\"docs/cli.md\"]\nrun = \"go run ./tools/gendocs/main.go -o docs/cli.md\"\n\n[tasks.\"docs:cli:check\"]\ndescription = \"Check that CLI docs are up to date\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\ntmpfile=$(mktemp)\ntrap 'rm -f \"$tmpfile\"' EXIT\n# Generate to temp file and compare\ngo run ./tools/gendocs/main.go -o \"$tmpfile\"\nif ! diff -q \"$tmpfile\" docs/cli.md > /dev/null 2>&1; then\n  echo \"ERROR: docs/cli.md is out of date. Run 'mise run docs:cli' to update.\"\n  exit 1\nfi\necho \"docs/cli.md is up to date\"\n\"\"\"\n\n# =============================================================================\n# CI tasks - granular for parallel execution and caching\n# =============================================================================\n\n# Build tasks (run first to produce artifacts)\n[tasks.\"ci:build\"]\ndescription = \"CI: Build all artifacts\"\nrun = [\n  { tasks = [\"ci:build:sdk\", \"ci:build:coglet\"] },\n]\n\n[tasks.\"ci:build:sdk\"]\ndescription = \"CI: Build SDK wheel and sdist\"\nrun = [\n  { task = \"_setup_dist\" },\n  \"\"\"\n#!/usr/bin/env bash\nset -euo pipefail\n# Version from Cargo.toml, converted to PEP 440\nRAW=$(grep '^version' crates/Cargo.toml | head -1 | sed 's/.*\"\\\\(.*\\\\)\"/\\\\1/')\nexport SETUPTOOLS_SCM_PRETEND_VERSION=$(echo \"$RAW\" | sed -E 's/-alpha/a/; s/-beta/b/; s/-rc/rc/; s/-dev/.dev/')\necho \"Building SDK wheel: $SETUPTOOLS_SCM_PRETEND_VERSION\"\nuv build --out-dir=dist .\n\"\"\",\n]\n\n[tasks.\"ci:build:coglet\"]\ndescription = \"CI: Build coglet wheel\"\nrun = [\n  { task = \"_setup_dist\" },\n  { task = \"build:coglet:wheel:linux-x64\" },\n]\n\n[tasks.\"ci:test:integration\"]\ndescription = \"CI: Run integration tests with GitHub Actions output (full suite)\"\n# exec ensures signals (SIGTERM from CI cancellation) go directly to gotestsum\nrun = \"exec gotestsum --format github-actions -- -tags integration -parallel ${TEST_PARALLEL:-4} -timeout 30m ./integration-tests/...\"\n\n# =============================================================================\n# Publish tasks (future)\n# =============================================================================\n\n[tasks.\"publish:coglet\"]\ndescription = \"Publish coglet to PyPI\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\necho \"TODO: Implement coglet PyPI publish\"\necho \"Wheels in dist/: $(ls dist/coglet-*.whl 2>/dev/null || echo 'none')\"\n\"\"\"\n\n[tasks.\"publish:sdk\"]\ndescription = \"Publish cog SDK to PyPI\"\nrun = \"\"\"\n#!/usr/bin/env bash\nset -e\necho \"TODO: Implement SDK PyPI publish\"\necho \"Wheels in dist/: $(ls dist/cog-*.whl 2>/dev/null || echo 'none')\"\n\"\"\"\n"
  },
  {
    "path": "mkdocs.yml",
    "content": "site_name: Cog\nrepo_url: https://github.com/replicate/cog\ndocs_dir: docs/\nnav:\n  - README: README.md\n  - Getting Started: getting-started.md\n  - Using your own model: getting-started-own-model.md\n  - Deploy your model: deploy.md\n  - YAML spec: yaml.md\n  - Prediction API: python.md\n  - Training API: training.md\n  - HTTP API: http.md\n  - CLI: cli.md\n  - Environment variables: environment.md\n  - Private registry: private-package-registry.md\n  - Notebooks: notebooks.md\n  - Windows: wsl2/wsl2.md\n  - Contributing: CONTRIBUTING.md\n  - License: https://github.com/replicate/cog/blob/main/LICENSE\n  - llms.txt: llms.txt\n\ntheme:\n  name: material\n  font:\n    text: \"Roboto\"\n    code: \"Roboto Mono\"\n  favicon: favicon.svg\n\n  # Display a link to edit pages right on GitHub\n  features:\n    - content.action.edit\n\n  icon:\n    logo: material/cog\n    repo: simple/github\n  palette:\n    # Palette toggle for light mode\n    - media: \"(prefers-color-scheme: light)\"\n      scheme: default\n      primary: black\n      toggle:\n        icon: material/weather-night\n        name: Switch to dark mode\n\n    # Palette toggle for dark mode\n    - media: \"(prefers-color-scheme: dark)\"\n      scheme: slate\n      primary: black\n      toggle:\n        icon: material/weather-sunny\n        name: Switch to light mode\n\nedit_uri: edit/main/docs/\n\nmarkdown_extensions:\n  - toc:\n      permalink: \"#\"\n  - markdown.extensions.codehilite:\n      guess_lang: true\n  - admonition\n  - codehilite\n  - extra\n  - pymdownx.highlight\n  - pymdownx.superfences\nextra_css:\n  - stylesheets/extra.css\n\nextra:\n  # Hide the \"made with material\" thing\n  generator: false\n  social:\n    - icon: simple/github\n      link: https://github.com/replicate/cog\n    - icon: simple/discord\n      link: https://discord.gg/replicate\n    - icon: simple/x\n      link: https://x.com/replicate\n    - icon: simple/youtube\n      link: https://youtube.com/@replicatehq\n\ncopyright: Cog is an open-source project from Replicate\n"
  },
  {
    "path": "noxfile.py",
    "content": "\"\"\"Nox sessions for cog Python SDK testing.\"\"\"\n\nimport glob\nimport platform\n\nimport nox\n\n# Use uv for venv creation and Python management (uv auto-downloads Python if needed)\nnox.options.default_venv_backend = \"uv\"\n\nPYTHON_VERSIONS = [\"3.10\", \"3.11\", \"3.12\", \"3.13\"]\nPYTHON_DEFAULT = \"3.13\"\n\n# Test dependencies (mirrored from pyproject.toml [dependency-groups].test)\nTEST_DEPS = [\n    \"pytest\",\n    \"pytest-timeout\",\n    \"pytest-xdist\",\n    \"pytest-cov\",\n]\n\n\ndef _find_compatible_wheel(pattern: str) -> str | None:\n    \"\"\"Find a wheel matching the current platform from dist/.\n\n    Returns None when no wheels exist at all.  Raises RuntimeError when\n    wheels exist but none are compatible — that means the build produced\n    the wrong platform and should be fixed, not silently papered over.\n    \"\"\"\n    wheels = glob.glob(pattern)\n    if not wheels:\n        return None\n\n    system = platform.system().lower()\n    machine = platform.machine().lower()\n    platform_tags = {\n        (\"darwin\", \"arm64\"): \"macosx\",\n        (\"darwin\", \"x86_64\"): \"macosx\",\n        (\"linux\", \"x86_64\"): \"manylinux\",\n        (\"linux\", \"aarch64\"): \"manylinux\",\n    }\n    tag = platform_tags.get((system, machine))\n    if tag:\n        for whl in wheels:\n            if tag in whl or \"none-any\" in whl:\n                return whl\n        raise RuntimeError(\n            f\"Found wheel(s) in dist/ but none compatible with {system}/{machine}:\\n\"\n            + \"\\n\".join(f\"  {w}\" for w in wheels)\n            + \"\\nRun 'mise run build:coglet:wheel' to build a native wheel.\"\n        )\n\n    # Unknown platform — let pip figure it out\n    return wheels[0]\n\n\ndef _install_coglet(session: nox.Session) -> None:\n    \"\"\"Install coglet wheel (required dependency).\"\"\"\n    whl = _find_compatible_wheel(\"dist/coglet-*.whl\")\n    if whl:\n        session.install(whl)\n    else:\n        session.error(\n            \"No coglet wheel found in dist/. Run 'mise run build:coglet:wheel' first.\"\n        )\n\n\ndef _install_package(session: nox.Session) -> None:\n    \"\"\"Install the cog SDK and coglet dependency.\"\"\"\n    _install_coglet(session)\n    whl = _find_compatible_wheel(\"dist/cog-*.whl\")\n    if whl:\n        session.install(whl)\n    else:\n        # No pre-built wheel — editable install from source.\n        # This fails in CI (setuptools_scm needs a full git checkout),\n        # so CI must run build:sdk first.\n        session.install(\"-e\", \".\")\n\n\n@nox.session(python=PYTHON_VERSIONS)\ndef tests(session: nox.Session) -> None:\n    \"\"\"Run the test suite.\"\"\"\n    _install_package(session)\n    session.install(*TEST_DEPS)\n    args = session.posargs or [\"-n\", \"auto\", \"-vv\"]\n    session.run(\n        \"pytest\",\n        \"python/tests\",\n        \"--cov=python/cog\",\n        \"--cov-report=term-missing:skip-covered\",\n        *args,\n    )\n\n\n@nox.session(python=PYTHON_DEFAULT)\ndef typecheck(session: nox.Session) -> None:\n    \"\"\"Run type checking with pyright.\"\"\"\n    _install_package(session)\n    session.install(\"pyright==1.1.375\")\n    session.run(\"pyright\", *session.posargs)\n\n\n@nox.session(name=\"coglet\", python=PYTHON_VERSIONS)\ndef coglet_tests(session: nox.Session) -> None:\n    \"\"\"Run coglet-python binding tests.\"\"\"\n    _install_package(session)\n    session.install(\"pytest\", \"requests\")\n    session.run(\"pytest\", \"crates/coglet-python/tests\", \"-v\", *session.posargs)\n"
  },
  {
    "path": "pkg/cli/baseimage.go",
    "content": "package cli\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"strings\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/dockercontext\"\n\t\"github.com/replicate/cog/pkg/dockerfile\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/update\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nvar (\n\tbaseImageCUDAVersion   string\n\tbaseImagePythonVersion string\n\tbaseImageTorchVersion  string\n)\n\nfunc NewBaseImageRootCommand() (*cobra.Command, error) {\n\trootCmd := cobra.Command{\n\t\tUse:     \"base-image\",\n\t\tShort:   \"Cog base image commands. This is an experimental feature with no guarantees of future support.\",\n\t\tVersion: fmt.Sprintf(\"%s (built %s)\", global.Version, global.BuildTime),\n\t\t// This stops errors being printed because we print them in cmd/cog/cog.go\n\t\tPersistentPreRun: func(cmd *cobra.Command, args []string) {\n\t\t\tif global.Debug {\n\t\t\t\tconsole.SetLevel(console.DebugLevel)\n\t\t\t}\n\t\t\tcmd.SilenceUsage = true\n\t\t\tif err := update.DisplayAndCheckForRelease(cmd.Context()); err != nil {\n\t\t\t\tconsole.Debugf(\"%s\", err)\n\t\t\t}\n\t\t},\n\t\tSilenceErrors: true,\n\t}\n\tsetPersistentFlags(&rootCmd)\n\n\trootCmd.AddCommand(\n\t\tnewBaseImageDockerfileCommand(),\n\t\tnewBaseImageBuildCommand(),\n\t\tnewBaseImageGenerateMatrix(),\n\t)\n\n\treturn &rootCmd, nil\n}\n\nfunc newBaseImageGenerateMatrix() *cobra.Command {\n\tvar cmd = &cobra.Command{\n\t\tUse:   \"generate-matrix\",\n\t\tShort: \"Generate a matrix of Cog base image versions (JSON)\",\n\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\tvalidCudaVersions := strings.FieldsFunc(baseImageCUDAVersion, func(c rune) bool {\n\t\t\t\treturn c == ','\n\t\t\t})\n\t\t\tvalidPythonVersions := strings.FieldsFunc(baseImagePythonVersion, func(c rune) bool {\n\t\t\t\treturn c == ','\n\t\t\t})\n\t\t\tvalidTorchVersions := strings.FieldsFunc(baseImageTorchVersion, func(c rune) bool {\n\t\t\t\treturn c == ','\n\t\t\t})\n\n\t\t\tallConfigurations := dockerfile.BaseImageConfigurations()\n\t\t\tfilteredMatrix := make([]dockerfile.BaseImageConfiguration, 0, len(allConfigurations))\n\t\t\tfor _, config := range allConfigurations {\n\t\t\t\tvar found bool\n\t\t\t\tif len(validCudaVersions) > 0 {\n\t\t\t\t\tfound = false\n\t\t\t\t\tfor _, validCudaVersion := range validCudaVersions {\n\t\t\t\t\t\tif config.CUDAVersion == validCudaVersion {\n\t\t\t\t\t\t\tfound = true\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\tif !found {\n\t\t\t\t\t\tcontinue\n\t\t\t\t\t}\n\t\t\t\t}\n\n\t\t\t\tif len(validPythonVersions) > 0 {\n\t\t\t\t\tfound = false\n\t\t\t\t\tfor _, validPythonVersion := range validPythonVersions {\n\t\t\t\t\t\tif config.PythonVersion == validPythonVersion {\n\t\t\t\t\t\t\tfound = true\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\tif !found {\n\t\t\t\t\t\tcontinue\n\t\t\t\t\t}\n\t\t\t\t}\n\n\t\t\t\tif len(validTorchVersions) > 0 {\n\t\t\t\t\tfound = false\n\t\t\t\t\tfor _, validTorchVersion := range validTorchVersions {\n\t\t\t\t\t\tif config.TorchVersion == validTorchVersion {\n\t\t\t\t\t\t\tfound = true\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\tif !found {\n\t\t\t\t\t\tcontinue\n\t\t\t\t\t}\n\t\t\t\t}\n\n\t\t\t\tfilteredMatrix = append(filteredMatrix, config)\n\t\t\t}\n\n\t\t\toutput, err := json.Marshal(filteredMatrix)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tfmt.Println(string(output))\n\t\t\treturn nil\n\t\t},\n\t\tArgs: cobra.MaximumNArgs(0),\n\t}\n\taddBaseImageFlags(cmd)\n\treturn cmd\n}\n\nfunc newBaseImageDockerfileCommand() *cobra.Command {\n\tvar cmd = &cobra.Command{\n\t\tUse:   \"dockerfile\",\n\t\tShort: \"Display Cog base image Dockerfile\",\n\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\tctx := cmd.Context()\n\n\t\t\tgenerator, err := baseImageGeneratorFromFlags(ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tdockerfile, err := generator.GenerateDockerfile(ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tfmt.Println(dockerfile)\n\t\t\treturn nil\n\t\t},\n\t\tArgs: cobra.MaximumNArgs(0),\n\t}\n\taddBaseImageFlags(cmd)\n\taddNoCacheFlag(cmd)\n\taddBuildProgressOutputFlag(cmd)\n\n\treturn cmd\n}\n\nfunc newBaseImageBuildCommand() *cobra.Command {\n\tvar cmd = &cobra.Command{\n\t\tUse:   \"build\",\n\t\tShort: \"Build Cog base image\",\n\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\tctx := cmd.Context()\n\n\t\t\tdockerClient, err := docker.NewClient(ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\n\t\t\tgenerator, err := baseImageGeneratorFromFlags(ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tdockerfileContents, err := generator.GenerateDockerfile(ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\n\t\t\tcwd, err := os.Getwd()\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tbaseImageName := dockerfile.BaseImageName(baseImageCUDAVersion, baseImagePythonVersion, baseImageTorchVersion)\n\n\t\t\tbuildOpts := command.ImageBuildOptions{\n\t\t\t\tWorkingDir:         cwd,\n\t\t\t\tDockerfileContents: dockerfileContents,\n\t\t\t\tImageName:          baseImageName,\n\t\t\t\tNoCache:            buildNoCache,\n\t\t\t\tProgressOutput:     buildProgressOutput,\n\t\t\t\tEpoch:              &config.BuildSourceEpochTimestamp,\n\t\t\t\tContextDir:         dockercontext.StandardBuildDirectory,\n\t\t\t}\n\t\t\tif _, err := dockerClient.ImageBuild(ctx, buildOpts); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tfmt.Println(\"Successfully built image: \" + baseImageName)\n\t\t\treturn nil\n\t\t},\n\t\tArgs: cobra.MaximumNArgs(0),\n\t}\n\taddBaseImageFlags(cmd)\n\n\treturn cmd\n}\n\nfunc addBaseImageFlags(cmd *cobra.Command) {\n\tcmd.Flags().StringVar(&baseImageCUDAVersion, \"cuda\", \"\", \"CUDA version\")\n\tcmd.Flags().StringVar(&baseImagePythonVersion, \"python\", \"\", \"Python version\")\n\tcmd.Flags().StringVar(&baseImageTorchVersion, \"torch\", \"\", \"Torch version\")\n\taddBuildTimestampFlag(cmd)\n}\n\nfunc baseImageGeneratorFromFlags(ctx context.Context) (*dockerfile.BaseImageGenerator, error) {\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tclient := registry.NewRegistryClient()\n\treturn dockerfile.NewBaseImageGenerator(\n\t\tctx,\n\t\tclient,\n\t\tbaseImageCUDAVersion,\n\t\tbaseImagePythonVersion,\n\t\tbaseImageTorchVersion,\n\t\tdockerClient,\n\t\ttrue,\n\t)\n}\n"
  },
  {
    "path": "pkg/cli/build.go",
    "content": "package cli\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\t\"strings\"\n\n\t\"github.com/spf13/cobra\"\n\t\"github.com/spf13/pflag\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/model\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nvar buildTag string\nvar buildSeparateWeights bool\nvar buildSecrets []string\nvar buildNoCache bool\nvar buildProgressOutput string\nvar buildSchemaFile string\nvar buildUseCudaBaseImage string\nvar buildDockerfileFile string\nvar buildUseCogBaseImage bool\nvar buildStrip bool\nvar buildPrecompile bool\nvar configFilename string\n\nconst useCogBaseImageFlagKey = \"use-cog-base-image\"\n\nfunc newBuildCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:   \"build\",\n\t\tShort: \"Build an image from cog.yaml\",\n\t\tLong: `Build a Docker image from the cog.yaml in the current directory.\n\nThe generated image contains your model code, dependencies, and the Cog\nruntime. It can be run locally with 'cog predict' or pushed to a registry\nwith 'cog push'.`,\n\t\tExample: `  # Build with default settings\n  cog build\n\n  # Build and tag the image\n  cog build -t my-model:latest\n\n  # Build without using the cache\n  cog build --no-cache\n\n  # Build with model weights in a separate layer\n  cog build --separate-weights -t my-model:v1`,\n\t\tArgs:    cobra.NoArgs,\n\t\tRunE:    buildCommand,\n\t\tPreRunE: checkMutuallyExclusiveFlags,\n\t}\n\taddBuildProgressOutputFlag(cmd)\n\taddSecretsFlag(cmd)\n\taddNoCacheFlag(cmd)\n\taddSeparateWeightsFlag(cmd)\n\taddSchemaFlag(cmd)\n\taddUseCudaBaseImageFlag(cmd)\n\taddDockerfileFlag(cmd)\n\taddUseCogBaseImageFlag(cmd)\n\taddBuildTimestampFlag(cmd)\n\taddStripFlag(cmd)\n\taddPrecompileFlag(cmd)\n\taddConfigFlag(cmd)\n\tcmd.Flags().StringVarP(&buildTag, \"tag\", \"t\", \"\", \"A name for the built image in the form 'repository:tag'\")\n\treturn cmd\n}\n\nfunc buildCommand(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tsrc, err := model.NewSource(configFilename)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\timageName := src.Config.Image\n\tif buildTag != \"\" {\n\t\timageName = buildTag\n\t}\n\tif imageName == \"\" {\n\t\timageName = config.DockerImageName(src.ProjectDir)\n\t}\n\n\tconsole.Infof(\"Building Docker image from environment in cog.yaml as %s...\", console.Bold(imageName))\n\tconsole.Info(\"\")\n\n\tresolver := model.NewResolver(dockerClient, registry.NewRegistryClient())\n\tm, err := resolver.Build(ctx, src, buildOptionsFromFlags(cmd, imageName, nil))\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tconsole.Info(\"\")\n\tconsole.Successf(\"Image built as %s\", console.Bold(m.ImageRef()))\n\n\treturn nil\n}\n\nfunc addBuildProgressOutputFlag(cmd *cobra.Command) {\n\tdefaultOutput := os.Getenv(\"BUILDKIT_PROGRESS\")\n\tif defaultOutput == \"\" {\n\t\tdefaultOutput = \"auto\"\n\t\tif os.Getenv(\"TERM\") == \"dumb\" {\n\t\t\tdefaultOutput = \"plain\"\n\t\t}\n\t}\n\tcmd.Flags().StringVar(&buildProgressOutput, \"progress\", defaultOutput, \"Set type of build progress output, 'auto' (default), 'tty', 'plain', or 'quiet'\")\n}\n\nfunc addSecretsFlag(cmd *cobra.Command) {\n\tcmd.Flags().StringArrayVar(&buildSecrets, \"secret\", []string{}, \"Secrets to pass to the build environment in the form 'id=foo,src=/path/to/file'\")\n}\n\nfunc addNoCacheFlag(cmd *cobra.Command) {\n\tcmd.Flags().BoolVar(&buildNoCache, \"no-cache\", false, \"Do not use cache when building the image\")\n}\n\nfunc addSeparateWeightsFlag(cmd *cobra.Command) {\n\tcmd.Flags().BoolVar(&buildSeparateWeights, \"separate-weights\", false, \"Separate model weights from code in image layers\")\n}\n\nfunc addSchemaFlag(cmd *cobra.Command) {\n\tcmd.Flags().StringVar(&buildSchemaFile, \"openapi-schema\", \"\", \"Load OpenAPI schema from a file\")\n}\n\nfunc addUseCudaBaseImageFlag(cmd *cobra.Command) {\n\tcmd.Flags().StringVar(&buildUseCudaBaseImage, \"use-cuda-base-image\", \"auto\", \"Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects\")\n}\n\nfunc addDockerfileFlag(cmd *cobra.Command) {\n\tcmd.Flags().StringVar(&buildDockerfileFile, \"dockerfile\", \"\", \"Path to a Dockerfile. If set, cog will use this Dockerfile instead of generating one from cog.yaml\")\n\tcmd.Flags().VisitAll(func(f *pflag.Flag) {\n\t\tif f.Name == \"dockerfile\" {\n\t\t\tf.Hidden = true\n\t\t}\n\t})\n}\n\nfunc addUseCogBaseImageFlag(cmd *cobra.Command) {\n\tcmd.Flags().BoolVar(&buildUseCogBaseImage, useCogBaseImageFlagKey, true, \"Use pre-built Cog base image for faster cold boots\")\n}\n\nfunc addBuildTimestampFlag(cmd *cobra.Command) {\n\tcmd.Flags().Int64Var(&config.BuildSourceEpochTimestamp, \"timestamp\", -1, \"Number of seconds since Epoch to use for the build timestamp; this rewrites the timestamp of each layer. Useful for reproducibility. (`-1` to disable timestamp rewrites)\")\n\t_ = cmd.Flags().MarkHidden(\"timestamp\")\n}\n\nfunc addStripFlag(cmd *cobra.Command) {\n\tconst stripFlag = \"strip\"\n\tcmd.Flags().BoolVar(&buildStrip, stripFlag, false, \"Whether to strip shared libraries for faster inference times\")\n\t_ = cmd.Flags().MarkHidden(stripFlag)\n}\n\nfunc addPrecompileFlag(cmd *cobra.Command) {\n\tconst precompileFlag = \"precompile\"\n\tcmd.Flags().BoolVar(&buildPrecompile, precompileFlag, false, \"Whether to precompile python files for faster load times\")\n\t_ = cmd.Flags().MarkHidden(precompileFlag)\n}\n\nfunc addConfigFlag(cmd *cobra.Command) {\n\tcmd.Flags().StringVarP(&configFilename, \"file\", \"f\", \"cog.yaml\", \"The name of the config file.\")\n}\n\nfunc checkMutuallyExclusiveFlags(cmd *cobra.Command, args []string) error {\n\tflags := []string{useCogBaseImageFlagKey, \"use-cuda-base-image\", \"dockerfile\"}\n\tvar flagsSet []string\n\tfor _, flag := range flags {\n\t\tif cmd.Flag(flag).Changed {\n\t\t\tflagsSet = append(flagsSet, \"--\"+flag)\n\t\t}\n\t}\n\tif len(flagsSet) > 1 {\n\t\treturn fmt.Errorf(\"The flags %s are mutually exclusive: you can only set one of them.\", strings.Join(flagsSet, \" and \"))\n\t}\n\treturn nil\n}\n\nfunc DetermineUseCogBaseImage(cmd *cobra.Command) *bool {\n\tif !cmd.Flags().Changed(useCogBaseImageFlagKey) {\n\t\treturn nil\n\t}\n\tuseCogBaseImage := new(bool)\n\t*useCogBaseImage = buildUseCogBaseImage\n\treturn useCogBaseImage\n}\n\n// buildOptionsFromFlags creates BuildOptions from the current CLI flag values.\n// The imageName and annotations parameters vary by command and must be provided.\nfunc buildOptionsFromFlags(cmd *cobra.Command, imageName string, annotations map[string]string) model.BuildOptions {\n\treturn model.BuildOptions{\n\t\tImageName:        imageName,\n\t\tSecrets:          buildSecrets,\n\t\tNoCache:          buildNoCache,\n\t\tSeparateWeights:  buildSeparateWeights,\n\t\tUseCudaBaseImage: buildUseCudaBaseImage,\n\t\tProgressOutput:   buildProgressOutput,\n\t\tSchemaFile:       buildSchemaFile,\n\t\tDockerfileFile:   buildDockerfileFile,\n\t\tUseCogBaseImage:  DetermineUseCogBaseImage(cmd),\n\t\tStrip:            buildStrip,\n\t\tPrecompile:       buildPrecompile,\n\t\tAnnotations:      annotations,\n\t\tOCIIndex:         model.OCIIndexEnabled(),\n\t}\n}\n"
  },
  {
    "path": "pkg/cli/debug.go",
    "content": "package cli\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/dockerfile\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nvar imageName string\n\nfunc newDebugCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:    \"debug\",\n\t\tHidden: true,\n\t\tShort:  \"Generate a Dockerfile from cog\",\n\t\tRunE:   cmdDockerfile,\n\t}\n\n\taddSeparateWeightsFlag(cmd)\n\taddUseCudaBaseImageFlag(cmd)\n\taddDockerfileFlag(cmd)\n\taddUseCogBaseImageFlag(cmd)\n\taddBuildTimestampFlag(cmd)\n\taddConfigFlag(cmd)\n\tcmd.Flags().StringVarP(&imageName, \"image-name\", \"\", \"\", \"The image name to use for the generated Dockerfile\")\n\n\treturn cmd\n}\n\nfunc cmdDockerfile(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\t// Find the root project directory\n\trootDir, err := config.GetProjectDir(configFilename)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tconfigPath := filepath.Join(rootDir, configFilename)\n\n\tf, err := os.Open(configPath)\n\tif err != nil {\n\t\treturn &config.ParseError{Filename: configFilename, Err: err}\n\t}\n\n\tresult, err := config.Load(f, rootDir)\n\tif err != nil {\n\t\t_ = f.Close()\n\t\treturn err\n\t}\n\n\t_ = f.Close()\n\n\tvar (\n\t\tcfg        = result.Config\n\t\tprojectDir = result.RootDir\n\t)\n\n\t// Display any deprecation warnings\n\tfor _, w := range result.Warnings {\n\t\tconsole.Warnf(\"%s\", w.Error())\n\t}\n\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tclient := registry.NewRegistryClient()\n\tgenerator, err := dockerfile.NewGenerator(cfg, projectDir, configFilename, dockerClient, client, true)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Error creating Dockerfile generator: %w\", err)\n\t}\n\tdefer func() {\n\t\tif err := generator.Cleanup(); err != nil {\n\t\t\tconsole.Warnf(\"Error cleaning up after build: %v\", err)\n\t\t}\n\t}()\n\n\tgenerator.SetUseCudaBaseImage(buildUseCudaBaseImage)\n\tuseCogBaseImage := DetermineUseCogBaseImage(cmd)\n\tif useCogBaseImage != nil {\n\t\tgenerator.SetUseCogBaseImage(*useCogBaseImage)\n\t}\n\n\tif buildSeparateWeights {\n\t\tif imageName == \"\" {\n\t\t\timageName = config.DockerImageName(projectDir)\n\t\t}\n\n\t\tweightsDockerfile, RunnerDockerfile, dockerignore, err := generator.GenerateModelBaseWithSeparateWeights(ctx, imageName)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tconsole.Output(fmt.Sprintf(\"=== Weights Dockerfile contents:\\n%s\\n===\\n\", weightsDockerfile))\n\t\tconsole.Output(fmt.Sprintf(\"=== Runner Dockerfile contents:\\n%s\\n===\\n\", RunnerDockerfile))\n\t\tconsole.Output(fmt.Sprintf(\"=== DockerIgnore contents:\\n%s===\\n\", dockerignore))\n\t} else {\n\t\tdockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights(ctx)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tconsole.Output(dockerfile)\n\t}\n\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/cli/init-templates/base/.dockerignore",
    "content": "# The .dockerignore file excludes files from the container build process.\n#\n# https://docs.docker.com/engine/reference/builder/#dockerignore-file\n\n# Exclude Git files\n**/.git\n**/.github\n**/.gitignore\n\n# Exclude Python tooling\n.python-version\n\n# Exclude Python cache files\n__pycache__\n.mypy_cache\n.pytest_cache\n.ruff_cache\n\n# Exclude Python virtual environment\n/venv\n"
  },
  {
    "path": "pkg/cli/init-templates/base/.github/workflows/push.yaml",
    "content": "name: Push to Replicate\n\non:\n  # Workflow dispatch allows you to manually trigger the workflow from GitHub.com\n  # Go to your repo, click \"Actions\", click \"Push to Replicate\", click \"Run workflow\"\n  workflow_dispatch:\n    inputs:\n      model_name:\n        description: 'Enter the model name, like \"alice/bunny-detector\". If unset, this will default to the value of `image` in cog.yaml.'\n  # # Uncomment these lines to trigger the workflow on every push to the main branch\n  # push:\n  #   branches:\n  #     - main\n\njobs:\n  push_to_replicate:\n    name: Push to Replicate\n\n    # If your model is large, the default GitHub Actions runner may not\n    # have enough disk space. If you need more space you can set up a\n    # bigger runner on GitHub.\n    runs-on: ubuntu-latest\n\n    steps:\n      # This action cleans up disk space to make more room for your\n      # model code, weights, etc.\n      - name: Free disk space\n        uses: jlumbroso/free-disk-space@v1.3.1\n        with:\n          tool-cache: false\n          docker-images: false\n\n      - name: Checkout\n        uses: actions/checkout@v4\n\n      # This action installs Docker buildx and Cog (and optionally CUDA)\n      - name: Setup Cog\n        uses: replicate/setup-cog@v2\n        with:\n          # If you add a CI auth token to your GitHub repository secrets,\n          # the action will authenticate with Replicate automatically so you\n          # can push your model without needing to pass in a token.\n          #\n          # To genereate a CLI auth token, run `cog login` or visit this page\n          # in your browser: https://replicate.com/account/api-token\n          token: ${{ secrets.REPLICATE_CLI_AUTH_TOKEN }}\n\n      # If you trigger the workflow manually, you can specify the model name.\n      # If you leave it blank (or if the workflow is triggered by a push), the\n      # model name will be derived from the `image` value in cog.yaml.\n      - name: Push to Replicate\n        run: |\n          if [ -n \"${{ inputs.model_name }}\" ]; then\n            cog push r8.im/${{ inputs.model_name }}\n          else\n            cog push\n          fi\n"
  },
  {
    "path": "pkg/cli/init-templates/base/cog.yaml",
    "content": "# Configuration for Cog ⚙️\n# Reference: https://cog.run/yaml\n\nbuild:\n  # set to true if your model requires a GPU\n  gpu: false\n\n  # a list of ubuntu apt packages to install\n  # system_packages:\n  #   - \"libgl1-mesa-glx\"\n  #   - \"libglib2.0-0\"\n\n  # python version in the form '3.11' or '3.11.4'\n  python_version: \"3.13\"\n\n  # path to a Python requirements.txt file\n  python_requirements: requirements.txt\n\n  # commands run after the environment is setup\n  # run:\n  #   - \"echo env is ready!\"\n  #   - \"echo another command if needed\"\n\n# predict.py defines how predictions are run on your model\npredict: \"predict.py:Predictor\"\n"
  },
  {
    "path": "pkg/cli/init-templates/base/predict.py",
    "content": "# Prediction interface for Cog ⚙️\n# https://cog.run/python\n\nfrom cog import BasePredictor, Input, Path\n\n\nclass Predictor(BasePredictor):\n    def setup(self) -> None:\n        \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n        # self.model = torch.load(\"./weights.pth\")\n\n    def predict(\n        self,\n        image: Path = Input(description=\"Grayscale input image\"),\n        scale: float = Input(\n            description=\"Factor to scale image by\", ge=0, le=10, default=1.5\n        ),\n    ) -> Path:\n        \"\"\"Run a single prediction on the model\"\"\"\n        # processed_input = preprocess(image)\n        # output = self.model(processed_image, scale)\n        # return postprocess(output)\n"
  },
  {
    "path": "pkg/cli/init-templates/base/requirements.txt",
    "content": "# This is a normal Python requirements.txt file.\n\n# You can add dependencies directly from PyPI:\n# \n# numpy==1.26.4\n# torch==2.2.1\n# torchvision==0.17.1\n\n\n# You can also add Git repos as dependencies, but you'll need to add git to the system_packages list in cog.yaml:\n# \n# build:\n#   system_packages:\n#     - \"git\"\n# \n# Then you can use a URL like this:\n# \n# git+https://github.com/huggingface/transformers\n\n\n# You can also pin Git repos to a specific commit:\n# \n# git+https://github.com/huggingface/transformers@2d1602a\n"
  },
  {
    "path": "pkg/cli/init.go",
    "content": "package cli\n\nimport (\n\t\"embed\"\n\t\"fmt\"\n\t\"io\"\n\t\"net/http\"\n\t\"os\"\n\t\"path\"\n\t\"time\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n\t\"github.com/replicate/cog/pkg/util/files\"\n)\n\n//go:embed init-templates/**/*\nvar initTemplates embed.FS\n\nfunc newInitCommand() *cobra.Command {\n\tvar cmd = &cobra.Command{\n\t\tUse:        \"init\",\n\t\tSuggestFor: []string{\"new\", \"start\"},\n\t\tShort:      \"Configure your project for use with Cog\",\n\t\tLong: `Create a cog.yaml and predict.py in the current directory.\n\nThese files provide a starting template for defining your model's environment\nand prediction interface. Edit them to match your model's requirements.`,\n\t\tExample: `  # Set up a new Cog project in the current directory\n  cog init`,\n\t\tRunE: initCommand,\n\t\tArgs: cobra.MaximumNArgs(0),\n\t}\n\n\treturn cmd\n}\n\nfunc initCommand(cmd *cobra.Command, args []string) error {\n\tconsole.Info(\"Setting up the current directory for use with Cog...\")\n\tconsole.Info(\"\")\n\n\tcwd, err := os.Getwd()\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tinitTemplate := \"base\"\n\n\t// Discover all files in the embedded template directory\n\ttemplateDir := path.Join(\"init-templates\", initTemplate)\n\tentries, err := initTemplates.ReadDir(templateDir)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Error reading template directory: %w\", err)\n\t}\n\n\tfor _, entry := range entries {\n\t\tif entry.IsDir() {\n\t\t\t// Recursively process subdirectories\n\t\t\tif err := processTemplateDirectory(initTemplates, templateDir, entry.Name(), cwd); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\n\t\t// Process individual files\n\t\tif err := processTemplateFile(initTemplates, templateDir, entry.Name(), cwd); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\tconsole.Successf(\"\\nDone! For next steps, check out the docs at https://cog.run/getting-started\")\n\n\treturn nil\n}\n\nfunc processTemplateDirectory(fs embed.FS, templateDir, subDir, cwd string) error {\n\tsubDirPath := path.Join(templateDir, subDir)\n\tentries, err := fs.ReadDir(subDirPath)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Error reading subdirectory %s: %w\", subDirPath, err)\n\t}\n\n\tfor _, entry := range entries {\n\t\tif entry.IsDir() {\n\t\t\t// Recursively process nested subdirectories\n\t\t\tif err := processTemplateDirectory(fs, subDirPath, entry.Name(), cwd); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\n\t\t// Process files in subdirectories\n\t\trelativePath := path.Join(subDir, entry.Name())\n\t\tif err := processTemplateFile(fs, templateDir, relativePath, cwd); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\treturn nil\n}\n\nfunc processTemplateFile(fs embed.FS, templateDir, filename, cwd string) error {\n\tfilePath := path.Join(cwd, filename)\n\tfileExists, err := files.Exists(filePath)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Error checking if %s exists: %w\", filePath, err)\n\t}\n\n\tif fileExists {\n\t\tconsole.Infof(\"Skipped existing %s\", filename)\n\t\treturn nil\n\t}\n\n\tdirPath := path.Dir(filePath)\n\tif err := os.MkdirAll(dirPath, os.ModePerm); err != nil {\n\t\treturn fmt.Errorf(\"Error creating directory %s: %w\", dirPath, err)\n\t}\n\n\tvar content []byte\n\n\t// Special handling for specific template files\n\tswitch filename {\n\tcase \"AGENTS.md\":\n\t\t// Try to download from Replicate docs\n\t\tdownloadedContent, err := downloadAgentsFile()\n\t\tif err != nil {\n\t\t\tconsole.Infof(\"Failed to download AGENTS.md: %v\", err)\n\t\t\tconsole.Infof(\"Using template version instead...\")\n\t\t\t// Fall back to template version\n\t\t\tcontent, err = fs.ReadFile(path.Join(templateDir, filename))\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"Error reading template %s: %w\", filename, err)\n\t\t\t}\n\t\t} else {\n\t\t\tcontent = downloadedContent\n\t\t}\n\tdefault:\n\t\t// Regular template file processing\n\t\tcontent, err = fs.ReadFile(path.Join(templateDir, filename))\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"Error reading %s: %w\", filename, err)\n\t\t}\n\t}\n\n\tconsole.Infof(\"Creating %s\", console.Bold(filename))\n\n\tif err := os.WriteFile(filePath, content, 0o644); err != nil {\n\t\treturn fmt.Errorf(\"Error writing %s: %w\", filePath, err)\n\t}\n\treturn nil\n}\n\nfunc downloadAgentsFile() ([]byte, error) {\n\tconst agentsURL = \"https://replicate.com/docs/reference/cog/llms.txt\"\n\n\tclient := &http.Client{\n\t\tTimeout: 10 * time.Second,\n\t}\n\n\tresp, err := client.Get(agentsURL)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"%w\", err)\n\t}\n\tdefer resp.Body.Close()\n\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn nil, fmt.Errorf(\"HTTP %d\", resp.StatusCode)\n\t}\n\n\tcontent, err := io.ReadAll(resp.Body)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to read response body: %w\", err)\n\t}\n\n\treturn content, nil\n}\n"
  },
  {
    "path": "pkg/cli/init_test.go",
    "content": "package cli\n\nimport (\n\t\"os\"\n\t\"path\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestInit(t *testing.T) {\n\tdir := t.TempDir()\n\n\trequire.NoError(t, os.Chdir(dir))\n\n\terr := initCommand(nil, []string{})\n\trequire.NoError(t, err)\n\n\trequire.FileExists(t, path.Join(dir, \".dockerignore\"))\n\trequire.FileExists(t, path.Join(dir, \"cog.yaml\"))\n\trequire.FileExists(t, path.Join(dir, \"predict.py\"))\n\trequire.FileExists(t, path.Join(dir, \"requirements.txt\"))\n}\n\nfunc TestInitSkipExisting(t *testing.T) {\n\tdir := t.TempDir()\n\n\trequire.NoError(t, os.Chdir(dir))\n\n\t// First run to create files\n\terr := initCommand(nil, []string{})\n\trequire.NoError(t, err)\n\n\trequire.FileExists(t, path.Join(dir, \".dockerignore\"))\n\trequire.FileExists(t, path.Join(dir, \"cog.yaml\"))\n\trequire.FileExists(t, path.Join(dir, \"predict.py\"))\n\n\t// update the file to show that its the same file after the second run\n\trequire.NoError(t, os.WriteFile(path.Join(dir, \"cog.yaml\"), []byte(\"test123\"), 0o644))\n\trequire.NoError(t, os.WriteFile(path.Join(dir, \"predict.py\"), []byte(\"test456\"), 0o644))\n\trequire.NoError(t, os.WriteFile(path.Join(dir, \".dockerignore\"), []byte(\"test789\"), 0o644))\n\n\t// Second run should skip the files that already exist\n\terr = initCommand(nil, []string{})\n\trequire.NoError(t, err)\n\n\trequire.FileExists(t, path.Join(dir, \".dockerignore\"))\n\trequire.FileExists(t, path.Join(dir, \"cog.yaml\"))\n\trequire.FileExists(t, path.Join(dir, \"predict.py\"))\n\n\t// check that the files are the same as the first run\n\tcontent, err := os.ReadFile(path.Join(dir, \"cog.yaml\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, []byte(\"test123\"), content)\n\n\tcontent, err = os.ReadFile(path.Join(dir, \"predict.py\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, []byte(\"test456\"), content)\n\n\tcontent, err = os.ReadFile(path.Join(dir, \".dockerignore\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, []byte(\"test789\"), content)\n}\n"
  },
  {
    "path": "pkg/cli/inspect.go",
    "content": "package cli\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"strings\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/model\"\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// InspectOutput is the structured output for cog inspect --json.\ntype InspectOutput struct {\n\tReference  string           `json:\"reference\"`\n\tType       string           `json:\"type\"` // \"image\" or \"index\"\n\tCogVersion string           `json:\"cogVersion\"`\n\tIndex      *InspectIndex    `json:\"index,omitempty\"`\n\tImage      *InspectManifest `json:\"image,omitempty\"`\n}\n\n// InspectIndex represents an OCI index in inspect output.\ntype InspectIndex struct {\n\tReference string            `json:\"reference\"`\n\tDigest    string            `json:\"digest\"`\n\tMediaType string            `json:\"mediaType\"`\n\tManifests []InspectManifest `json:\"manifests\"`\n}\n\n// InspectManifest represents a manifest entry in inspect output.\ntype InspectManifest struct {\n\tType        string            `json:\"type\"`           // \"image\" or \"weights\"\n\tName        string            `json:\"name,omitempty\"` // weight name from AnnotationWeightName\n\tDigest      string            `json:\"digest\"`\n\tMediaType   string            `json:\"mediaType\"`\n\tSize        int64             `json:\"size\"`\n\tPlatform    string            `json:\"platform,omitempty\"` // \"linux/amd64\"\n\tTarget      string            `json:\"target,omitempty\"`   // weight mount path from AnnotationWeightDest\n\tAnnotations map[string]string `json:\"annotations,omitempty\"`\n\tLayers      []InspectLayer    `json:\"layers\"`\n}\n\n// InspectLayer represents a layer in inspect output.\ntype InspectLayer struct {\n\tDigest    string `json:\"digest\"`\n\tSize      int64  `json:\"size\"`\n\tMediaType string `json:\"mediaType\"`\n}\n\nfunc newInspectCommand() *cobra.Command {\n\tvar (\n\t\tlocalOnly  bool\n\t\tremoteOnly bool\n\t\tjsonOutput bool\n\t\trawOutput  bool\n\t)\n\n\tcmd := &cobra.Command{\n\t\tUse:    \"inspect \",\n\t\tShort:  \"Inspect a model image or OCI index\",\n\t\tArgs:   cobra.ExactArgs(1),\n\t\tHidden: true,\n\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\treturn inspectCommand(cmd, args, localOnly, remoteOnly, jsonOutput, rawOutput)\n\t\t},\n\t}\n\n\tcmd.Flags().BoolVar(&localOnly, \"local\", false, \"Only inspect local docker daemon\")\n\tcmd.Flags().BoolVar(&remoteOnly, \"remote\", false, \"Only inspect remote registry\")\n\tcmd.Flags().BoolVar(&jsonOutput, \"json\", false, \"Output as JSON\")\n\tcmd.Flags().BoolVar(&rawOutput, \"raw\", false, \"Output raw JSON fragments (one per line)\")\n\n\treturn cmd\n}\n\nfunc inspectCommand(cmd *cobra.Command, args []string, localOnly, remoteOnly, jsonOutput, rawOutput bool) error {\n\tctx := cmd.Context()\n\n\tref, err := model.ParseRef(args[0])\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tregClient := registry.NewRegistryClient()\n\tresolver := model.NewResolver(dockerClient, regClient)\n\n\t// Build resolve options\n\tvar opts []model.Option\n\tswitch {\n\tcase localOnly:\n\t\topts = append(opts, model.LocalOnly())\n\tcase remoteOnly:\n\t\topts = append(opts, model.RemoteOnly())\n\t}\n\n\tm, err := resolver.Inspect(ctx, ref, opts...)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\t// Build output\n\tout, err := buildInspectOutput(ctx, ref.String(), m, regClient)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tswitch {\n\tcase rawOutput:\n\t\treturn streamRaw(ctx, ref.String(), m, regClient)\n\tcase jsonOutput:\n\t\tenc := json.NewEncoder(os.Stdout)\n\t\tenc.SetIndent(\"\", \"  \")\n\t\treturn enc.Encode(out)\n\tdefault:\n\t\tprintInspectText(out)\n\t\treturn nil\n\t}\n}\n\nfunc buildInspectOutput(ctx context.Context, reference string, m *model.Model, reg registry.Client) (*InspectOutput, error) {\n\tout := &InspectOutput{\n\t\tReference:  reference,\n\t\tCogVersion: m.CogVersion,\n\t}\n\n\tif m.Index != nil {\n\t\tout.Type = \"index\"\n\t\tidx := &InspectIndex{\n\t\t\tReference: m.Index.Reference,\n\t\t\tDigest:    m.Index.Digest,\n\t\t\tMediaType: m.Index.MediaType,\n\t\t}\n\n\t\tfor _, im := range m.Index.Manifests {\n\t\t\tmanifest := buildManifestEntry(im)\n\n\t\t\t// Try to fetch layers from registry\n\t\t\tlayers, err := fetchLayers(ctx, reference, im.Digest, reg)\n\t\t\tif err == nil {\n\t\t\t\tmanifest.Layers = layers\n\t\t\t}\n\n\t\t\tidx.Manifests = append(idx.Manifests, manifest)\n\t\t}\n\n\t\tout.Index = idx\n\t} else {\n\t\tout.Type = \"image\"\n\t\tif m.Image != nil {\n\t\t\tmanifest := &InspectManifest{\n\t\t\t\tType:   \"image\",\n\t\t\t\tDigest: m.Image.Digest,\n\t\t\t}\n\t\t\tif m.Image.Platform != nil {\n\t\t\t\tparts := []string{m.Image.Platform.OS, m.Image.Platform.Architecture}\n\t\t\t\tif m.Image.Platform.Variant != \"\" {\n\t\t\t\t\tparts = append(parts, m.Image.Platform.Variant)\n\t\t\t\t}\n\t\t\t\tmanifest.Platform = strings.Join(parts, \"/\")\n\t\t\t}\n\n\t\t\t// Try to fetch layers\n\t\t\tif m.Image.Digest != \"\" {\n\t\t\t\tlayers, err := fetchLayers(ctx, reference, m.Image.Digest, reg)\n\t\t\t\tif err == nil {\n\t\t\t\t\tmanifest.Layers = layers\n\t\t\t\t}\n\t\t\t}\n\n\t\t\tout.Image = manifest\n\t\t}\n\t}\n\n\treturn out, nil\n}\n\nfunc buildManifestEntry(im model.IndexManifest) InspectManifest {\n\tmanifest := InspectManifest{\n\t\tDigest:      im.Digest,\n\t\tMediaType:   im.MediaType,\n\t\tSize:        im.Size,\n\t\tAnnotations: im.Annotations,\n\t}\n\n\tswitch im.Type {\n\tcase model.ManifestTypeWeights:\n\t\tmanifest.Type = \"weights\"\n\t\tmanifest.Name = im.Annotations[model.AnnotationWeightName]\n\t\tmanifest.Target = im.Annotations[model.AnnotationWeightDest]\n\tdefault:\n\t\tmanifest.Type = \"image\"\n\t\tif im.Platform != nil {\n\t\t\tparts := []string{im.Platform.OS, im.Platform.Architecture}\n\t\t\tif im.Platform.Variant != \"\" {\n\t\t\t\tparts = append(parts, im.Platform.Variant)\n\t\t\t}\n\t\t\tmanifest.Platform = strings.Join(parts, \"/\")\n\t\t}\n\t}\n\n\treturn manifest\n}\n\nfunc fetchLayers(ctx context.Context, reference, digest string, reg registry.Client) ([]InspectLayer, error) {\n\t// Build a digest reference from the repo\n\tref, err := model.ParseRef(reference)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdigestRef := ref.Ref.Context().String() + \"@\" + digest\n\n\timg, err := reg.GetImage(ctx, digestRef, nil)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tmanifest, err := img.Manifest()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tvar layers []InspectLayer\n\tfor _, l := range manifest.Layers {\n\t\tlayers = append(layers, InspectLayer{\n\t\t\tDigest:    l.Digest.String(),\n\t\t\tSize:      l.Size,\n\t\t\tMediaType: string(l.MediaType),\n\t\t})\n\t}\n\n\treturn layers, nil\n}\n\ntype rawStep struct {\n\tStep     string `json:\"step\"`\n\tData     any    `json:\"data,omitempty\"`\n\tManifest any    `json:\"manifest,omitempty\"`\n}\n\nfunc streamRaw(ctx context.Context, reference string, m *model.Model, reg registry.Client) error {\n\tenc := json.NewEncoder(os.Stdout)\n\n\t// Step 1: resolve\n\t_ = enc.Encode(rawStep{\n\t\tStep: \"resolve\",\n\t\tData: map[string]any{\n\t\t\t\"reference\":  reference,\n\t\t\t\"cogVersion\": m.CogVersion,\n\t\t\t\"type\": func() string {\n\t\t\t\tif m.Index != nil {\n\t\t\t\t\treturn \"index\"\n\t\t\t\t}\n\t\t\t\treturn \"image\"\n\t\t\t}(),\n\t\t},\n\t})\n\n\tif m.Index != nil {\n\t\t// Step 2: index\n\t\t_ = enc.Encode(rawStep{\n\t\t\tStep: \"index\",\n\t\t\tData: map[string]any{\n\t\t\t\t\"digest\":    m.Index.Digest,\n\t\t\t\t\"mediaType\": m.Index.MediaType,\n\t\t\t\t\"count\":     len(m.Index.Manifests),\n\t\t\t},\n\t\t})\n\n\t\t// Step 3: per-child manifests\n\t\tfor _, im := range m.Index.Manifests {\n\t\t\tentry := buildManifestEntry(im)\n\n\t\t\tref, err := model.ParseRef(reference)\n\t\t\tif err == nil {\n\t\t\t\tdigestRef := ref.Ref.Context().String() + \"@\" + im.Digest\n\t\t\t\timg, err := reg.GetImage(ctx, digestRef, nil)\n\t\t\t\tif err == nil {\n\t\t\t\t\trawManifest, err := img.RawManifest()\n\t\t\t\t\tif err == nil {\n\t\t\t\t\t\tvar parsed any\n\t\t\t\t\t\tif jsonErr := json.Unmarshal(rawManifest, &parsed); jsonErr == nil {\n\t\t\t\t\t\t\t_ = enc.Encode(rawStep{\n\t\t\t\t\t\t\t\tStep:     \"manifest\",\n\t\t\t\t\t\t\t\tData:     entry,\n\t\t\t\t\t\t\t\tManifest: parsed,\n\t\t\t\t\t\t\t})\n\t\t\t\t\t\t\tcontinue\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\n\t\t\t// Fallback: output without raw manifest\n\t\t\t_ = enc.Encode(rawStep{\n\t\t\t\tStep: \"manifest\",\n\t\t\t\tData: entry,\n\t\t\t})\n\t\t}\n\t}\n\n\t// Final step: model summary\n\t_ = enc.Encode(rawStep{\n\t\tStep: \"model\",\n\t\tData: map[string]any{\n\t\t\t\"reference\":  reference,\n\t\t\t\"cogVersion\": m.CogVersion,\n\t\t},\n\t})\n\n\treturn nil\n}\n\nfunc printInspectText(out *InspectOutput) {\n\tfmt.Printf(\"Model: %s\\n\", out.Reference)\n\tif out.Type == \"index\" {\n\t\tfmt.Println(\"Type:  Model Bundle (OCI Index)\")\n\t} else {\n\t\tfmt.Println(\"Type:  Image\")\n\t}\n\tfmt.Printf(\"Cog:   %s\\n\", out.CogVersion)\n\tfmt.Println()\n\n\tif out.Index != nil {\n\t\t// Build the digest reference: repo@sha256:...\n\t\tdigestRef := out.Index.Digest\n\t\tif out.Index.Reference != \"\" && out.Index.Digest != \"\" {\n\t\t\t// Extract repo from the reference (strip tag/digest)\n\t\t\trepo := out.Index.Reference\n\t\t\tif idx := strings.LastIndex(repo, \":\"); idx != -1 {\n\t\t\t\t// Only strip if it looks like a tag (no @)\n\t\t\t\tif !strings.Contains(repo[idx:], \"@\") {\n\t\t\t\t\trepo = repo[:idx]\n\t\t\t\t}\n\t\t\t}\n\t\t\tdigestRef = repo + \"@\" + out.Index.Digest\n\t\t}\n\t\tfmt.Printf(\"Index: %s\\n\", digestRef)\n\t\tfmt.Printf(\"  Tag:       %s\\n\", out.Reference)\n\t\tfmt.Printf(\"  Digest:    %s\\n\", out.Index.Digest)\n\t\tfmt.Printf(\"  MediaType: %s\\n\", out.Index.MediaType)\n\t\tfmt.Printf(\"  Manifests: %d\\n\", len(out.Index.Manifests))\n\t\tfmt.Println()\n\n\t\tfor _, m := range out.Index.Manifests {\n\t\t\tprintManifestText(m, \"  \")\n\t\t\tfmt.Println()\n\t\t}\n\t} else if out.Image != nil {\n\t\tprintManifestText(*out.Image, \"\")\n\t}\n}\n\nfunc printManifestText(m InspectManifest, indent string) {\n\tif m.Type == \"weights\" {\n\t\tname := m.Name\n\t\tif name == \"\" {\n\t\t\tname = \"(unnamed)\"\n\t\t}\n\t\tfmt.Printf(\"%s[weights] %s\\n\", indent, name)\n\t} else {\n\t\tplatform := m.Platform\n\t\tif platform == \"\" {\n\t\t\tplatform = \"(unknown)\"\n\t\t}\n\t\tfmt.Printf(\"%s[image] %s\\n\", indent, platform)\n\t}\n\n\tfmt.Printf(\"%s  Digest: %s\\n\", indent, m.Digest)\n\n\t// Show manifest size + total layer size if layers are available\n\tif len(m.Layers) > 0 {\n\t\tvar layerTotal int64\n\t\tfor _, l := range m.Layers {\n\t\t\tlayerTotal += l.Size\n\t\t}\n\t\tfmt.Printf(\"%s  Size:   %s (Layers: %s)\\n\", indent, formatSize(m.Size), formatSize(layerTotal))\n\t} else {\n\t\tfmt.Printf(\"%s  Size:   %s\\n\", indent, formatSize(m.Size))\n\t}\n\n\tif m.Target != \"\" {\n\t\tfmt.Printf(\"%s  Target: %s\\n\", indent, m.Target)\n\t}\n\n\tif m.MediaType != \"\" {\n\t\tfmt.Printf(\"%s  Type:   %s\\n\", indent, m.MediaType)\n\t}\n\n\tif len(m.Layers) > 0 {\n\t\tfmt.Printf(\"%s  Layers: %d\\n\", indent, len(m.Layers))\n\t\tfor _, l := range m.Layers {\n\t\t\tfmt.Printf(\"%s    %s  %s  %s\\n\", indent, l.Digest, formatSize(l.Size), l.MediaType)\n\t\t}\n\t}\n}\n"
  },
  {
    "path": "pkg/cli/login.go",
    "content": "package cli\n\nimport (\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/provider\"\n\t\"github.com/replicate/cog/pkg/provider/setup\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc newLoginCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:        \"login\",\n\t\tSuggestFor: []string{\"auth\", \"authenticate\", \"authorize\"},\n\t\tShort:      \"Log in to a container registry\",\n\t\tLong: `Log in to a container registry.\n\nFor Replicate's registry (r8.im), this command handles authentication\nthrough Replicate's token-based flow.\n\nFor other registries, this command prompts for username and password,\nthen stores credentials using Docker's credential system.`,\n\t\tRunE: login,\n\t\tArgs: cobra.MaximumNArgs(0),\n\t}\n\n\tcmd.Flags().Bool(\"token-stdin\", false, \"Pass login token on stdin instead of opening a browser. You can find your Replicate login token at https://replicate.com/auth/token\")\n\n\treturn cmd\n}\n\nfunc login(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\t// Initialize the provider registry\n\tsetup.Init()\n\n\t// Use global registry host (can be set via --registry flag or COG_REGISTRY_HOST env var)\n\tregistryHost := global.ReplicateRegistryHost\n\n\ttokenStdin, err := cmd.Flags().GetBool(\"token-stdin\")\n\tif err != nil {\n\t\treturn err\n\t}\n\n\t// Look up the provider for this registry\n\tp := provider.DefaultRegistry().ForHost(registryHost)\n\tif p == nil {\n\t\t// This shouldn't happen since GenericProvider matches everything\n\t\tconsole.Warnf(\"No provider found for registry '%s'.\", registryHost)\n\t\tconsole.Infof(\"Please use 'docker login %s' to authenticate.\", registryHost)\n\t\treturn nil\n\t}\n\n\treturn p.Login(ctx, provider.LoginOptions{\n\t\tTokenStdin: tokenStdin,\n\t\tHost:       registryHost,\n\t})\n}\n"
  },
  {
    "path": "pkg/cli/predict.go",
    "content": "package cli\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/base64\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n\t\"os/signal\"\n\t\"path\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"syscall\"\n\t\"time\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\t\"github.com/mitchellh/go-homedir\"\n\t\"github.com/spf13/cobra\"\n\t\"golang.org/x/sys/unix\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/model\"\n\tr8_path \"github.com/replicate/cog/pkg/path\"\n\t\"github.com/replicate/cog/pkg/predict\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n\t\"github.com/replicate/cog/pkg/util/files\"\n\t\"github.com/replicate/cog/pkg/util/mime\"\n)\n\nconst StdinPath = \"-\"\n\nvar (\n\tenvFlags             []string\n\tinputFlags           []string\n\toutPath              string\n\tsetupTimeout         uint32\n\tuseReplicateAPIToken bool\n\tinputJSON            string\n)\n\nfunc newPredictCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:   \"predict [image]\",\n\t\tShort: \"Run a prediction\",\n\t\tLong: `Run a prediction.\n\nIf 'image' is passed, it will run the prediction on that Docker image.\nIt must be an image that has been built by Cog.\n\nOtherwise, it will build the model in the current directory and run\nthe prediction on that.`,\n\t\tExample: `  # Run a prediction with named inputs\n  cog predict -i prompt=\"a photo of a cat\"\n\n  # Pass a file as input\n  cog predict -i image=@photo.jpg\n\n  # Save output to a file\n  cog predict -i image=@input.jpg -o output.png\n\n  # Pass multiple inputs\n  cog predict -i prompt=\"sunset\" -i width=1024 -i height=768\n\n  # Run against a pre-built image\n  cog predict r8.im/your-username/my-model -i prompt=\"hello\"\n\n  # Pass inputs as JSON\n  echo '{\"prompt\": \"a cat\"}' | cog predict --json @-`,\n\t\tRunE:       cmdPredict,\n\t\tArgs:       cobra.MaximumNArgs(1),\n\t\tSuggestFor: []string{\"infer\"},\n\t}\n\n\taddUseCudaBaseImageFlag(cmd)\n\taddUseCogBaseImageFlag(cmd)\n\taddBuildProgressOutputFlag(cmd)\n\taddDockerfileFlag(cmd)\n\taddGpusFlag(cmd)\n\taddSetupTimeoutFlag(cmd)\n\taddConfigFlag(cmd)\n\n\tcmd.Flags().StringArrayVarP(&inputFlags, \"input\", \"i\", []string{}, \"Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg\")\n\tcmd.Flags().StringVarP(&outPath, \"output\", \"o\", \"\", \"Output path\")\n\tcmd.Flags().StringArrayVarP(&envFlags, \"env\", \"e\", []string{}, \"Environment variables, in the form name=value\")\n\tcmd.Flags().BoolVar(&useReplicateAPIToken, \"use-replicate-token\", false, \"Pass REPLICATE_API_TOKEN from local environment into the model context\")\n\tcmd.Flags().StringVar(&inputJSON, \"json\", \"\", \"Pass inputs as JSON object, read from file (@inputs.json) or via stdin (@-)\")\n\n\treturn cmd\n}\n\nfunc readStdin() (string, error) {\n\t// Read from stdin\n\tdata, err := io.ReadAll(os.Stdin)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"Failed to read JSON from stdin: %w\", err)\n\t}\n\treturn string(data), nil\n}\n\nfunc parseJSONInput(jsonInput string) (map[string]any, error) {\n\tvar jsonStr string\n\n\tswitch {\n\tcase strings.HasPrefix(jsonInput, \"@\"):\n\t\t// Read from file or stdin\n\t\tsource := jsonInput[1:]\n\n\t\tif source == StdinPath {\n\t\t\tjsonStdinStr, err := readStdin()\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tjsonStr = jsonStdinStr\n\t\t} else {\n\t\t\t// Read from file\n\t\t\tdata, err := os.ReadFile(source)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"Failed to read JSON from file %q: %w\", source, err)\n\t\t\t}\n\t\t\tjsonStr = string(data)\n\t\t}\n\tcase jsonInput == StdinPath:\n\t\tjsonStdinStr, err := readStdin()\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tjsonStr = jsonStdinStr\n\tdefault:\n\t\t// Direct JSON string\n\t\tjsonStr = jsonInput\n\t}\n\n\tvar inputs map[string]any\n\tif err := json.Unmarshal([]byte(jsonStr), &inputs); err != nil {\n\t\treturn nil, fmt.Errorf(\"Failed to parse JSON: %w\", err)\n\t}\n\n\treturn inputs, nil\n}\n\nfunc transformPathsToBase64URLs(inputs map[string]any) (map[string]any, error) {\n\tresult := make(map[string]any)\n\n\tfor key, value := range inputs {\n\t\tif strValue, ok := value.(string); ok && strings.HasPrefix(strValue, \"@\") {\n\t\t\t// This is a file path, convert to base64 data URL\n\t\t\tfilePath := strValue[1:]\n\n\t\t\t// Read file\n\t\t\tdata, err := os.ReadFile(filePath)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"Failed to read file %q: %w\", filePath, err)\n\t\t\t}\n\n\t\t\t// Get MIME type\n\t\t\tmimeType := mime.TypeByExtension(filepath.Ext(filePath))\n\t\t\tif mimeType == \"\" {\n\t\t\t\tmimeType = \"application/octet-stream\"\n\t\t\t}\n\n\t\t\t// Create base64 data URL\n\t\t\tbase64Data := base64.StdEncoding.EncodeToString(data)\n\t\t\tdataURL := fmt.Sprintf(\"data:%s;base64,%s\", mimeType, base64Data)\n\n\t\t\tresult[key] = dataURL\n\t\t} else {\n\t\t\tresult[key] = value\n\t\t}\n\t}\n\n\treturn result, nil\n}\n\nfunc cmdPredict(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\timageName := \"\"\n\tvolumes := []command.Volume{}\n\tgpus := gpusFlag\n\n\tresolver := model.NewResolver(dockerClient, registry.NewRegistryClient())\n\n\tif len(args) == 0 {\n\t\t// Build image\n\t\tsrc, err := model.NewSource(configFilename)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tconsole.Info(\"Building Docker image from environment in cog.yaml...\")\n\t\tconsole.Info(\"\")\n\t\tm, err := resolver.Build(ctx, src, serveBuildOptions(cmd))\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\timageName = m.ImageRef()\n\n\t\t// ExcludeSource build doesn't have /src in it, so mount as volume\n\t\tvolumes = append(volumes, command.Volume{\n\t\t\tSource:      src.ProjectDir,\n\t\t\tDestination: \"/src\",\n\t\t})\n\n\t\tif gpus == \"\" && m.HasGPU() {\n\t\t\tgpus = \"all\"\n\t\t}\n\t} else {\n\t\t// Use existing image\n\t\timageName = args[0]\n\n\t\t// If the image name contains '=', then it's probably a mistake\n\t\tif strings.Contains(imageName, \"=\") {\n\t\t\treturn fmt.Errorf(\"Invalid image name '%s'. Did you forget `-i`?\", imageName)\n\t\t}\n\n\t\t// Pull the image (if needed) and validate it's a Cog model\n\t\tref, err := model.ParseRef(imageName)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tm, err := resolver.Pull(ctx, ref)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tif gpus == \"\" && m.HasGPU() {\n\t\t\tgpus = \"all\"\n\t\t}\n\t}\n\n\tconsole.Info(\"\")\n\tconsole.Info(\"Starting Docker image and running setup()...\")\n\n\t// Automatically propagate RUST_LOG for Rust coglet debugging\n\tenv := envFlags\n\tif rustLog := os.Getenv(\"RUST_LOG\"); rustLog != \"\" {\n\t\tenv = append(env, \"RUST_LOG=\"+rustLog)\n\t}\n\n\tpredictor, err := predict.NewPredictor(ctx, command.RunOptions{\n\t\tGPUs:    gpus,\n\t\tImage:   imageName,\n\t\tVolumes: volumes,\n\t\tEnv:     env,\n\t}, false, dockerClient)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tgo func() {\n\t\tcaptureSignal := make(chan os.Signal, 1)\n\t\tsignal.Notify(captureSignal, syscall.SIGINT)\n\n\t\t<-captureSignal\n\n\t\tconsole.Info(\"Stopping container...\")\n\t\tif err := predictor.Stop(ctx); err != nil {\n\t\t\tconsole.Warnf(\"Failed to stop container: %s\", err)\n\t\t}\n\t}()\n\n\ttimeout := time.Duration(setupTimeout) * time.Second\n\tif err := predictor.Start(ctx, os.Stderr, timeout); err != nil {\n\t\t// Only retry if we're using a GPU but the user didn't explicitly select a GPU with --gpus\n\t\t// If the user specified the wrong GPU, they are explicitly selecting a GPU and they'll want to hear about it\n\t\tif gpus == \"all\" && errors.Is(err, docker.ErrMissingDeviceDriver) {\n\t\t\tconsole.Info(\"Missing device driver, re-trying without GPU\")\n\n\t\t\t_ = predictor.Stop(ctx)\n\t\t\tpredictor, err = predict.NewPredictor(ctx, command.RunOptions{\n\t\t\t\tImage:   imageName,\n\t\t\t\tVolumes: volumes,\n\t\t\t\tEnv:     env,\n\t\t\t}, false, dockerClient)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\n\t\t\tif err := predictor.Start(ctx, os.Stderr, timeout); err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t} else {\n\t\t\treturn err\n\t\t}\n\t}\n\n\t// FIXME: will not run on signal\n\tdefer func() {\n\t\tconsole.Debugf(\"Stopping container...\")\n\t\t// use background context to ensure stop signal is still sent after root context is canceled\n\t\tif err := predictor.Stop(context.Background()); err != nil {\n\t\t\tconsole.Warnf(\"Failed to stop container: %s\", err)\n\t\t}\n\t}()\n\n\tif inputJSON != \"\" {\n\t\tif len(inputFlags) > 0 {\n\t\t\treturn fmt.Errorf(\"Must use one of --json or --input to provide model inputs\")\n\t\t}\n\n\t\treturn predictJSONInputs(*predictor, inputJSON, outPath, false)\n\t}\n\treturn predictIndividualInputs(*predictor, inputFlags, outPath, false)\n}\n\nfunc isURI(ref *openapi3.Schema) bool {\n\treturn ref != nil && ref.Type.Is(\"string\") && ref.Format == \"uri\"\n}\n\nfunc predictJSONInputs(predictor predict.Predictor, jsonInput string, outputPath string, isTrain bool) error {\n\tjsonInputs, err := parseJSONInput(jsonInput)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\ttransformedInputs, err := transformPathsToBase64URLs(jsonInputs)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\t// Convert to predict.Inputs format\n\tinputs := make(predict.Inputs)\n\tfor key, value := range transformedInputs {\n\t\tif strValue, ok := value.(string); ok {\n\t\t\tinputs[key] = predict.Input{String: &strValue}\n\t\t} else {\n\t\t\t// For non-string values, marshal to JSON\n\t\t\tjsonBytes, err := json.Marshal(value)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"Failed to marshal input %q to JSON: %w\", key, err)\n\t\t\t}\n\t\t\tjsonRaw := json.RawMessage(jsonBytes)\n\t\t\tinputs[key] = predict.Input{Json: &jsonRaw}\n\t\t}\n\t}\n\n\treturn runPrediction(predictor, inputs, outputPath, isTrain, true)\n}\n\nfunc predictIndividualInputs(predictor predict.Predictor, inputFlags []string, outputPath string, isTrain bool) error {\n\tschema, err := predictor.GetSchema()\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tinputs, err := parseInputFlags(inputFlags, schema, isTrain)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treturn runPrediction(predictor, inputs, outputPath, isTrain, false)\n}\n\nfunc runPrediction(predictor predict.Predictor, inputs predict.Inputs, outputPath string, isTrain bool, needsJSON bool) error {\n\tif isTrain {\n\t\tconsole.Info(\"Running training...\")\n\t} else {\n\t\tconsole.Info(\"Running prediction...\")\n\t}\n\tconsole.Info(\"\")\n\n\t// Generate output depending on type in schema\n\turl := \"/predictions\"\n\tif isTrain {\n\t\turl = \"/trainings\"\n\t}\n\n\twriteOutputToDisk := outputPath != \"\"\n\tfallbackPath := \"output\"\n\tif needsJSON {\n\t\tfallbackPath = \"output.json\"\n\t}\n\n\toutputPath, err := ensureOutputWriteable(strings.TrimPrefix(outputPath, \"@\"), fallbackPath)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Output path is not writable: %w\", err)\n\t}\n\n\tif needsJSON && !strings.HasSuffix(outputPath, \".json\") {\n\t\tconsole.Warnf(\"--output value does not have a .json suffix: %s\", path.Base(outputPath))\n\t}\n\n\tcontext := predict.RequestContext{}\n\n\tif useReplicateAPIToken {\n\t\tcontext.ReplicateAPIToken = os.Getenv(\"REPLICATE_API_TOKEN\")\n\t\tif context.ReplicateAPIToken == \"\" {\n\t\t\treturn fmt.Errorf(\"Failed to find REPLICATE_API_TOKEN in the current environment when called with --use-replicate-token\")\n\t\t}\n\t}\n\n\tprediction, err := predictor.Predict(inputs, context)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Failed to run prediction: %w\", err)\n\t}\n\n\tschema, err := predictor.GetSchema()\n\tif err != nil {\n\t\treturn err\n\t}\n\n\t// Safely extract output schema with nil checks to avoid panics on malformed schemas\n\tvar outputSchema *openapi3.Schema\n\tif pathItem := schema.Paths.Value(url); pathItem != nil {\n\t\tif pathItem.Post != nil {\n\t\t\tif resp := pathItem.Post.Responses.Value(\"200\"); resp != nil && resp.Value != nil {\n\t\t\t\tif content, ok := resp.Value.Content[\"application/json\"]; ok && content.Schema != nil {\n\t\t\t\t\tif content.Schema.Value != nil {\n\t\t\t\t\t\tif outputProp, ok := content.Schema.Value.Properties[\"output\"]; ok && outputProp != nil {\n\t\t\t\t\t\t\toutputSchema = outputProp.Value\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\tif outputSchema == nil {\n\t\treturn fmt.Errorf(\"invalid OpenAPI schema: missing output definition for %s\", url)\n\t}\n\n\tfileOutputPath := outputPath\n\tif needsJSON {\n\t\t// Strip the suffix when in JSON mode.\n\t\tfileOutputPath = r8_path.TrimExt(fileOutputPath)\n\t}\n\n\tif prediction.Status == \"succeeded\" && prediction.Output != nil {\n\t\ttransformed, err := processFileOutputs(*prediction.Output, outputSchema, fileOutputPath)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tprediction.Output = &transformed\n\t}\n\n\tif needsJSON {\n\t\trawJSON, err := json.Marshal(prediction)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"Failed to encode prediction output as JSON: %w\", err)\n\t\t}\n\t\tvar indentedJSON bytes.Buffer\n\t\tif err := json.Indent(&indentedJSON, rawJSON, \"\", \"  \"); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tif writeOutputToDisk {\n\t\t\tpath, err := files.WriteFile(indentedJSON.Bytes(), outputPath)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"Failed to write output: %w\", err)\n\t\t\t}\n\t\t\tconsole.Infof(\"Written output to: %s\", path)\n\t\t} else {\n\t\t\tconsole.Output(indentedJSON.String())\n\t\t}\n\n\t\t// Exit with non-zero code if the prediction has failed.\n\t\tif prediction.Status != \"succeeded\" {\n\t\t\tos.Exit(1)\n\t\t}\n\n\t\treturn nil\n\t}\n\n\tif prediction.Status != \"succeeded\" {\n\t\treturn fmt.Errorf(\"Prediction failed with status %q: %s\", prediction.Status, prediction.Error)\n\t}\n\n\tif prediction.Output == nil {\n\t\tconsole.Warn(\"No output generated\")\n\t\treturn nil\n\t}\n\n\t// Handle default presentation of output types.\n\t// 1. For Path and list[Path] do nothing. We already print info for each file write.\n\t// 2. For everything else we want to print the raw value.\n\tswitch {\n\tcase isURI(outputSchema):\n\t\treturn nil\n\tcase outputSchema.Type.Is(\"array\") && isURI(outputSchema.Items.Value):\n\t\treturn nil\n\tcase outputSchema.Type.Is(\"string\"):\n\t\t// Output the raw string.\n\t\ts, ok := (*prediction.Output).(string)\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"Failed to convert prediction output to string\")\n\t\t}\n\n\t\tif writeOutputToDisk {\n\t\t\tpath, err := files.WriteFile([]byte(s), outputPath)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"Failed to write output: %w\", err)\n\t\t\t}\n\t\t\tconsole.Infof(\"Written output to: %s\", path)\n\t\t} else {\n\t\t\tconsole.Output(s)\n\t\t}\n\n\t\treturn nil\n\tdefault:\n\t\t// Treat everything else as JSON -- ints, floats, bools will all be presented\n\t\t// as raw values. Lists and objects will be pretty printed JSON.\n\t\toutput, err := prettyJSONMarshal(prediction.Output)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\t// No special handling for needsJSON here.\n\t\tif writeOutputToDisk {\n\t\t\tpath, err := files.WriteFile(output, outputPath)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"Failed to write output: %w\", err)\n\t\t\t}\n\t\t\tconsole.Infof(\"Written output to: %s\", path)\n\t\t} else {\n\t\t\tconsole.Output(string(output))\n\t\t}\n\n\t\treturn nil\n\t}\n}\n\n// Ensures the path (or fallback) provided is writable. Returns path, error\nfunc ensureOutputWriteable(outputPath string, fallbackPath string) (string, error) {\n\t// If no outputPath is provided use fallback path and track.\n\tusingFallback := false\n\tif outputPath == \"\" {\n\t\toutputPath = fallbackPath\n\t\tusingFallback = true\n\t}\n\n\toutputPath, err := homedir.Expand(outputPath)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tstat, err := os.Stat(outputPath)\n\n\t// If the file doesn't exist, use the parent directory with given filename.\n\tif os.IsNotExist(err) {\n\t\tif err = unix.Access(path.Dir(outputPath), unix.W_OK); err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Output directory is not writable: %s\", path.Dir(outputPath))\n\t\t}\n\t\treturn outputPath, nil\n\t} else if err != nil {\n\t\treturn \"\", fmt.Errorf(\"Unexpected error checking output path: %w\", err)\n\t}\n\n\t// If a directory was provided, use that with the fallback filename\n\tif stat.IsDir() {\n\t\t// If the fallback path already exists as a directory error.\n\t\tif usingFallback {\n\t\t\treturn \"\", fmt.Errorf(\"Default output name %q conflicts with directory, provide --output\", outputPath)\n\t\t}\n\t\terr := unix.Access(outputPath, unix.W_OK)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\treturn path.Join(outputPath, path.Base(fallbackPath)), nil\n\t}\n\n\tif err = unix.Access(outputPath, unix.W_OK); err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn outputPath, nil\n}\n\nfunc prettyJSONMarshal(v any) ([]byte, error) {\n\traw, err := json.Marshal(v)\n\tif err != nil {\n\t\treturn []byte(\"\"), fmt.Errorf(\"Failed to encode JSON: %w\", err)\n\t}\n\tvar formatted bytes.Buffer\n\tif err := json.Indent(&formatted, raw, \"\", \"  \"); err != nil {\n\t\treturn []byte(\"\"), err\n\t}\n\treturn formatted.Bytes(), nil\n}\n\nfunc processFileOutputs(output any, schema *openapi3.Schema, destination string) (any, error) {\n\t// TODO: This doesn't currently support arbitrary objects.\n\tswitch {\n\tcase isURI(schema):\n\t\toutputStr, ok := output.(string)\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"Failed to convert prediction output to string: %v\", output)\n\t\t}\n\n\t\tpath, err := files.WriteDataURLToFile(outputStr, destination)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"Failed to write output: %w\", err)\n\t\t}\n\t\tconsole.Infof(\"Written output to: %s\", path)\n\n\t\treturn any(path), nil\n\tcase schema.Type.Is(\"array\") && isURI(schema.Items.Value):\n\t\toutputs, ok := (output).([]any)\n\t\tif !ok {\n\t\t\treturn nil, fmt.Errorf(\"Failed to decode output: %v\", output)\n\t\t}\n\n\t\tclone := []any{}\n\t\tfor i, output := range outputs {\n\t\t\titemDestination := fmt.Sprintf(\"%s.%d%s\", r8_path.TrimExt(destination), i, path.Ext(destination))\n\t\t\titem, err := processFileOutputs(output, schema.Items.Value, itemDestination)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"Failed to write output %d: %w\", i, err)\n\t\t\t}\n\n\t\t\tclone = append(clone, item)\n\t\t}\n\n\t\treturn clone, nil\n\t}\n\n\treturn output, nil\n}\n\nfunc parseInputFlags(inputs []string, schema *openapi3.T, isTrain ...bool) (predict.Inputs, error) {\n\tkeyVals := map[string][]string{}\n\tfor _, input := range inputs {\n\t\tvar name, value string\n\n\t\t// Default input name is \"input\"\n\t\tif !strings.Contains(input, \"=\") {\n\t\t\treturn nil, fmt.Errorf(\"Failed to parse input '%s', expected format is 'name=value'\", input)\n\t\t}\n\n\t\tsplit := strings.SplitN(input, \"=\", 2)\n\t\tname = split[0]\n\t\tvalue = split[1]\n\n\t\tif strings.HasPrefix(value, `\"`) && strings.HasSuffix(value, `\"`) {\n\t\t\tvalue = value[1 : len(value)-1]\n\t\t}\n\n\t\t// Append new values to the slice associated with the key\n\t\tkeyVals[name] = append(keyVals[name], value)\n\t}\n\n\ttrain := len(isTrain) > 0 && isTrain[0]\n\treturn predict.NewInputsForMode(keyVals, schema, train)\n}\n\nfunc addSetupTimeoutFlag(cmd *cobra.Command) {\n\tcmd.Flags().Uint32Var(&setupTimeout, \"setup-timeout\", 5*60, \"The timeout for a container to setup (in seconds).\")\n}\n"
  },
  {
    "path": "pkg/cli/predict_test.go",
    "content": "package cli\n\nimport (\n\t\"testing\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestExtractOutputSchemaFromMalformedSchema(t *testing.T) {\n\t// Test that we don't panic when extracting output schema from malformed OpenAPI schemas\n\ttestCases := []struct {\n\t\tname   string\n\t\tschema *openapi3.T\n\t}{\n\t\t{\n\t\t\tname:   \"nil schema\",\n\t\t\tschema: nil,\n\t\t},\n\t\t{\n\t\t\tname:   \"empty schema\",\n\t\t\tschema: &openapi3.T{},\n\t\t},\n\t\t{\n\t\t\tname: \"schema with nil paths\",\n\t\t\tschema: &openapi3.T{\n\t\t\t\tPaths: nil,\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"schema with empty paths\",\n\t\t\tschema: &openapi3.T{\n\t\t\t\tPaths: &openapi3.Paths{},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"schema with path but no post\",\n\t\t\tschema: &openapi3.T{\n\t\t\t\tPaths: &openapi3.Paths{\n\t\t\t\t\tExtensions: map[string]any{},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"schema with post but no responses\",\n\t\t\tschema: func() *openapi3.T {\n\t\t\t\ts := &openapi3.T{\n\t\t\t\t\tPaths: openapi3.NewPaths(),\n\t\t\t\t}\n\t\t\t\ts.Paths.Set(\"/predictions\", &openapi3.PathItem{\n\t\t\t\t\tPost: &openapi3.Operation{},\n\t\t\t\t})\n\t\t\t\treturn s\n\t\t\t}(),\n\t\t},\n\t\t{\n\t\t\tname: \"schema with response but no content\",\n\t\t\tschema: func() *openapi3.T {\n\t\t\t\ts := &openapi3.T{\n\t\t\t\t\tPaths: openapi3.NewPaths(),\n\t\t\t\t}\n\t\t\t\ts.Paths.Set(\"/predictions\", &openapi3.PathItem{\n\t\t\t\t\tPost: &openapi3.Operation{\n\t\t\t\t\t\tResponses: &openapi3.Responses{},\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t\treturn s\n\t\t\t}(),\n\t\t},\n\t\t{\n\t\t\tname: \"schema with content but no output property\",\n\t\t\tschema: func() *openapi3.T {\n\t\t\t\ts := &openapi3.T{\n\t\t\t\t\tPaths: openapi3.NewPaths(),\n\t\t\t\t}\n\t\t\t\tresponses := openapi3.NewResponses()\n\t\t\t\tresponses.Set(\"200\", &openapi3.ResponseRef{\n\t\t\t\t\tValue: &openapi3.Response{\n\t\t\t\t\t\tContent: openapi3.Content{\n\t\t\t\t\t\t\t\"application/json\": &openapi3.MediaType{\n\t\t\t\t\t\t\t\tSchema: &openapi3.SchemaRef{\n\t\t\t\t\t\t\t\t\tValue: &openapi3.Schema{\n\t\t\t\t\t\t\t\t\t\tProperties: openapi3.Schemas{},\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t\ts.Paths.Set(\"/predictions\", &openapi3.PathItem{\n\t\t\t\t\tPost: &openapi3.Operation{\n\t\t\t\t\t\tResponses: responses,\n\t\t\t\t\t},\n\t\t\t\t})\n\t\t\t\treturn s\n\t\t\t}(),\n\t\t},\n\t}\n\n\tfor _, tc := range testCases {\n\t\tt.Run(tc.name, func(t *testing.T) {\n\t\t\t// This should not panic - it should return an error or nil gracefully\n\t\t\toutputSchema := safeExtractOutputSchema(tc.schema, \"/predictions\")\n\t\t\t// We expect nil for all malformed schemas\n\t\t\trequire.Nil(t, outputSchema, \"expected nil output schema for malformed input\")\n\t\t})\n\t}\n}\n\n// safeExtractOutputSchema extracts the output schema safely without panicking\nfunc safeExtractOutputSchema(schema *openapi3.T, url string) *openapi3.Schema {\n\tif schema == nil || schema.Paths == nil {\n\t\treturn nil\n\t}\n\tpathItem := schema.Paths.Value(url)\n\tif pathItem == nil || pathItem.Post == nil {\n\t\treturn nil\n\t}\n\tif pathItem.Post.Responses == nil {\n\t\treturn nil\n\t}\n\tresp := pathItem.Post.Responses.Value(\"200\")\n\tif resp == nil || resp.Value == nil {\n\t\treturn nil\n\t}\n\tcontent, ok := resp.Value.Content[\"application/json\"]\n\tif !ok || content == nil || content.Schema == nil || content.Schema.Value == nil {\n\t\treturn nil\n\t}\n\toutputProp, ok := content.Schema.Value.Properties[\"output\"]\n\tif !ok || outputProp == nil {\n\t\treturn nil\n\t}\n\treturn outputProp.Value\n}\n\nfunc TestExtractOutputSchemaFromValidSchema(t *testing.T) {\n\t// Test that we correctly extract output schema from a valid OpenAPI schema\n\ts := &openapi3.T{\n\t\tPaths: openapi3.NewPaths(),\n\t}\n\tresponses := openapi3.NewResponses()\n\tresponses.Set(\"200\", &openapi3.ResponseRef{\n\t\tValue: &openapi3.Response{\n\t\t\tContent: openapi3.Content{\n\t\t\t\t\"application/json\": &openapi3.MediaType{\n\t\t\t\t\tSchema: &openapi3.SchemaRef{\n\t\t\t\t\t\tValue: &openapi3.Schema{\n\t\t\t\t\t\t\tProperties: openapi3.Schemas{\n\t\t\t\t\t\t\t\t\"output\": &openapi3.SchemaRef{\n\t\t\t\t\t\t\t\t\tValue: &openapi3.Schema{\n\t\t\t\t\t\t\t\t\t\tType: &openapi3.Types{\"string\"},\n\t\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\ts.Paths.Set(\"/predictions\", &openapi3.PathItem{\n\t\tPost: &openapi3.Operation{\n\t\t\tResponses: responses,\n\t\t},\n\t})\n\n\toutputSchema := safeExtractOutputSchema(s, \"/predictions\")\n\trequire.NotNil(t, outputSchema, \"expected non-nil output schema for valid input\")\n\trequire.Contains(t, outputSchema.Type.Slice(), \"string\", \"expected string type\")\n}\n"
  },
  {
    "path": "pkg/cli/push.go",
    "content": "package cli\n\nimport (\n\t\"fmt\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/go/uuid\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/model\"\n\t\"github.com/replicate/cog/pkg/provider\"\n\t\"github.com/replicate/cog/pkg/provider/setup\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc newPushCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:   \"push [IMAGE]\",\n\t\tShort: \"Build and push model in current directory to a Docker registry\",\n\t\tLong: `Build a Docker image from cog.yaml and push it to a container registry.\n\nCog can push to any OCI-compliant registry. When pushing to Replicate's\nregistry (r8.im), run 'cog login' first to authenticate.`,\n\t\tExample: `  # Push to Replicate\n  cog push r8.im/your-username/my-model\n\n  # Push to any OCI registry\n  cog push registry.example.com/your-username/model-name\n\n  # Push with model weights in a separate layer (Replicate only)\n  cog push r8.im/your-username/my-model --separate-weights`,\n\t\tRunE: push,\n\t\tArgs: cobra.MaximumNArgs(1),\n\t}\n\taddSecretsFlag(cmd)\n\taddNoCacheFlag(cmd)\n\taddSeparateWeightsFlag(cmd)\n\taddSchemaFlag(cmd)\n\taddUseCudaBaseImageFlag(cmd)\n\taddDockerfileFlag(cmd)\n\taddBuildProgressOutputFlag(cmd)\n\taddUseCogBaseImageFlag(cmd)\n\taddStripFlag(cmd)\n\taddPrecompileFlag(cmd)\n\taddConfigFlag(cmd)\n\n\treturn cmd\n}\n\nfunc push(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\t// Initialize the provider registry\n\tsetup.Init()\n\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tsrc, err := model.NewSource(configFilename)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\timageName := src.Config.Image\n\tif len(args) > 0 {\n\t\timageName = args[0]\n\t}\n\n\tif imageName == \"\" {\n\t\treturn fmt.Errorf(\"To push images, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog push registry.example.com/your-username/model-name'\")\n\t}\n\n\t// Look up the provider for the target registry\n\tp := provider.DefaultRegistry().ForImage(imageName)\n\tif p == nil {\n\t\treturn fmt.Errorf(\"no provider found for image '%s'\", imageName)\n\t}\n\n\tpushOpts := provider.PushOptions{\n\t\tImage:      imageName,\n\t\tConfig:     src.Config,\n\t\tProjectDir: src.ProjectDir,\n\t}\n\n\t// Build the image\n\tbuildID, _ := uuid.NewV7()\n\tannotations := map[string]string{}\n\tif buildID.String() != \"\" {\n\t\tannotations[\"run.cog.push_id\"] = buildID.String()\n\t}\n\n\tregClient := registry.NewRegistryClient()\n\tresolver := model.NewResolver(dockerClient, regClient)\n\n\t// Build the model\n\tconsole.Infof(\"Building Docker image from environment in cog.yaml as %s...\", console.Bold(imageName))\n\tconsole.Info(\"\")\n\tbuildOpts := buildOptionsFromFlags(cmd, imageName, annotations)\n\tm, err := resolver.Build(ctx, src, buildOpts)\n\tif err != nil {\n\t\t// Call PostPush to handle error logging/analytics\n\t\t_ = p.PostPush(ctx, pushOpts, err)\n\t\treturn err\n\t}\n\n\t// Log weights info\n\tweights := m.WeightArtifacts()\n\tif len(weights) > 0 {\n\t\tconsole.Infof(\"\\n%d weight artifact(s)\", len(weights))\n\t}\n\n\t// Push the model (image + optional weights)\n\tconsole.Infof(\"\\nPushing image %s...\", console.Bold(m.ImageRef()))\n\n\t// Set up progress display using Docker's jsonmessage rendering. This uses the\n\t// same cursor movement and progress display as `docker push`, which handles\n\t// terminal resizing correctly (each line is erased and rewritten individually,\n\t// rather than relying on a bulk cursor-up count that can desync on resize).\n\tpw := docker.NewProgressWriter()\n\tdefer pw.Close()\n\n\tpushErr := resolver.Push(ctx, m, model.PushOptions{\n\t\tImageProgressFn: func(prog model.PushProgress) {\n\t\t\t// Truncate digest for display: \"sha256:abc123...\" → \"abc123...\"\n\t\t\tdisplayDigest := prog.LayerDigest\n\t\t\tif len(displayDigest) > 7+12 { // \"sha256:\" + 12 hex chars\n\t\t\t\tdisplayDigest = displayDigest[7:19] + \"...\"\n\t\t\t}\n\n\t\t\tpw.Write(displayDigest, \"Pushing\", prog.Complete, prog.Total)\n\t\t},\n\t\tOnFallback: func() {\n\t\t\t// Close progress writer to finalize OCI progress bars before Docker\n\t\t\t// push starts its own output. Without this, stale OCI progress lines\n\t\t\t// remain on screen above Docker's progress output.\n\t\t\tpw.Close()\n\t\t},\n\t})\n\n\tpw.Close()\n\n\t// PostPush: the provider handles formatting errors and showing success messages\n\tif err := p.PostPush(ctx, pushOpts, pushErr); err != nil {\n\t\treturn err\n\t}\n\n\t// If there was a push error but PostPush didn't return one,\n\t// return a generic error\n\tif pushErr != nil {\n\t\treturn fmt.Errorf(\"failed to push image: %w\", pushErr)\n\t}\n\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/cli/root.go",
    "content": "package cli\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/update\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc NewRootCommand() (*cobra.Command, error) {\n\trootCmd := cobra.Command{\n\t\tUse:   \"cog\",\n\t\tShort: \"Cog: Containers for machine learning\",\n\t\tLong: `Containers for machine learning.\n\nTo get started, take a look at the documentation:\nhttps://github.com/replicate/cog`,\n\t\tExample: `   To run a command inside a Docker environment defined with Cog:\n      $ cog run echo hello world`,\n\t\tVersion: fmt.Sprintf(\"%s (built %s)\", global.Version, global.BuildTime),\n\t\t// This stops errors being printed because we print them in cmd/cog/cog.go\n\t\tPersistentPreRun: func(cmd *cobra.Command, args []string) {\n\t\t\tif global.Debug {\n\t\t\t\tconsole.SetLevel(console.DebugLevel)\n\t\t\t}\n\t\t\tif global.NoColor || !console.ShouldUseColor() {\n\t\t\t\tconsole.SetColor(false)\n\t\t\t}\n\t\t\tif global.NoColor {\n\t\t\t\tos.Setenv(\"NO_COLOR\", \"1\") //nolint:errcheck,gosec // best-effort\n\t\t\t}\n\t\t\tcmd.SilenceUsage = true\n\t\t\tif err := update.DisplayAndCheckForRelease(cmd.Context()); err != nil {\n\t\t\t\tconsole.Debugf(\"%s\", err)\n\t\t\t}\n\t\t},\n\t\tSilenceErrors: true,\n\t}\n\tsetPersistentFlags(&rootCmd)\n\n\trootCmd.AddCommand(\n\t\tnewBuildCommand(),\n\t\tnewDebugCommand(),\n\t\tnewInitCommand(),\n\t\tnewInspectCommand(),\n\t\tnewLoginCommand(),\n\t\tnewPredictCommand(),\n\t\tnewPushCommand(),\n\t\tnewRunCommand(),\n\t\tnewServeCommand(),\n\t\tnewTrainCommand(),\n\t\tnewWeightsCommand(),\n\t)\n\n\treturn &rootCmd, nil\n}\n\nfunc setPersistentFlags(cmd *cobra.Command) {\n\tcmd.PersistentFlags().BoolVar(&global.Debug, \"debug\", false, \"Show debugging output\")\n\tcmd.PersistentFlags().BoolVar(&global.NoColor, \"no-color\", false, \"Disable colored output\")\n\tcmd.PersistentFlags().BoolVar(&global.ProfilingEnabled, \"profile\", false, \"Enable profiling\")\n\tcmd.PersistentFlags().Bool(\"version\", false, \"Show version of Cog\")\n\tcmd.PersistentFlags().StringVar(&global.ReplicateRegistryHost, \"registry\", global.ReplicateRegistryHost, \"Registry host\")\n\t_ = cmd.PersistentFlags().MarkHidden(\"profile\")\n\t_ = cmd.PersistentFlags().MarkHidden(\"registry\")\n}\n"
  },
  {
    "path": "pkg/cli/run.go",
    "content": "package cli\n\nimport (\n\t\"os\"\n\t\"strconv\"\n\t\"strings\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/model\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nvar (\n\trunPorts []string\n\tgpusFlag string\n)\n\nfunc addGpusFlag(cmd *cobra.Command) {\n\tcmd.Flags().StringVar(&gpusFlag, \"gpus\", \"\", \"GPU devices to add to the container, in the same format as `docker run --gpus`.\")\n}\n\nfunc newRunCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:   \"run  [arg...]\",\n\t\tShort: \"Run a command inside a Docker environment\",\n\t\tLong: `Run a command inside a Docker environment defined by cog.yaml.\n\nCog builds a temporary image from your cog.yaml configuration and runs the\ngiven command inside it. This is useful for debugging, running scripts, or\nexploring the environment your model will run in.`,\n\t\tExample: `  # Open a Python interpreter inside the model environment\n  cog run python\n\n  # Run a script\n  cog run python train.py\n\n  # Run with environment variables\n  cog run -e HUGGING_FACE_HUB_TOKEN=abc123 python download.py\n\n  # Expose a port (e.g. for Jupyter)\n  cog run -p 8888 jupyter notebook`,\n\t\tRunE:    run,\n\t\tPreRunE: checkMutuallyExclusiveFlags,\n\t\tArgs:    cobra.MinimumNArgs(1),\n\t}\n\taddBuildProgressOutputFlag(cmd)\n\taddDockerfileFlag(cmd)\n\taddUseCudaBaseImageFlag(cmd)\n\taddUseCogBaseImageFlag(cmd)\n\taddGpusFlag(cmd)\n\taddConfigFlag(cmd)\n\n\tflags := cmd.Flags()\n\t// Flags after first argument are considered args and passed to command\n\n\t// This is called `publish` for consistency with `docker run`\n\tcmd.Flags().StringArrayVarP(&runPorts, \"publish\", \"p\", []string{}, \"Publish a container's port to the host, e.g. -p 8000\")\n\tcmd.Flags().StringArrayVarP(&envFlags, \"env\", \"e\", []string{}, \"Environment variables, in the form name=value\")\n\n\tflags.SetInterspersed(false)\n\n\treturn cmd\n}\n\nfunc run(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tsrc, err := model.NewSource(configFilename)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tresolver := model.NewResolver(dockerClient, registry.NewRegistryClient())\n\n\tconsole.Info(\"Building Docker image from environment in cog.yaml...\")\n\tconsole.Info(\"\")\n\topts := serveBuildOptions(cmd)\n\topts.SkipSchemaValidation = true\n\tm, err := resolver.Build(ctx, src, opts)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tgpus := \"\"\n\tif gpusFlag != \"\" {\n\t\tgpus = gpusFlag\n\t} else if m.HasGPU() {\n\t\tgpus = \"all\"\n\t}\n\n\t// Automatically propagate RUST_LOG for Rust coglet debugging\n\tenv := envFlags\n\tif rustLog := os.Getenv(\"RUST_LOG\"); rustLog != \"\" {\n\t\tenv = append(env, \"RUST_LOG=\"+rustLog)\n\t}\n\n\trunOptions := command.RunOptions{\n\t\tArgs:    args,\n\t\tEnv:     env,\n\t\tGPUs:    gpus,\n\t\tImage:   m.ImageRef(),\n\t\tVolumes: []command.Volume{{Source: src.ProjectDir, Destination: \"/src\"}},\n\t\tWorkdir: \"/src\",\n\t}\n\n\tfor _, portString := range runPorts {\n\t\tport, err := strconv.Atoi(portString)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\trunOptions.Ports = append(runOptions.Ports, command.Port{HostPort: port, ContainerPort: port})\n\t}\n\n\tconsole.Info(\"\")\n\tconsole.Infof(\"Running %s in Docker with the current directory mounted as a volume...\", console.Bold(strings.Join(args, \" \")))\n\tconsole.Info(\"\")\n\n\terr = docker.Run(ctx, dockerClient, runOptions)\n\t// Only retry if we're using a GPU but the user didn't explicitly select a GPU with --gpus\n\t// If the user specified the wrong GPU, they are explicitly selecting a GPU and they'll want to hear about it\n\tif runOptions.GPUs == \"all\" && err == docker.ErrMissingDeviceDriver {\n\t\tconsole.Info(\"Missing device driver, re-trying without GPU\")\n\n\t\trunOptions.GPUs = \"\"\n\t\terr = docker.Run(ctx, dockerClient, runOptions)\n\t}\n\n\treturn err\n}\n"
  },
  {
    "path": "pkg/cli/serve.go",
    "content": "package cli\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\t\"strings\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/model\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nvar (\n\tport      = 8393\n\tuploadURL = \"\"\n)\n\nfunc newServeCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:   \"serve\",\n\t\tShort: \"Run a prediction HTTP server\",\n\t\tLong: `Run a prediction HTTP server.\n\nBuilds the model and starts an HTTP server that exposes the model's inputs\nand outputs as a REST API. Compatible with the Cog HTTP protocol.`,\n\t\tExample: `  # Start the server on the default port (8393)\n  cog serve\n\n  # Start on a custom port\n  cog serve -p 5000\n\n  # Test the server\n  curl http://localhost:8393/predictions \\\n    -X POST \\\n    -H 'Content-Type: application/json' \\\n    -d '{\"input\": {\"prompt\": \"a cat\"}}'`,\n\t\tRunE:       cmdServe,\n\t\tArgs:       cobra.MaximumNArgs(0),\n\t\tSuggestFor: []string{\"http\"},\n\t}\n\n\taddBuildProgressOutputFlag(cmd)\n\taddUseCudaBaseImageFlag(cmd)\n\taddUseCogBaseImageFlag(cmd)\n\taddGpusFlag(cmd)\n\taddConfigFlag(cmd)\n\n\tcmd.Flags().IntVarP(&port, \"port\", \"p\", port, \"Port on which to listen\")\n\tcmd.Flags().StringVar(&uploadURL, \"upload-url\", \"\", \"Upload URL for file outputs (e.g. https://example.com/upload/)\")\n\n\treturn cmd\n}\n\n// serveBuildOptions creates BuildOptions for cog serve.\n// Same build path as cog build, but with ExcludeSource so COPY . /src is\n// skipped — source is volume-mounted at runtime instead. All other layers\n// (wheels, apt, etc.) share Docker layer cache with cog build.\nfunc serveBuildOptions(cmd *cobra.Command) model.BuildOptions {\n\treturn model.BuildOptions{\n\t\tUseCudaBaseImage: buildUseCudaBaseImage,\n\t\tUseCogBaseImage:  DetermineUseCogBaseImage(cmd),\n\t\tProgressOutput:   buildProgressOutput,\n\t\tExcludeSource:    true,\n\t\tSkipLabels:       true,\n\t}\n}\n\nfunc cmdServe(cmd *cobra.Command, arg []string) error {\n\tctx := cmd.Context()\n\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tsrc, err := model.NewSource(configFilename)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tconsole.Info(\"Building Docker image from environment in cog.yaml...\")\n\tconsole.Info(\"\")\n\tresolver := model.NewResolver(dockerClient, registry.NewRegistryClient())\n\tm, err := resolver.Build(ctx, src, serveBuildOptions(cmd))\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tgpus := \"\"\n\tif gpusFlag != \"\" {\n\t\tgpus = gpusFlag\n\t} else if m.HasGPU() {\n\t\tgpus = \"all\"\n\t}\n\n\targs := []string{\n\t\t\"python\",\n\t\t\"--check-hash-based-pycs\", \"never\",\n\t\t\"-m\", \"cog.server.http\",\n\t\t\"--await-explicit-shutdown\", \"true\",\n\t}\n\n\tif uploadURL != \"\" {\n\t\targs = append(args, \"--upload-url\", uploadURL)\n\t}\n\n\t// Automatically propagate RUST_LOG for Rust coglet debugging\n\tenv := envFlags\n\tif rustLog := os.Getenv(\"RUST_LOG\"); rustLog != \"\" {\n\t\tenv = append(env, \"RUST_LOG=\"+rustLog)\n\t}\n\n\trunOptions := command.RunOptions{\n\t\tArgs:    args,\n\t\tEnv:     env,\n\t\tGPUs:    gpus,\n\t\tImage:   m.ImageRef(),\n\t\tVolumes: []command.Volume{{Source: src.ProjectDir, Destination: \"/src\"}},\n\t\tWorkdir: \"/src\",\n\t}\n\n\t// On Linux, host.docker.internal is not available by default — add it.\n\t// This allows the container to reach services running on the host,\n\t// e.g. when --upload-url points to a local upload server.\n\tif uploadURL != \"\" {\n\t\trunOptions.ExtraHosts = []string{\"host.docker.internal:host-gateway\"}\n\t}\n\n\trunOptions.Ports = append(runOptions.Ports, command.Port{HostPort: port, ContainerPort: 5000})\n\n\tconsole.Info(\"\")\n\tconsole.Infof(\"Running %[1]s in Docker with the current directory mounted as a volume...\", console.Bold(strings.Join(args, \" \")))\n\tconsole.Info(\"\")\n\tconsole.Infof(\"Serving at %s\", console.Bold(fmt.Sprintf(\"http://127.0.0.1:%v\", port)))\n\tconsole.Info(\"\")\n\n\terr = docker.Run(ctx, dockerClient, runOptions)\n\t// Only retry if we're using a GPU but the user didn't explicitly select a GPU with --gpus\n\t// If the user specified the wrong GPU, they are explicitly selecting a GPU and they'll want to hear about it\n\tif runOptions.GPUs == \"all\" && err == docker.ErrMissingDeviceDriver {\n\t\tconsole.Info(\"Missing device driver, re-trying without GPU\")\n\n\t\trunOptions.GPUs = \"\"\n\t\terr = docker.Run(ctx, dockerClient, runOptions)\n\t}\n\n\treturn err\n}\n"
  },
  {
    "path": "pkg/cli/train.go",
    "content": "package cli\n\nimport (\n\t\"context\"\n\t\"os\"\n\t\"os/signal\"\n\t\"syscall\"\n\t\"time\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/model\"\n\t\"github.com/replicate/cog/pkg/predict\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nvar (\n\ttrainEnvFlags   []string\n\ttrainInputFlags []string\n\ttrainOutPath    string\n)\n\nfunc newTrainCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:   \"train [image]\",\n\t\tShort: \"Run a training\",\n\t\tLong: `Run a training.\n\nIf 'image' is passed, it will run the training on that Docker image.\nIt must be an image that has been built by Cog.\n\nOtherwise, it will build the model in the current directory and train it.`,\n\t\tRunE:       cmdTrain,\n\t\tArgs:       cobra.MaximumNArgs(1),\n\t\tHidden:     true,\n\t\tDeprecated: \"the train command will be removed in a future version of Cog\",\n\t}\n\n\taddBuildProgressOutputFlag(cmd)\n\taddDockerfileFlag(cmd)\n\taddUseCudaBaseImageFlag(cmd)\n\taddGpusFlag(cmd)\n\taddUseCogBaseImageFlag(cmd)\n\taddConfigFlag(cmd)\n\n\tcmd.Flags().StringArrayVarP(&trainInputFlags, \"input\", \"i\", []string{}, \"Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg\")\n\tcmd.Flags().StringArrayVarP(&trainEnvFlags, \"env\", \"e\", []string{}, \"Environment variables, in the form name=value\")\n\tcmd.Flags().StringVarP(&trainOutPath, \"output\", \"o\", \"weights\", \"Output path\")\n\n\treturn cmd\n}\n\nfunc cmdTrain(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\tdockerClient, err := docker.NewClient(ctx)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\timageName := \"\"\n\tvolumes := []command.Volume{}\n\tgpus := gpusFlag\n\n\tresolver := model.NewResolver(dockerClient, registry.NewRegistryClient())\n\n\tif len(args) == 0 {\n\t\t// Build image\n\t\tsrc, err := model.NewSource(configFilename)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tconsole.Info(\"Building Docker image from environment in cog.yaml...\")\n\t\tconsole.Info(\"\")\n\t\tm, err := resolver.Build(ctx, src, serveBuildOptions(cmd))\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\timageName = m.ImageRef()\n\n\t\t// ExcludeSource build doesn't have /src in it, so mount as volume\n\t\tvolumes = append(volumes, command.Volume{\n\t\t\tSource:      src.ProjectDir,\n\t\t\tDestination: \"/src\",\n\t\t})\n\n\t\tif gpus == \"\" && m.HasGPU() {\n\t\t\tgpus = \"all\"\n\t\t}\n\t} else {\n\t\t// Use existing image\n\t\timageName = args[0]\n\n\t\t// Pull the image (if needed) and validate it's a Cog model\n\t\tref, err := model.ParseRef(imageName)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tm, err := resolver.Pull(ctx, ref)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tif gpus == \"\" && m.HasGPU() {\n\t\t\tgpus = \"all\"\n\t\t}\n\t}\n\n\tconsole.Info(\"\")\n\tconsole.Info(\"Starting Docker image and running setup()...\")\n\n\tpredictor, err := predict.NewPredictor(ctx, command.RunOptions{\n\t\tGPUs:    gpus,\n\t\tImage:   imageName,\n\t\tVolumes: volumes,\n\t\tEnv:     trainEnvFlags,\n\t\tArgs:    []string{\"python\", \"-m\", \"cog.server.http\", \"--x-mode\", \"train\"},\n\t}, true, dockerClient)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tgo func() {\n\t\tcaptureSignal := make(chan os.Signal, 1)\n\t\tsignal.Notify(captureSignal, syscall.SIGINT)\n\n\t\t<-captureSignal\n\n\t\tconsole.Info(\"Stopping container...\")\n\t\tif err := predictor.Stop(ctx); err != nil {\n\t\t\tconsole.Warnf(\"Failed to stop container: %s\", err)\n\t\t}\n\t}()\n\n\tif err := predictor.Start(ctx, os.Stderr, time.Duration(setupTimeout)*time.Second); err != nil {\n\t\treturn err\n\t}\n\n\t// FIXME: will not run on signal\n\tdefer func() {\n\t\tconsole.Debugf(\"Stopping container...\")\n\t\t// use background context to ensure stop signal is still sent after root context is canceled\n\t\tif err := predictor.Stop(context.Background()); err != nil {\n\t\t\tconsole.Warnf(\"Failed to stop container: %s\", err)\n\t\t}\n\t}()\n\n\treturn predictIndividualInputs(*predictor, trainInputFlags, trainOutPath, true)\n}\n"
  },
  {
    "path": "pkg/cli/train_test.go",
    "content": "package cli\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestTrainCommandIsDeprecated(t *testing.T) {\n\tcmd := newTrainCommand()\n\trequire.NotEmpty(t, cmd.Deprecated, \"train command should have a deprecation message\")\n\trequire.Contains(t, cmd.Deprecated, \"will be removed in a future version\")\n}\n"
  },
  {
    "path": "pkg/cli/weights.go",
    "content": "package cli\n\nimport (\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"time\"\n\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\t\"github.com/spf13/cobra\"\n\t\"golang.org/x/sync/errgroup\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/model\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc newWeightsCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:    \"weights\",\n\t\tShort:  \"Manage model weights\",\n\t\tLong:   \"Commands for managing model weight files.\",\n\t\tHidden: true,\n\t}\n\n\tcmd.AddCommand(newWeightsBuildCommand())\n\tcmd.AddCommand(newWeightsInspectCommand())\n\tcmd.AddCommand(newWeightsPushCommand())\n\treturn cmd\n}\n\nfunc newWeightsBuildCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:   \"build\",\n\t\tShort: \"Generate weights.lock from weight sources in cog.yaml\",\n\t\tLong: `Reads the weights section from cog.yaml, processes each weight source,\nand generates a weights.lock file containing metadata (digests, sizes) for each file.`,\n\t\tArgs: cobra.NoArgs,\n\t\tRunE: weightsBuildCommand,\n\t}\n\n\taddConfigFlag(cmd)\n\treturn cmd\n}\n\nfunc weightsBuildCommand(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\tsrc, err := model.NewSource(configFilename)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to read config: %w\", err)\n\t}\n\n\tif len(src.Config.Weights) == 0 {\n\t\treturn fmt.Errorf(\"no weights defined in %s\", configFilename)\n\t}\n\n\t// Extract weight specs from the source\n\tvar weightSpecs []*model.WeightSpec\n\tfor _, spec := range src.ArtifactSpecs() {\n\t\tif ws, ok := spec.(*model.WeightSpec); ok {\n\t\t\tweightSpecs = append(weightSpecs, ws)\n\t\t}\n\t}\n\n\tconsole.Infof(\"Processing %d weight source(s)...\", len(weightSpecs))\n\n\tlockPath := filepath.Join(src.ProjectDir, model.WeightsLockFilename)\n\tbuilder := model.NewWeightBuilder(src, global.Version, lockPath)\n\n\t// Build each weight artifact (hashes file, updates lockfile)\n\tvar totalSize int64\n\tfor _, ws := range weightSpecs {\n\t\tartifact, buildErr := builder.Build(ctx, ws)\n\t\tif buildErr != nil {\n\t\t\treturn fmt.Errorf(\"failed to build weight %q: %w\", ws.Name(), buildErr)\n\t\t}\n\n\t\twa, ok := artifact.(*model.WeightArtifact)\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"unexpected artifact type %T for weight %q\", artifact, ws.Name())\n\t\t}\n\t\tsize := wa.Descriptor().Size\n\t\ttotalSize += size\n\t\tconsole.Infof(\"  %s -> %s (%s)\", wa.Name(), wa.Target, formatSize(size))\n\t}\n\n\tconsole.Infof(\"\\nGenerated %s with %d file(s) (%s total)\",\n\t\tmodel.WeightsLockFilename, len(weightSpecs), formatSize(totalSize))\n\n\treturn nil\n}\n\nfunc formatSize(bytes int64) string {\n\tconst (\n\t\tkb = 1024\n\t\tmb = kb * 1024\n\t\tgb = mb * 1024\n\t)\n\n\tswitch {\n\tcase bytes >= gb:\n\t\treturn fmt.Sprintf(\"%.1fGB\", float64(bytes)/float64(gb))\n\tcase bytes >= mb:\n\t\treturn fmt.Sprintf(\"%.1fMB\", float64(bytes)/float64(mb))\n\tcase bytes >= kb:\n\t\treturn fmt.Sprintf(\"%.1fKB\", float64(bytes)/float64(kb))\n\tdefault:\n\t\treturn fmt.Sprintf(\"%dB\", bytes)\n\t}\n}\n\nfunc newWeightsPushCommand() *cobra.Command {\n\tcmd := &cobra.Command{\n\t\tUse:   \"push [IMAGE]\",\n\t\tShort: \"Push weights to a registry\",\n\t\tLong: `Reads weights.lock and pushes weight files as an OCI artifact to a registry.\n\nThe registry is determined from the image name, which can be:\n- Specified as an argument: cog weights push registry.example.com/user/model\n- Set in cog.yaml as the 'image' field`,\n\t\tArgs: cobra.MaximumNArgs(1),\n\t\tRunE: weightsPushCommand,\n\t}\n\n\taddConfigFlag(cmd)\n\treturn cmd\n}\n\nfunc weightsPushCommand(cmd *cobra.Command, args []string) error {\n\tctx := cmd.Context()\n\n\tsrc, err := model.NewSource(configFilename)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to read config: %w\", err)\n\t}\n\n\tcfg := src.Config\n\n\t// Determine image name\n\timageName := cfg.Image\n\tif len(args) > 0 {\n\t\timageName = args[0]\n\t}\n\tif imageName == \"\" {\n\t\treturn fmt.Errorf(\"To push weights, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog weights push registry.example.com/your-username/model-name'\")\n\t}\n\n\t// Parse as repository only — reject tags/digests since weight tags are auto-generated.\n\tparsedRepo, err := name.NewRepository(imageName, name.Insecure)\n\tif err != nil {\n\t\t// NewRepository fails for inputs with :tag or @digest — check if it's a valid ref\n\t\tif ref, refErr := name.ParseReference(imageName, name.Insecure); refErr == nil {\n\t\t\treturn fmt.Errorf(\"image reference %q includes a tag or digest — provide only the repository (e.g., %q)\", imageName, ref.Context().Name())\n\t\t}\n\t\treturn fmt.Errorf(\"invalid repository %q: %w\", imageName, err)\n\t}\n\trepo := parsedRepo.Name()\n\n\tif len(cfg.Weights) == 0 {\n\t\treturn fmt.Errorf(\"no weights defined in %s\", configFilename)\n\t}\n\n\t// Build weight artifacts (reads lockfile as cache, hashes files)\n\tlockPath := filepath.Join(src.ProjectDir, model.WeightsLockFilename)\n\tbuilder := model.NewWeightBuilder(src, global.Version, lockPath)\n\n\tvar artifacts []*model.WeightArtifact\n\tfor _, spec := range src.ArtifactSpecs() {\n\t\tws, ok := spec.(*model.WeightSpec)\n\t\tif !ok {\n\t\t\tcontinue\n\t\t}\n\t\tartifact, buildErr := builder.Build(ctx, ws)\n\t\tif buildErr != nil {\n\t\t\treturn fmt.Errorf(\"failed to build weight %q: %w\", ws.Name(), buildErr)\n\t\t}\n\t\twa, ok := artifact.(*model.WeightArtifact)\n\t\tif !ok {\n\t\t\treturn fmt.Errorf(\"unexpected artifact type %T for weight %q\", artifact, ws.Name())\n\t\t}\n\t\tartifacts = append(artifacts, wa)\n\t}\n\n\tif len(artifacts) == 0 {\n\t\treturn fmt.Errorf(\"no weight artifacts to push\")\n\t}\n\n\tconsole.Infof(\"Pushing %d weight file(s) to %s...\", len(artifacts), repo)\n\n\tregClient := registry.NewRegistryClient()\n\tpusher := model.NewWeightPusher(regClient)\n\n\t// Set up progress display using Docker's jsonmessage rendering.\n\tpw := docker.NewProgressWriter()\n\tdefer pw.Close()\n\n\t// Push each weight artifact concurrently using errgroup for\n\t// bounded concurrency and first-error cancellation.\n\ttype pushResult struct {\n\t\tref  string\n\t\tsize int64\n\t}\n\n\tordered := make([]pushResult, len(artifacts))\n\n\tg, ctx := errgroup.WithContext(ctx)\n\tg.SetLimit(model.GetPushConcurrency())\n\n\tfor i, wa := range artifacts {\n\t\tartName := wa.Name()\n\t\tartSize := wa.Descriptor().Size\n\n\t\tg.Go(func() error {\n\t\t\tresult, pushErr := pusher.Push(ctx, repo, wa, model.WeightPushOptions{\n\t\t\t\tProgressFn: func(prog model.PushProgress) {\n\t\t\t\t\tpw.Write(artName, \"Pushing\", prog.Complete, prog.Total)\n\t\t\t\t},\n\t\t\t\tRetryFn: func(event model.WeightRetryEvent) bool {\n\t\t\t\t\tstatus := fmt.Sprintf(\"Retrying (%d/%d) in %s\",\n\t\t\t\t\t\tevent.Attempt, event.MaxAttempts,\n\t\t\t\t\t\tevent.NextRetryIn.Round(time.Second))\n\t\t\t\t\tpw.WriteStatus(event.Name, status)\n\t\t\t\t\t// In non-TTY mode, also log the error detail since the\n\t\t\t\t\t// progress writer output won't be visible.\n\t\t\t\t\tif !console.IsTerminal() {\n\t\t\t\t\t\tconsole.Warnf(\"  %s: retrying (%d/%d) in %s: %v\",\n\t\t\t\t\t\t\tevent.Name, event.Attempt, event.MaxAttempts,\n\t\t\t\t\t\t\tevent.NextRetryIn.Round(time.Second), event.Err)\n\t\t\t\t\t}\n\t\t\t\t\treturn true\n\t\t\t\t},\n\t\t\t})\n\n\t\t\tif pushErr != nil {\n\t\t\t\tpw.WriteStatus(artName, \"FAILED\")\n\t\t\t\treturn fmt.Errorf(\"push weight %q: %w\", artName, pushErr)\n\t\t\t}\n\n\t\t\tpw.WriteStatus(artName, \"Pushed\")\n\t\t\tordered[i] = pushResult{ref: result.Ref, size: artSize}\n\t\t\treturn nil\n\t\t})\n\t}\n\n\tif err := g.Wait(); err != nil {\n\t\tpw.Close()\n\t\treturn err\n\t}\n\n\t// Close progress display\n\tpw.Close()\n\n\t// Print final summary\n\tvar totalSize int64\n\tfor i, wa := range artifacts {\n\t\tconsole.Infof(\"  %s: %s\", wa.Name(), ordered[i].ref)\n\t\ttotalSize += ordered[i].size\n\t}\n\n\tconsole.Infof(\"\\nPushed %d weight artifact(s) to %s\", len(artifacts), repo)\n\tconsole.Infof(\"Total: %s\", formatSize(totalSize))\n\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/cli/weights_inspect.go",
    "content": "package cli\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/model\"\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// localWeight tracks the local state of a weight from cog.yaml + weights.lock.\ntype localWeight struct {\n\ttarget   string\n\tsource   string\n\tlockFile *model.WeightFile\n}\n\n// WeightsInspectOutput is the structured output for cog weights inspect --json.\ntype WeightsInspectOutput struct {\n\tReference string               `json:\"reference\"`\n\tWeights   []WeightInspectEntry `json:\"weights\"`\n}\n\n// WeightInspectEntry represents one weight's comparison between local and remote state.\ntype WeightInspectEntry struct {\n\tName   string             `json:\"name\"`\n\tStatus string             `json:\"status\"` // synced, local-only, remote-only, digest-mismatch, missing-lockfile\n\tLocal  *WeightLocalState  `json:\"local,omitempty\"`\n\tRemote *WeightRemoteState `json:\"remote,omitempty\"`\n}\n\n// WeightLocalState represents the local state of a weight from cog.yaml + weights.lock.\ntype WeightLocalState struct {\n\tDigest     string `json:\"digest\"`\n\tSize       int64  `json:\"size\"`\n\tTarget     string `json:\"target\"`\n\tFileExists bool   `json:\"fileExists\"`\n}\n\n// WeightRemoteLayer represents a single layer in a remote weight manifest.\ntype WeightRemoteLayer struct {\n\tDigest    string `json:\"digest\"`\n\tSize      int64  `json:\"size\"`\n\tMediaType string `json:\"mediaType\"`\n}\n\n// WeightRemoteState represents the remote state of a weight from the registry.\ntype WeightRemoteState struct {\n\tRef              string              `json:\"ref\"`\n\tTag              string              `json:\"tag\"`\n\tDigest           string              `json:\"digest\"`\n\tSize             int64               `json:\"size\"`\n\tMediaType        string              `json:\"mediaType\"`\n\tLayers           []WeightRemoteLayer `json:\"layers,omitempty\"`\n\tMatchedByContent bool                `json:\"matchedByContent,omitempty\"`\n}\n\nfunc newWeightsInspectCommand() *cobra.Command {\n\tvar jsonOutput bool\n\n\tcmd := &cobra.Command{\n\t\tUse:   \"inspect \",\n\t\tShort: \"Compare local weights against remote registry state\",\n\t\tArgs:  cobra.ExactArgs(1),\n\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\treturn weightsInspectCommand(cmd, args, jsonOutput)\n\t\t},\n\t}\n\n\tcmd.Flags().BoolVar(&jsonOutput, \"json\", false, \"Output as JSON\")\n\taddConfigFlag(cmd)\n\n\treturn cmd\n}\n\nfunc weightsInspectCommand(cmd *cobra.Command, args []string, jsonOutput bool) error {\n\tctx := cmd.Context()\n\n\t// 1. Load local state\n\tsrc, err := model.NewSource(configFilename)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to read config: %w\", err)\n\t}\n\n\tlockPath := filepath.Join(src.ProjectDir, model.WeightsLockFilename)\n\tlock, lockErr := model.LoadWeightsLock(lockPath)\n\t// lockErr is OK — lockfile may not exist yet\n\n\t// Build local weight map: name -> (lockfile entry, source file path)\n\tlocalWeights := make(map[string]*localWeight)\n\tfor _, w := range src.Config.Weights {\n\t\tlw := &localWeight{\n\t\t\ttarget: w.Target,\n\t\t\tsource: w.Source,\n\t\t}\n\t\tlocalWeights[w.Name] = lw\n\t}\n\n\t// Fill in lockfile data\n\tif lockErr == nil && lock != nil {\n\t\tfor i := range lock.Files {\n\t\t\tf := &lock.Files[i]\n\t\t\tif lw, ok := localWeights[f.Name]; ok {\n\t\t\t\tlw.lockFile = f\n\t\t\t}\n\t\t}\n\t}\n\n\t// 2. Resolve remote state — accept repo only (tags are auto-generated for weights).\n\tparsedRepo, err := name.NewRepository(args[0], name.Insecure)\n\tif err != nil {\n\t\tif ref, refErr := name.ParseReference(args[0], name.Insecure); refErr == nil {\n\t\t\treturn fmt.Errorf(\"image reference %q includes a tag or digest — provide only the repository (e.g., %q)\", args[0], ref.Context().Name())\n\t\t}\n\t\treturn fmt.Errorf(\"invalid repository %q: %w\", args[0], err)\n\t}\n\trepo := parsedRepo.Name()\n\n\tregClient := registry.NewRegistryClient()\n\tremoteWeights := resolveWeightsByTag(ctx, repo, localWeights, regClient)\n\n\t// 3. Build comparison\n\tout := &WeightsInspectOutput{\n\t\tReference: repo,\n\t}\n\n\t// Track which remote weights we've matched\n\tmatchedRemote := make(map[string]bool)\n\n\t// Process local weights\n\tfor _, w := range src.Config.Weights {\n\t\tentry := WeightInspectEntry{Name: w.Name}\n\t\tlw := localWeights[w.Name]\n\n\t\tif lw.lockFile == nil {\n\t\t\t// No lockfile entry — needs `cog weights build`\n\t\t\tentry.Status = \"missing-lockfile\"\n\t\t\tentry.Local = &WeightLocalState{\n\t\t\t\tTarget:     lw.target,\n\t\t\t\tFileExists: fileExists(filepath.Join(src.ProjectDir, lw.source)),\n\t\t\t}\n\t\t} else {\n\t\t\t// Check if source file exists on disk\n\t\t\texists := fileExists(filepath.Join(src.ProjectDir, lw.source))\n\t\t\tentry.Local = &WeightLocalState{\n\t\t\t\tDigest:     lw.lockFile.Digest,\n\t\t\t\tSize:       lw.lockFile.Size,\n\t\t\t\tTarget:     lw.lockFile.Dest,\n\t\t\t\tFileExists: exists,\n\t\t\t}\n\n\t\t\tif remote, ok := remoteWeights[w.Name]; ok {\n\t\t\t\tmatchedRemote[w.Name] = true\n\t\t\t\tentry.Remote = remote\n\n\t\t\t\tif remote.MatchedByContent || lw.lockFile.Digest == remote.Digest {\n\t\t\t\t\tentry.Status = \"synced\"\n\t\t\t\t} else {\n\t\t\t\t\tentry.Status = \"digest-mismatch\"\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tentry.Status = \"local-only\"\n\t\t\t}\n\t\t}\n\n\t\tout.Weights = append(out.Weights, entry)\n\t}\n\n\t// Add remote-only weights\n\tfor name, remote := range remoteWeights {\n\t\tif matchedRemote[name] {\n\t\t\tcontinue\n\t\t}\n\t\tout.Weights = append(out.Weights, WeightInspectEntry{\n\t\t\tName:   name,\n\t\t\tStatus: \"remote-only\",\n\t\t\tRemote: remote,\n\t\t})\n\t}\n\n\t// 4. Output\n\tif jsonOutput {\n\t\tenc := json.NewEncoder(os.Stdout)\n\t\tenc.SetIndent(\"\", \"  \")\n\t\treturn enc.Encode(out)\n\t}\n\n\tprintWeightsInspectText(out)\n\treturn nil\n}\n\n// resolveWeightsByTag checks for each local weight's tag in the registry.\n// This is the fallback path when no OCI index exists (e.g., after `cog weights push`\n// but before `cog push`).\n//\n// It looks up the combined tag :weights-- which encodes both\n// the weight name and its content digest. A match means the exact content is synced.\nfunc resolveWeightsByTag(ctx context.Context, repo string, localWeights map[string]*localWeight, reg registry.Client) map[string]*WeightRemoteState {\n\tresult := make(map[string]*WeightRemoteState)\n\tfor weightName, lw := range localWeights {\n\t\tif lw.lockFile == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\ttag := model.WeightTag(weightName, lw.lockFile.Digest)\n\t\ttagRef := repo + \":\" + tag\n\n\t\t// Use GetImage to fetch the full manifest (not just HEAD) so we can read layer sizes.\n\t\timg, err := reg.GetImage(ctx, tagRef, nil)\n\t\tif err != nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tmanifest, err := img.Manifest()\n\t\tif err != nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tdigest, err := img.Digest()\n\t\tif err != nil {\n\t\t\tcontinue\n\t\t}\n\n\t\trawManifest, err := img.RawManifest()\n\t\tif err != nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tstate := &WeightRemoteState{\n\t\t\tRef:              tagRef,\n\t\t\tTag:              tag,\n\t\t\tDigest:           digest.String(),\n\t\t\tSize:             int64(len(rawManifest)),\n\t\t\tMediaType:        string(manifest.MediaType),\n\t\t\tMatchedByContent: true,\n\t\t}\n\n\t\tfor _, layer := range manifest.Layers {\n\t\t\tstate.Layers = append(state.Layers, WeightRemoteLayer{\n\t\t\t\tDigest:    layer.Digest.String(),\n\t\t\t\tSize:      layer.Size,\n\t\t\t\tMediaType: string(layer.MediaType),\n\t\t\t})\n\t\t}\n\n\t\tresult[weightName] = state\n\t}\n\tif len(result) == 0 {\n\t\treturn nil\n\t}\n\treturn result\n}\n\nfunc printWeightsInspectText(out *WeightsInspectOutput) {\n\tfmt.Printf(\"Weights for: %s\\n\\n\", out.Reference)\n\n\tfor _, w := range out.Weights {\n\t\tif w.Remote != nil && w.Remote.Tag != \"\" {\n\t\t\tfmt.Printf(\"  %s  :%s\\n\", w.Name, w.Remote.Tag)\n\t\t} else {\n\t\t\tfmt.Printf(\"  %s\\n\", w.Name)\n\t\t}\n\t\tfmt.Printf(\"    Status:  %s\", w.Status)\n\n\t\tswitch w.Status {\n\t\tcase \"local-only\":\n\t\t\tfmt.Print(\" (not pushed)\")\n\t\tcase \"remote-only\":\n\t\t\tfmt.Print(\" (not in cog.yaml)\")\n\t\tcase \"missing-lockfile\":\n\t\t\tfmt.Print(\" (run cog weights build)\")\n\t\t}\n\t\tfmt.Println()\n\n\t\tif w.Local != nil {\n\t\t\tif w.Local.Digest != \"\" {\n\t\t\t\tfmt.Printf(\"    Local:   %s (%s) -> %s\\n\", w.Local.Digest, formatSize(w.Local.Size), w.Local.Target)\n\t\t\t} else {\n\t\t\t\tfmt.Printf(\"    Local:   (no lockfile entry) -> %s\\n\", w.Local.Target)\n\t\t\t}\n\t\t} else {\n\t\t\tfmt.Println(\"    Local:   -\")\n\t\t}\n\n\t\tif w.Remote != nil {\n\t\t\tfor _, layer := range w.Remote.Layers {\n\t\t\t\tfmt.Printf(\"    Layer:   %s (%s)\\n\", layer.Digest, formatSize(layer.Size))\n\t\t\t}\n\t\t} else {\n\t\t\tfmt.Println(\"    Remote:  -\")\n\t\t}\n\n\t\tfmt.Println()\n\t}\n}\n\nfunc fileExists(path string) bool {\n\t_, err := os.Stat(path)\n\treturn err == nil\n}\n"
  },
  {
    "path": "pkg/config/build_options.go",
    "content": "package config\n\n// BuildOptions contains runtime options passed via CLI flags, not from cog.yaml.\n// These are separate from the Config struct because they are not part of the\n// model configuration - they are build-time settings that affect how the\n// container is built but not what's in it.\ntype BuildOptions struct {\n\t// SourceEpochTimestamp is the number of seconds since Unix epoch to use\n\t// for the build timestamp. Set to -1 to disable timestamp rewrites.\n\t// This is useful for reproducible builds.\n\tSourceEpochTimestamp int64\n\n\t// XCachePath is the path to the BuildKit cache directory.\n\t// If empty, inline caching is used instead of local cache.\n\tXCachePath string\n}\n\n// DefaultBuildOptions returns BuildOptions with sensible defaults.\nfunc DefaultBuildOptions() BuildOptions {\n\treturn BuildOptions{\n\t\tSourceEpochTimestamp: -1,\n\t\tXCachePath:           \"\",\n\t}\n}\n"
  },
  {
    "path": "pkg/config/compatibility.go",
    "content": "package config\n\nimport (\n\t// blank import for embeds\n\t_ \"embed\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"sort\"\n\t\"strings\"\n\n\t\"golang.org/x/exp/slices\"\n\n\t\"github.com/replicate/cog/pkg/requirements\"\n\t\"github.com/replicate/cog/pkg/util\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n\n\t\"github.com/replicate/cog/pkg/util/version\"\n)\n\n// TODO(andreas): check tf/py versions. tf 1.5.0 didn't install on py 3.10\n// TODO(andreas): support more tf versions. No matching tensorflow CPU package for version 1.15.4, etc.\n// TODO(andreas): allow user to install versions that aren't compatible\n// TODO(andreas): allow user to install tf cpu package on gpu\n\ntype TFCompatibility struct {\n\tTF           string\n\tTFCPUPackage string\n\tTFGPUPackage string\n\tCUDA         string\n\tCuDNN        string\n\tPythons      []string\n}\n\nfunc (compat *TFCompatibility) UnmarshalJSON(data []byte) error {\n\t// to avoid unmarshalling stack overflow https://stackoverflow.com/questions/34859449/unmarshaljson-results-in-stack-overflow\n\ttype tempType TFCompatibility\n\tc := new(tempType)\n\tif err := json.Unmarshal(data, c); err != nil {\n\t\treturn err\n\t}\n\tcuda := version.MustVersion(c.CUDA)\n\tcuDNN := version.MustVersion(c.CuDNN)\n\tcompat.TF = c.TF\n\tcompat.TFCPUPackage = c.TFCPUPackage\n\tcompat.TFGPUPackage = c.TFGPUPackage\n\t// include minor version\n\tcompat.CUDA = fmt.Sprintf(\"%d.%d\", cuda.Major, cuda.Minor)\n\t// strip cuDNN minor version to match nvidia images\n\tcompat.CuDNN = fmt.Sprintf(\"%d\", cuDNN.Major)\n\tcompat.Pythons = c.Pythons\n\treturn nil\n}\n\ntype TorchCompatibility struct {\n\tTorch         string\n\tTorchvision   string\n\tTorchaudio    string\n\tFindLinks     string\n\tExtraIndexURL string\n\tCUDA          *string\n\tPythons       []string\n}\n\nfunc (c *TorchCompatibility) TorchVersion() string {\n\treturn version.StripModifier(c.Torch)\n}\n\nfunc (c *TorchCompatibility) TorchvisionVersion() string {\n\treturn version.StripModifier(c.Torchvision)\n}\n\ntype CUDABaseImage struct {\n\tTag     string\n\tCUDA    string\n\tCuDNN   string\n\tIsDevel bool\n\tUbuntu  string\n}\n\nfunc (i *CUDABaseImage) ImageTag() string {\n\treturn \"nvidia/cuda:\" + i.Tag\n}\n\n//go:embed cuda_compatibility.json\nvar cudaBaseImagesData []byte\nvar CUDABaseImages []CUDABaseImage\n\n//go:embed tf_compatibility.json\nvar tfCompatibilityMatrixData []byte\nvar TFCompatibilityMatrix []TFCompatibility\n\n//go:embed torch_compatibility.json\nvar torchCompatibilityMatrixData []byte\nvar TorchCompatibilityMatrix []TorchCompatibility\n\nfunc init() {\n\tif err := json.Unmarshal(cudaBaseImagesData, &CUDABaseImages); err != nil {\n\t\tconsole.Fatalf(\"Failed to load embedded CUDA base images: %s\", err)\n\t}\n\n\tif err := json.Unmarshal(tfCompatibilityMatrixData, &TFCompatibilityMatrix); err != nil {\n\t\tconsole.Fatalf(\"Failed to load embedded Tensorflow compatibility matrix: %s\", err)\n\t}\n\n\tvar torchCompatibilityMatrix []TorchCompatibility\n\tif err := json.Unmarshal(torchCompatibilityMatrixData, &torchCompatibilityMatrix); err != nil {\n\t\tconsole.Fatalf(\"Failed to load embedded PyTorch compatibility matrix: %s\", err)\n\t}\n\tfilteredTorchCompatibilityMatrix := []TorchCompatibility{}\n\tfor _, compat := range torchCompatibilityMatrix {\n\t\tfor _, cudaBaseImage := range CUDABaseImages {\n\t\t\tif compat.CUDA == nil || version.Matches(*compat.CUDA, cudaBaseImage.CUDA) {\n\t\t\t\tfilteredTorchCompatibilityMatrix = append(filteredTorchCompatibilityMatrix, compat)\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\tTorchCompatibilityMatrix = filteredTorchCompatibilityMatrix\n}\n\nfunc cudaVersionFromTorchPlusVersion(ver string) (string, string) {\n\tconst cudaVersionPrefix = \"cu\"\n\n\t// Split the version string by the '+' character.\n\tversionParts := strings.Split(ver, \"+\")\n\n\t// If there is no '+' in the version string, return the original string with an empty CUDA version.\n\tif len(versionParts) <= 1 {\n\t\treturn \"\", ver\n\t}\n\n\t// Extract the part after the last '+'.\n\tcudaVersionPart := versionParts[len(versionParts)-1]\n\n\t// Check if the extracted part has the CUDA version prefix.\n\tif !strings.HasPrefix(cudaVersionPart, cudaVersionPrefix) {\n\t\treturn \"\", ver\n\t}\n\n\t// Trim the CUDA version prefix and reformat the version string.\n\tcleanVersion := strings.TrimPrefix(cudaVersionPart, cudaVersionPrefix)\n\tif len(cleanVersion) < 2 {\n\t\treturn \"\", ver // Handle case where cleanVersion is too short to reformat.\n\t}\n\n\t// Insert a dot before the last character to format it as expected.\n\tcleanVersion = cleanVersion[:len(cleanVersion)-1] + \".\" + cleanVersion[len(cleanVersion)-1:]\n\n\t// Return the reformatted CUDA version and the main version.\n\treturn cleanVersion, versionParts[0]\n}\n\nfunc cudasFromTorch(ver string) ([]string, error) {\n\tif ver == \"\" {\n\t\treturn nil, errors.New(\n\t\t\t\"torch version must be specified when using CUDA\",\n\t\t)\n\t}\n\tcudas := []string{}\n\n\t// Check the version modifier on torch (such as +cu118)\n\tcudaVer, ver := cudaVersionFromTorchPlusVersion(ver)\n\tif len(cudaVer) > 0 {\n\t\tfor _, compat := range TorchCompatibilityMatrix {\n\t\t\tif compat.CUDA == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif version.Matches(ver, compat.TorchVersion()) && *compat.CUDA == cudaVer {\n\t\t\t\tcudas = append(cudas, *compat.CUDA)\n\t\t\t\treturn cudas, nil\n\t\t\t}\n\t\t}\n\t}\n\n\tfor _, compat := range TorchCompatibilityMatrix {\n\t\tif version.Matches(ver, compat.TorchVersion()) && compat.CUDA != nil {\n\t\t\tcudas = append(cudas, *compat.CUDA)\n\t\t}\n\t}\n\tslices.Sort(cudas)\n\n\treturn cudas, nil\n}\n\nfunc cudaFromTF(ver string) (cuda string, cuDNN string, err error) {\n\tfor _, compat := range TFCompatibilityMatrix {\n\t\tif ver == compat.TF {\n\t\t\treturn compat.CUDA, compat.CuDNN, nil\n\t\t}\n\t}\n\treturn \"\", \"\", nil\n}\n\nfunc compatibleCuDNNsForCUDA(cuda string) []string {\n\tcuDNNs := []string{}\n\tfor _, image := range CUDABaseImages {\n\t\tif image.CUDA == cuda {\n\t\t\tcuDNNs = append(cuDNNs, image.CuDNN)\n\t\t}\n\t}\n\treturn cuDNNs\n}\n\nfunc defaultCUDA() string {\n\t// TODO: change this to latestTF().CUDA once replicate supports >= 12 everywhere\n\treturn \"11.8\"\n}\n\nfunc latestCUDAFrom(cudas []string) string {\n\tlatest := \"\"\n\tfor _, cuda := range cudas {\n\t\tif latest == \"\" {\n\t\t\tlatest = cuda\n\t\t} else {\n\t\t\tgreater, err := versionGreater(cuda, latest)\n\t\t\tif err != nil {\n\t\t\t\t// should never happen\n\t\t\t\tpanic(fmt.Sprintf(\"Invalid CUDA version: %s\", err))\n\t\t\t}\n\t\t\tif greater {\n\t\t\t\tlatest = cuda\n\t\t\t}\n\t\t}\n\t}\n\treturn latest\n}\n\nfunc latestCuDNNForCUDA(cuda string) (string, error) {\n\tcuDNNs := []string{}\n\tfor _, image := range CUDABaseImages {\n\t\tif version.Matches(cuda, image.CUDA) {\n\t\t\tcuDNNs = append(cuDNNs, image.CuDNN)\n\t\t}\n\t}\n\tsort.Slice(cuDNNs, func(i, j int) bool {\n\t\treturn version.Greater(cuDNNs[i], cuDNNs[j])\n\t})\n\tif len(cuDNNs) == 0 {\n\t\t// TODO: return a list of supported cuda versions\n\t\treturn \"\", fmt.Errorf(\"CUDA %s is not supported by Cog\", cuda)\n\t}\n\treturn cuDNNs[0], nil\n}\n\nfunc versionGreater(a string, b string) (bool, error) {\n\t// TODO(andreas): use library\n\taVer, err := version.NewVersion(a)\n\tif err != nil {\n\t\treturn false, err\n\t}\n\tbVer, err := version.NewVersion(b)\n\tif err != nil {\n\t\treturn false, err\n\t}\n\treturn aVer.Greater(bVer), nil\n}\n\nfunc cudaBaseImageFor(cuda string, cuDNN string) (string, error) {\n\tvar images []CUDABaseImage\n\tfor _, image := range CUDABaseImages {\n\t\tif version.Matches(cuda, image.CUDA) && image.CuDNN == cuDNN {\n\t\t\timages = append(images, image)\n\t\t}\n\t}\n\tif len(images) == 0 {\n\t\treturn \"\", fmt.Errorf(\"no matching base image for CUDA %s and CuDNN %s\", cuda, cuDNN)\n\t}\n\n\tsort.Slice(images, func(i, j int) bool {\n\t\tif images[i].CUDA != images[j].CUDA {\n\t\t\treturn version.MustVersion(images[i].CUDA).Greater(version.MustVersion(images[j].CUDA))\n\t\t}\n\t\treturn images[i].Ubuntu > images[j].Ubuntu\n\t})\n\n\treturn images[0].ImageTag(), nil\n}\n\nfunc tfGPUPackage(ver string, cuda string) (name string, cpuVersion string, err error) {\n\tfor _, compat := range TFCompatibilityMatrix {\n\t\tif compat.TF == ver && version.Equal(compat.CUDA, cuda) {\n\t\t\tname, cpuVersion, _, _, err = requirements.SplitPinnedPythonRequirement(compat.TFGPUPackage)\n\t\t\treturn name, cpuVersion, err\n\t\t}\n\t}\n\t// We've already warned user if they're doing something stupid in validateAndCompleteCUDA(), so fail silently\n\treturn \"\", \"\", nil\n}\n\nfunc torchCPUPackage(ver, goos, goarch string) (name, cpuVersion, findLinks, extraIndexURL string, err error) {\n\tfor _, compat := range TorchCompatibilityMatrix {\n\t\tif compat.TorchVersion() == ver && compat.CUDA == nil {\n\t\t\treturn \"torch\", torchStripCPUSuffixForM1(compat.Torch, goos, goarch), compat.FindLinks, compat.ExtraIndexURL, nil\n\t\t}\n\t}\n\n\t// Fall back to just installing default version. For older pytorch versions, they don't have any CPU versions.\n\treturn \"torch\", ver, \"\", \"\", nil\n}\n\nfunc torchGPUPackage(ver string, cuda string) (name, cpuVersion, findLinks, extraIndexURL string, err error) {\n\t// find the torch package that has the requested torch version and the latest cuda version\n\t// that is at most as high as the requested cuda version\n\tvar latest *TorchCompatibility\n\tfor _, compat := range TorchCompatibilityMatrix {\n\t\tif !version.Matches(compat.TorchVersion(), ver) || compat.CUDA == nil {\n\t\t\tcontinue\n\t\t}\n\t\tgreater, err := versionGreater(*compat.CUDA, cuda)\n\t\tif err != nil {\n\t\t\tpanic(fmt.Sprintf(\"Invalid CUDA version: %s\", err))\n\t\t}\n\n\t\tif greater {\n\t\t\tcontinue\n\t\t}\n\t\tif latest == nil {\n\t\t\tlatest = &compat\n\t\t} else {\n\t\t\tgreater, err := versionGreater(*compat.CUDA, *latest.CUDA)\n\t\t\tif err != nil {\n\t\t\t\t// should never happen\n\t\t\t\tpanic(fmt.Sprintf(\"Invalid CUDA version: %s\", err))\n\t\t\t}\n\t\t\tif greater {\n\t\t\t\tlatest = &compat\n\t\t\t}\n\t\t}\n\t}\n\tif latest == nil {\n\t\t// We've already warned user if they're doing something stupid in validateAndCompleteCUDA()\n\t\treturn \"torch\", ver, \"\", \"\", nil\n\t}\n\n\treturn \"torch\", version.StripModifier(latest.Torch), latest.FindLinks, latest.ExtraIndexURL, nil\n}\n\nfunc torchvisionCPUPackage(ver, goos, goarch string) (name, cpuVersion, findLinks, extraIndexURL string, err error) {\n\tfor _, compat := range TorchCompatibilityMatrix {\n\t\tif compat.TorchvisionVersion() == ver && compat.CUDA == nil {\n\t\t\treturn \"torchvision\", torchStripCPUSuffixForM1(compat.Torchvision, goos, goarch), compat.FindLinks, compat.ExtraIndexURL, nil\n\t\t}\n\t}\n\t// Fall back to just installing default version. For older torchvision versions, they don't have any CPU versions.\n\treturn \"torchvision\", ver, \"\", \"\", nil\n}\n\nfunc torchvisionGPUPackage(ver, cuda string) (name, cpuVersion, findLinks, extraIndexURL string, err error) {\n\t// find the torchvision package that has the requested\n\t// torchvision version and the latest cuda version that is at\n\t// most as high as the requested cuda version\n\tvar latest *TorchCompatibility\n\tfor _, compat := range TorchCompatibilityMatrix {\n\t\tif compat.TorchvisionVersion() != ver || compat.CUDA == nil {\n\t\t\tcontinue\n\t\t}\n\t\tgreater, err := versionGreater(*compat.CUDA, cuda)\n\t\tif err != nil {\n\t\t\tpanic(fmt.Sprintf(\"Invalid CUDA version: %s\", err))\n\t\t}\n\t\tif greater {\n\t\t\tcontinue\n\t\t}\n\t\tif latest == nil {\n\t\t\tlatest = &compat\n\t\t} else {\n\t\t\tgreater, err := versionGreater(*compat.CUDA, *latest.CUDA)\n\t\t\tif err != nil {\n\t\t\t\t// should never happen\n\t\t\t\tpanic(fmt.Sprintf(\"Invalid CUDA version: %s\", err))\n\t\t\t}\n\t\t\tif greater {\n\t\t\t\tlatest = &compat\n\t\t\t}\n\t\t}\n\t}\n\tif latest == nil {\n\t\t// TODO: can we suggest a CUDA version known to be compatible?\n\t\tconsole.Warnf(\"Cog doesn't know if CUDA %s is compatible with torchvision %s. This might cause CUDA problems.\", cuda, ver)\n\t\treturn \"torchvision\", ver, \"\", \"\", nil\n\t}\n\n\treturn \"torchvision\", version.StripModifier(latest.Torchvision), latest.FindLinks, latest.ExtraIndexURL, nil\n}\n\n// aarch64 packages don't have +cpu suffix: https://download.pytorch.org/whl/torch_stable.html\n// TODO(andreas): clean up this hack by actually parsing the torch_stable.html list in the generator\nfunc torchStripCPUSuffixForM1(version string, goos string, goarch string) string {\n\t// TODO(andreas): clean up this hack\n\tif util.IsAppleSiliconMac(goos, goarch) {\n\t\treturn strings.ReplaceAll(version, \"+cpu\", \"\")\n\t}\n\treturn version\n}\n"
  },
  {
    "path": "pkg/config/compatibility_test.go",
    "content": "package config\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestLatestCuDNNForCUDA(t *testing.T) {\n\tactual, err := latestCuDNNForCUDA(\"11.8\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"8\", actual)\n}\n\nfunc TestCudasFromTorchWithCUVersionModifier(t *testing.T) {\n\tcudas, err := cudasFromTorch(\"2.0.1+cu118\")\n\trequire.GreaterOrEqual(t, len(cudas), 1)\n\trequire.Equal(t, cudas[0], \"11.8\")\n\trequire.Nil(t, err)\n}\n"
  },
  {
    "path": "pkg/config/config.go",
    "content": "package config\n\nimport (\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"regexp\"\n\t\"slices\"\n\t\"strconv\"\n\t\"strings\"\n\n\t\"go.yaml.in/yaml/v4\"\n\n\t\"github.com/replicate/cog/pkg/requirements\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n\t\"github.com/replicate/cog/pkg/util/version\"\n)\n\nvar (\n\tBuildSourceEpochTimestamp int64 = -1\n\tBuildXCachePath           string\n\tPipPackageNameRegex       = regexp.MustCompile(`^([^>=<~ \\n[#]+)`)\n)\n\n// TODO(andreas): support conda packages\n// TODO(andreas): support dockerfiles\n// TODO(andreas): custom cpu/gpu installs\n// TODO(andreas): suggest valid torchvision versions (e.g. if the user wants to use 0.8.0, suggest 0.8.1)\n\nconst (\n\tMinimumMajorPythonVersion               int    = 3\n\tMinimumMinorPythonVersion               int    = 10\n\tMinimumMinorPythonVersionForConcurrency int    = 11\n\tMinimumMajorCudaVersion                 int    = 11\n\tDefaultPythonVersion                    string = \"3.13\"\n)\n\ntype RunItem struct {\n\tCommand string `json:\"command,omitempty\" yaml:\"command\"`\n\tMounts  []struct {\n\t\tType   string `json:\"type,omitempty\" yaml:\"type\"`\n\t\tID     string `json:\"id,omitempty\" yaml:\"id\"`\n\t\tTarget string `json:\"target,omitempty\" yaml:\"target\"`\n\t} `json:\"mounts,omitempty\" yaml:\"mounts,omitempty\"`\n}\n\ntype Build struct {\n\tGPU                bool      `json:\"gpu,omitempty\" yaml:\"gpu,omitempty\"`\n\tPythonVersion      string    `json:\"python_version,omitempty\" yaml:\"python_version\"`\n\tPythonRequirements string    `json:\"python_requirements,omitempty\" yaml:\"python_requirements,omitempty\"`\n\tPythonPackages     []string  `json:\"python_packages,omitempty\" yaml:\"python_packages,omitempty\"` // Deprecated, but included for backwards compatibility\n\tRun                []RunItem `json:\"run,omitempty\" yaml:\"run,omitempty\"`\n\tSystemPackages     []string  `json:\"system_packages,omitempty\" yaml:\"system_packages,omitempty\"`\n\tPreInstall         []string  `json:\"pre_install,omitempty\" yaml:\"pre_install,omitempty\"` // Deprecated, but included for backwards compatibility\n\tCUDA               string    `json:\"cuda,omitempty\" yaml:\"cuda,omitempty\"`\n\tCuDNN              string    `json:\"cudnn,omitempty\" yaml:\"cudnn,omitempty\"`\n\t// SDKVersion pins the cog Python SDK version installed in the container.\n\t// Accepts a PEP 440 version string (e.g. \"0.18.0\" or \"0.18.0a1\").\n\t// When empty the latest release is installed. Overridden by COG_SDK_WHEEL env var.\n\tSDKVersion string `json:\"sdk_version,omitempty\" yaml:\"sdk_version,omitempty\"`\n\n\tpythonRequirementsContent []string\n}\n\ntype Concurrency struct {\n\tMax int `json:\"max,omitempty\" yaml:\"max\"`\n}\n\n// WeightSource defines a weight file or directory to include in the model.\ntype WeightSource struct {\n\tName   string `json:\"name,omitempty\" yaml:\"name,omitempty\"`\n\tSource string `json:\"source\" yaml:\"source\"`\n\tTarget string `json:\"target,omitempty\" yaml:\"target,omitempty\"`\n}\n\ntype Config struct {\n\tBuild       *Build         `json:\"build\" yaml:\"build\"`\n\tImage       string         `json:\"image,omitempty\" yaml:\"image,omitempty\"`\n\tPredict     string         `json:\"predict,omitempty\" yaml:\"predict\"`\n\tTrain       string         `json:\"train,omitempty\" yaml:\"train,omitempty\"`\n\tConcurrency *Concurrency   `json:\"concurrency,omitempty\" yaml:\"concurrency,omitempty\"`\n\tEnvironment []string       `json:\"environment,omitempty\" yaml:\"environment,omitempty\"`\n\tWeights     []WeightSource `json:\"weights,omitempty\" yaml:\"weights,omitempty\"`\n\n\tparsedEnvironment map[string]string\n}\n\nfunc defaultConfig() *Config {\n\treturn &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           false,\n\t\t\tPythonVersion: \"3.13\",\n\t\t},\n\t}\n}\n\nfunc (r *RunItem) UnmarshalYAML(unmarshal func(any) error) error {\n\tvar commandOrMap any\n\tif err := unmarshal(&commandOrMap); err != nil {\n\t\treturn err\n\t}\n\n\tswitch v := commandOrMap.(type) {\n\tcase string:\n\t\tr.Command = v\n\tcase map[string]any:\n\t\tvar data []byte\n\t\tvar err error\n\n\t\tif data, err = yaml.Marshal(v); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\taux := struct {\n\t\t\tCommand string `yaml:\"command\"`\n\t\t\tMounts  []struct {\n\t\t\t\tType   string `yaml:\"type\"`\n\t\t\t\tID     string `yaml:\"id\"`\n\t\t\t\tTarget string `yaml:\"target\"`\n\t\t\t} `yaml:\"mounts,omitempty\"`\n\t\t}{}\n\n\t\tif err := yaml.Unmarshal(data, &aux); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\t*r = RunItem(aux)\n\tdefault:\n\t\treturn fmt.Errorf(\"unexpected type %T for RunItem\", v)\n\t}\n\n\treturn nil\n}\n\nfunc (r *RunItem) UnmarshalJSON(data []byte) error {\n\tvar commandOrMap any\n\tif err := json.Unmarshal(data, &commandOrMap); err != nil {\n\t\treturn err\n\t}\n\n\tswitch v := commandOrMap.(type) {\n\tcase string:\n\t\tr.Command = v\n\tcase map[string]any:\n\t\taux := struct {\n\t\t\tCommand string `json:\"command\"`\n\t\t\tMounts  []struct {\n\t\t\t\tType   string `json:\"type\"`\n\t\t\t\tID     string `json:\"id\"`\n\t\t\t\tTarget string `json:\"target\"`\n\t\t\t} `json:\"mounts,omitempty\"`\n\t\t}{}\n\n\t\tjsonData, err := json.Marshal(v)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tif err := json.Unmarshal(jsonData, &aux); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\t*r = RunItem(aux)\n\tdefault:\n\t\treturn fmt.Errorf(\"unexpected type %T for RunItem\", v)\n\t}\n\n\treturn nil\n}\n\nfunc (c *Config) CUDABaseImageTag() (string, error) {\n\treturn cudaBaseImageFor(c.Build.CUDA, c.Build.CuDNN)\n}\n\nfunc (c *Config) TorchVersion() (string, bool) {\n\treturn c.pythonPackageVersion(\"torch\")\n}\n\nfunc (c *Config) TorchvisionVersion() (string, bool) {\n\treturn c.pythonPackageVersion(\"torchvision\")\n}\n\nfunc (c *Config) TorchaudioVersion() (string, bool) {\n\treturn c.pythonPackageVersion(\"torchaudio\")\n}\n\nfunc (c *Config) TensorFlowVersion() (string, bool) {\n\treturn c.pythonPackageVersion(\"tensorflow\")\n}\n\nfunc (c *Config) cudasFromTorch() (torchVersion string, torchCUDAs []string, err error) {\n\tif version, ok := c.TorchVersion(); ok {\n\t\tcudas, err := cudasFromTorch(version)\n\t\tif err != nil {\n\t\t\treturn \"\", nil, err\n\t\t}\n\t\treturn version, cudas, nil\n\t}\n\treturn \"\", nil, nil\n}\n\nfunc (c *Config) cudaFromTF() (tfVersion string, tfCUDA string, tfCuDNN string, err error) {\n\tif version, ok := c.TensorFlowVersion(); ok {\n\t\tcuda, cudnn, err := cudaFromTF(version)\n\t\tif err != nil {\n\t\t\treturn \"\", \"\", \"\", err\n\t\t}\n\t\treturn version, cuda, cudnn, nil\n\t}\n\treturn \"\", \"\", \"\", nil\n}\n\nfunc (c *Config) pythonPackageVersion(name string) (version string, ok bool) {\n\tfor _, pkg := range c.Build.pythonRequirementsContent {\n\t\tpkgName := requirements.PackageName(pkg)\n\t\tif pkgName == name {\n\t\t\tversions := requirements.Versions(pkg)\n\t\t\tif len(versions) > 0 {\n\t\t\t\treturn versions[0], true\n\t\t\t}\n\t\t\treturn \"\", true\n\t\t}\n\t}\n\treturn \"\", false\n}\n\nfunc splitPythonVersion(version string) (major int, minor int, err error) {\n\tversion = strings.TrimSpace(version)\n\tparts := strings.SplitN(version, \".\", 3)\n\tif len(parts) < 2 {\n\t\treturn 0, 0, fmt.Errorf(\"missing minor version in %s\", version)\n\t}\n\tmajorStr, minorStr := parts[0], parts[1]\n\tmajor, err = strconv.Atoi(majorStr)\n\tif err != nil {\n\t\treturn 0, 0, err\n\t}\n\tminor, err = strconv.Atoi(minorStr)\n\tif err != nil {\n\t\treturn 0, 0, err\n\t}\n\treturn major, minor, nil\n}\n\n// Complete performs CUDA resolution, requirements loading, and environment loading for a Config.\n// Use this when building a Config struct directly (not from YAML).\n// For configs loaded from YAML, use Load() instead which handles validation and completion.\nfunc (c *Config) Complete(projectDir string) error {\n\t// Validate mutual exclusion of python_packages and python_requirements\n\tif len(c.Build.PythonPackages) > 0 && c.Build.PythonRequirements != \"\" {\n\t\treturn fmt.Errorf(\"only one of python_packages or python_requirements can be set in your cog.yaml, not both\")\n\t}\n\n\t// Load python_requirements into memory to simplify reading it multiple times\n\tif c.Build.PythonRequirements != \"\" {\n\t\trequirementsFilePath := c.Build.PythonRequirements\n\t\tif !strings.HasPrefix(requirementsFilePath, \"/\") {\n\t\t\trequirementsFilePath = filepath.Join(projectDir, c.Build.PythonRequirements)\n\t\t}\n\t\treqs, err := requirements.ReadRequirements(requirementsFilePath)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"failed to open python_requirements file: %w\", err)\n\t\t}\n\t\tc.Build.pythonRequirementsContent = reqs\n\t} else if len(c.Build.PythonPackages) > 0 {\n\t\t// Backwards compatibility: if using deprecated python_packages, populate requirements content\n\t\tc.Build.pythonRequirementsContent = c.Build.PythonPackages\n\t}\n\n\t// Resolve CUDA/CuDNN versions if GPU is enabled\n\tif c.Build.GPU {\n\t\tif err := c.validateAndCompleteCUDA(); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\t// Parse and validate environment variables\n\tif err := c.loadEnvironment(); err != nil {\n\t\treturn err\n\t}\n\n\treturn nil\n}\n\n// PythonRequirementsForArch returns a requirements.txt file with all the GPU packages resolved for given OS and architecture.\nfunc (c *Config) PythonRequirementsForArch(goos string, goarch string, includePackages []string) (string, error) {\n\tpackages := []string{}\n\tfindLinksSet := map[string]bool{}\n\textraIndexURLSet := map[string]bool{}\n\n\tincludePackageNames := []string{}\n\tfor _, pkg := range includePackages {\n\t\tpackageName := requirements.PackageName(pkg)\n\t\tincludePackageNames = append(includePackageNames, packageName)\n\t}\n\n\t// Include all the requirements and remove our include packages if they exist\n\tfor _, pkg := range c.Build.pythonRequirementsContent {\n\t\tarchPkg, findLinksList, extraIndexURLs, err := c.pythonPackageForArch(pkg, goos, goarch)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tpackages = append(packages, archPkg)\n\t\tif len(findLinksList) > 0 {\n\t\t\tfor _, fl := range findLinksList {\n\t\t\t\tfindLinksSet[fl] = true\n\t\t\t}\n\t\t}\n\t\tif len(extraIndexURLs) > 0 {\n\t\t\tfor _, u := range extraIndexURLs {\n\t\t\t\textraIndexURLSet[u] = true\n\t\t\t}\n\t\t}\n\n\t\tpackageName := requirements.PackageName(archPkg)\n\t\tif packageName != \"\" {\n\t\t\tfoundIdx := -1\n\t\t\tfor i, includePkg := range includePackageNames {\n\t\t\t\tif includePkg == packageName {\n\t\t\t\t\tfoundIdx = i\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\t\tif foundIdx != -1 {\n\t\t\t\tincludePackageNames = append(includePackageNames[:foundIdx], includePackageNames[foundIdx+1:]...)\n\t\t\t\tincludePackages = append(includePackages[:foundIdx], includePackages[foundIdx+1:]...)\n\t\t\t}\n\t\t}\n\t}\n\n\t// If we still have some include packages add them in\n\tpackages = append(packages, includePackages...)\n\n\t// Create final requirements.txt output\n\t// Put index URLs first\n\tlines := []string{}\n\tfor findLinks := range findLinksSet {\n\t\tlines = append(lines, \"--find-links \"+findLinks)\n\t}\n\tfor extraIndexURL := range extraIndexURLSet {\n\t\tlines = append(lines, \"--extra-index-url \"+extraIndexURL)\n\t}\n\n\t// Then, everything else\n\tlines = append(lines, packages...)\n\n\treturn strings.Join(lines, \"\\n\"), nil\n}\n\n// pythonPackageForArch takes a package==version line and\n// returns a package==version and index URL resolved to the correct GPU package for the given OS and architecture\nfunc (c *Config) pythonPackageForArch(pkg, goos, goarch string) (actualPackage string, findLinksList []string, extraIndexURLs []string, err error) {\n\tname, version, findLinksList, extraIndexURLs, err := requirements.SplitPinnedPythonRequirement(pkg)\n\tif err != nil {\n\t\t// It's not pinned, so just return the line verbatim\n\t\treturn pkg, []string{}, []string{}, nil\n\t}\n\tif len(extraIndexURLs) > 0 {\n\t\treturn name + \"==\" + version, findLinksList, extraIndexURLs, nil\n\t}\n\n\textraIndexURL := \"\"\n\tfindLinks := \"\"\n\tswitch name {\n\tcase \"tensorflow\":\n\t\tif c.Build.GPU {\n\t\t\tname, version, err = tfGPUPackage(version, c.Build.CUDA)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", nil, nil, err\n\t\t\t}\n\t\t}\n\t\t// There is no CPU case for tensorflow because the default package is just the CPU package, so no transformation of version is needed\n\tcase \"torch\":\n\t\tif c.Build.GPU {\n\t\t\tname, version, findLinks, extraIndexURL, err = torchGPUPackage(version, c.Build.CUDA)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", nil, nil, err\n\t\t\t}\n\t\t} else {\n\t\t\tname, version, findLinks, extraIndexURL, err = torchCPUPackage(version, goos, goarch)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", nil, nil, err\n\t\t\t}\n\t\t}\n\tcase \"torchvision\":\n\t\tif c.Build.GPU {\n\t\t\tname, version, findLinks, extraIndexURL, err = torchvisionGPUPackage(version, c.Build.CUDA)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", nil, nil, err\n\t\t\t}\n\t\t} else {\n\t\t\tname, version, findLinks, extraIndexURL, err = torchvisionCPUPackage(version, goos, goarch)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", nil, nil, err\n\t\t\t}\n\t\t}\n\t}\n\tpkgWithVersion := name\n\tif version != \"\" {\n\t\tpkgWithVersion += \"==\" + version\n\t}\n\tif extraIndexURL != \"\" {\n\t\textraIndexURLs = []string{extraIndexURL}\n\t}\n\tif findLinks != \"\" {\n\t\tfindLinksList = []string{findLinks}\n\t}\n\treturn pkgWithVersion, findLinksList, extraIndexURLs, nil\n}\n\nfunc validateCudaVersion(cudaVersion string) error {\n\tparts := strings.Split(cudaVersion, \".\")\n\tif len(parts) < 2 {\n\t\treturn fmt.Errorf(\"CUDA version %q must include both major and minor versions\", cudaVersion)\n\t}\n\n\tmajor, err := strconv.Atoi(parts[0])\n\tif err != nil {\n\t\treturn fmt.Errorf(\"invalid major version in CUDA version %q\", cudaVersion)\n\t}\n\n\tif major < MinimumMajorCudaVersion {\n\t\treturn fmt.Errorf(\"minimum supported CUDA version is %d, requested %q\", MinimumMajorCudaVersion, cudaVersion)\n\t}\n\treturn nil\n}\n\nfunc (c *Config) validateAndCompleteCUDA() error {\n\tif c.Build.CUDA != \"\" {\n\t\tif err := validateCudaVersion(c.Build.CUDA); err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\tif c.Build.CUDA != \"\" && c.Build.CuDNN != \"\" {\n\t\tcompatibleCuDNNs := compatibleCuDNNsForCUDA(c.Build.CUDA)\n\t\tif !slices.Contains(compatibleCuDNNs, c.Build.CuDNN) {\n\t\t\treturn fmt.Errorf(`the specified CUDA version %s is not compatible with CuDNN %s.\nCompatible CuDNN versions are: %s`, c.Build.CUDA, c.Build.CuDNN, strings.Join(compatibleCuDNNs, \",\"))\n\t\t}\n\t}\n\n\ttorchVersion, torchCUDAs, err := c.cudasFromTorch()\n\tif err != nil {\n\t\treturn err\n\t}\n\ttfVersion, tfCUDA, tfCuDNN, err := c.cudaFromTF()\n\tif err != nil {\n\t\treturn err\n\t}\n\t// The pre-compiled TensorFlow binaries requires specific CUDA/CuDNN versions to be\n\t// installed, but Torch bundles their own CUDA/CuDNN libraries.\n\n\tswitch {\n\tcase tfVersion != \"\":\n\t\tswitch {\n\t\tcase c.Build.CUDA == \"\":\n\t\t\tif tfCuDNN == \"\" {\n\t\t\t\treturn fmt.Errorf(\"cog doesn't know what CUDA version is compatible with tensorflow==%s. You might need to upgrade Cog: https://github.com/replicate/cog#upgrade\\n\\nIf that doesn't work, you need to set the 'cuda' option in cog.yaml to set what version to use. You might be able to find this out from https://www.tensorflow.org/\", tfVersion)\n\t\t\t}\n\t\t\tconsole.Debugf(\"Setting CUDA to version %s from Tensorflow version\", tfCUDA)\n\t\t\tc.Build.CUDA = tfCUDA\n\t\tcase tfCUDA == \"\" || version.EqualMinor(tfCUDA, c.Build.CUDA):\n\t\t\tconsole.Warnf(\"Cog doesn't know if CUDA %s is compatible with Tensorflow %s. This might cause CUDA problems.\", c.Build.CUDA, tfVersion)\n\t\t\tif tfCUDA != \"\" {\n\t\t\t\tconsole.Warnf(\"Try %s instead?\", tfCUDA)\n\t\t\t}\n\t\t}\n\n\t\tswitch {\n\t\tcase c.Build.CuDNN == \"\" && tfCuDNN != \"\":\n\t\t\tconsole.Debugf(\"Setting CuDNN to version %s from Tensorflow version\", tfCuDNN)\n\t\t\tc.Build.CuDNN = tfCuDNN\n\t\tcase c.Build.CuDNN == \"\":\n\t\t\tc.Build.CuDNN, err = latestCuDNNForCUDA(c.Build.CUDA)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tconsole.Debugf(\"Setting CuDNN to version %s\", c.Build.CUDA)\n\t\tcase tfCuDNN != c.Build.CuDNN:\n\t\t\tconsole.Warnf(\"Cog doesn't know if cuDNN %s is compatible with Tensorflow %s. This might cause CUDA problems.\", c.Build.CuDNN, tfVersion)\n\t\t\treturn fmt.Errorf(`the specified cuDNN version %s is not compatible with tensorflow==%s.\nCompatible cuDNN version is: %s`, c.Build.CuDNN, tfVersion, tfCuDNN)\n\t\t}\n\tcase torchVersion != \"\":\n\t\tswitch {\n\t\tcase c.Build.CUDA == \"\":\n\t\t\tif len(torchCUDAs) == 0 {\n\t\t\t\treturn fmt.Errorf(\"cog doesn't know what CUDA version is compatible with torch==%s. You might need to upgrade Cog: https://github.com/replicate/cog#upgrade\\n\\nIf that doesn't work, you need to set the 'cuda' option in cog.yaml to set what version to use. You might be able to find this out from https://pytorch.org/\", torchVersion)\n\t\t\t}\n\t\t\tc.Build.CUDA = latestCUDAFrom(torchCUDAs)\n\t\t\tconsole.Debugf(\"Setting CUDA to version %s from Torch version\", c.Build.CUDA)\n\t\tcase !slices.ContainsFunc(torchCUDAs, func(torchCUDA string) bool { return version.EqualMinor(torchCUDA, c.Build.CUDA) }):\n\t\t\t// TODO: can we suggest a CUDA version known to be compatible?\n\t\t\tconsole.Warnf(\"Cog doesn't know if CUDA %s is compatible with PyTorch %s. This might cause CUDA problems.\", c.Build.CUDA, torchVersion)\n\t\t\tif len(torchCUDAs) > 0 {\n\t\t\t\tconsole.Warnf(\"Try %s instead?\", torchCUDAs[len(torchCUDAs)-1])\n\t\t\t}\n\t\t}\n\n\t\tif c.Build.CuDNN == \"\" {\n\t\t\tc.Build.CuDNN, err = latestCuDNNForCUDA(c.Build.CUDA)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tconsole.Debugf(\"Setting CuDNN to version %s\", c.Build.CUDA)\n\t\t}\n\tdefault:\n\t\tif c.Build.CUDA == \"\" {\n\t\t\tc.Build.CUDA = defaultCUDA()\n\t\t\tconsole.Debugf(\"Setting CUDA to version %s\", c.Build.CUDA)\n\t\t}\n\t\tif c.Build.CuDNN == \"\" {\n\t\t\tc.Build.CuDNN, err = latestCuDNNForCUDA(c.Build.CUDA)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tconsole.Debugf(\"Setting CuDNN to version %s\", c.Build.CUDA)\n\t\t}\n\t}\n\n\treturn nil\n}\n\nfunc (c *Config) RequirementsFile(projectDir string) string {\n\treturn filepath.Join(projectDir, c.Build.PythonRequirements)\n}\n\nfunc (c *Config) ParsedEnvironment() map[string]string {\n\treturn c.parsedEnvironment\n}\n\nfunc (c *Config) loadEnvironment() error {\n\tenv, err := parseAndValidateEnvironment(c.Environment)\n\tif err != nil {\n\t\treturn err\n\t}\n\tc.parsedEnvironment = env\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/config/config_file.go",
    "content": "package config\n\nimport (\n\t\"encoding/json\"\n\t\"fmt\"\n\n\t\"go.yaml.in/yaml/v4\"\n)\n\n// configFile represents the raw cog.yaml as written by users.\n// All fields are pointers/omitempty to distinguish \"not set\" from \"set to zero value\".\n// This struct is only used during parsing - validation produces errors,\n// completion produces a Config.\ntype configFile struct {\n\tBuild       *buildFile       `json:\"build,omitempty\" yaml:\"build,omitempty\"`\n\tImage       *string          `json:\"image,omitempty\" yaml:\"image,omitempty\"`\n\tPredict     *string          `json:\"predict,omitempty\" yaml:\"predict,omitempty\"`\n\tTrain       *string          `json:\"train,omitempty\" yaml:\"train,omitempty\"`\n\tConcurrency *concurrencyFile `json:\"concurrency,omitempty\" yaml:\"concurrency,omitempty\"`\n\tEnvironment []string         `json:\"environment,omitempty\" yaml:\"environment,omitempty\"`\n\tWeights     []weightFile     `json:\"weights,omitempty\" yaml:\"weights,omitempty\"`\n}\n\n// buildFile represents the raw build configuration from cog.yaml.\ntype buildFile struct {\n\tGPU                *bool         `json:\"gpu,omitempty\" yaml:\"gpu,omitempty\"`\n\tPythonVersion      *string       `json:\"python_version,omitempty\" yaml:\"python_version,omitempty\"`\n\tPythonRequirements *string       `json:\"python_requirements,omitempty\" yaml:\"python_requirements,omitempty\"`\n\tRun                []runItemFile `json:\"run,omitempty\" yaml:\"run,omitempty\"`\n\tSystemPackages     []string      `json:\"system_packages,omitempty\" yaml:\"system_packages,omitempty\"`\n\tCUDA               *string       `json:\"cuda,omitempty\" yaml:\"cuda,omitempty\"`\n\tCuDNN              *string       `json:\"cudnn,omitempty\" yaml:\"cudnn,omitempty\"`\n\tSDKVersion         *string       `json:\"sdk_version,omitempty\" yaml:\"sdk_version,omitempty\"`\n\n\t// Deprecated fields - parsed with warnings\n\tPythonPackages []string `json:\"python_packages,omitempty\" yaml:\"python_packages,omitempty\"`\n\tPreInstall     []string `json:\"pre_install,omitempty\" yaml:\"pre_install,omitempty\"`\n}\n\n// runItemFile represents a run command which can be either a string or an object.\ntype runItemFile struct {\n\tCommand string      `json:\"command,omitempty\" yaml:\"command,omitempty\"`\n\tMounts  []mountFile `json:\"mounts,omitempty\" yaml:\"mounts,omitempty\"`\n}\n\n// mountFile represents a mount configuration in a run command.\ntype mountFile struct {\n\tType   string `json:\"type,omitempty\" yaml:\"type,omitempty\"`\n\tID     string `json:\"id,omitempty\" yaml:\"id,omitempty\"`\n\tTarget string `json:\"target,omitempty\" yaml:\"target,omitempty\"`\n}\n\n// weightFile represents a weight source configuration.\ntype weightFile struct {\n\tName   string `json:\"name,omitempty\" yaml:\"name,omitempty\"`\n\tSource string `json:\"source\" yaml:\"source\"`\n\tTarget string `json:\"target,omitempty\" yaml:\"target,omitempty\"`\n}\n\n// concurrencyFile represents concurrency configuration.\ntype concurrencyFile struct {\n\tMax *int `json:\"max,omitempty\" yaml:\"max,omitempty\"`\n}\n\n// UnmarshalYAML implements custom YAML unmarshaling for runItemFile\n// to support both string and object forms.\nfunc (r *runItemFile) UnmarshalYAML(unmarshal func(any) error) error {\n\tvar commandOrMap any\n\tif err := unmarshal(&commandOrMap); err != nil {\n\t\treturn err\n\t}\n\n\tswitch v := commandOrMap.(type) {\n\tcase string:\n\t\tr.Command = v\n\tcase map[string]any:\n\t\tvar data []byte\n\t\tvar err error\n\n\t\tif data, err = yaml.Marshal(v); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\taux := struct {\n\t\t\tCommand string `yaml:\"command\"`\n\t\t\tMounts  []struct {\n\t\t\t\tType   string `yaml:\"type\"`\n\t\t\t\tID     string `yaml:\"id\"`\n\t\t\t\tTarget string `yaml:\"target\"`\n\t\t\t} `yaml:\"mounts,omitempty\"`\n\t\t}{}\n\n\t\tif err := yaml.Unmarshal(data, &aux); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tr.Command = aux.Command\n\t\tr.Mounts = make([]mountFile, len(aux.Mounts))\n\t\tfor i, m := range aux.Mounts {\n\t\t\tr.Mounts[i] = mountFile{\n\t\t\t\tType:   m.Type,\n\t\t\t\tID:     m.ID,\n\t\t\t\tTarget: m.Target,\n\t\t\t}\n\t\t}\n\tdefault:\n\t\treturn fmt.Errorf(\"unexpected type %T for runItemFile\", v)\n\t}\n\n\treturn nil\n}\n\n// UnmarshalJSON implements custom JSON unmarshaling for runItemFile\n// to support both string and object forms.\nfunc (r *runItemFile) UnmarshalJSON(data []byte) error {\n\tvar commandOrMap any\n\tif err := json.Unmarshal(data, &commandOrMap); err != nil {\n\t\treturn err\n\t}\n\n\tswitch v := commandOrMap.(type) {\n\tcase string:\n\t\tr.Command = v\n\tcase map[string]any:\n\t\taux := struct {\n\t\t\tCommand string `json:\"command\"`\n\t\t\tMounts  []struct {\n\t\t\t\tType   string `json:\"type\"`\n\t\t\t\tID     string `json:\"id\"`\n\t\t\t\tTarget string `json:\"target\"`\n\t\t\t} `json:\"mounts,omitempty\"`\n\t\t}{}\n\n\t\tjsonData, err := json.Marshal(v)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tif err := json.Unmarshal(jsonData, &aux); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tr.Command = aux.Command\n\t\tr.Mounts = make([]mountFile, len(aux.Mounts))\n\t\tfor i, m := range aux.Mounts {\n\t\t\tr.Mounts[i] = mountFile{\n\t\t\t\tType:   m.Type,\n\t\t\t\tID:     m.ID,\n\t\t\t\tTarget: m.Target,\n\t\t\t}\n\t\t}\n\tdefault:\n\t\treturn fmt.Errorf(\"unexpected type %T for runItemFile\", v)\n\t}\n\n\treturn nil\n}\n\n// Helper functions for working with configFile\n\n// GetGPU returns the GPU setting, defaulting to false if not set.\nfunc (b *buildFile) GetGPU() bool {\n\tif b == nil || b.GPU == nil {\n\t\treturn false\n\t}\n\treturn *b.GPU\n}\n"
  },
  {
    "path": "pkg/config/config_test.go",
    "content": "package config\n\nimport (\n\t\"encoding/json\"\n\t\"os\"\n\t\"path\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/hashicorp/go-version\"\n\t\"github.com/stretchr/testify/require\"\n\t\"go.yaml.in/yaml/v4\"\n)\n\nfunc TestValidateCudaVersion(t *testing.T) {\n\ttestCases := []struct {\n\t\tname        string\n\t\tinput       string\n\t\texpectedErr bool\n\t}{\n\t\t{\n\t\t\tname:        \"ValidVersion\",\n\t\t\tinput:       \"12.4\",\n\t\t\texpectedErr: false,\n\t\t},\n\t\t{\n\t\t\tname:        \"MinimumVersion\",\n\t\t\tinput:       \"11.0\",\n\t\t\texpectedErr: false,\n\t\t},\n\t\t{\n\t\t\tname:        \"FullyQualifiedVersion\",\n\t\t\tinput:       \"12.4.1\",\n\t\t\texpectedErr: false,\n\t\t},\n\t\t{\n\t\t\tname:        \"InvalidFormat\",\n\t\t\tinput:       \"11-2\",\n\t\t\texpectedErr: true,\n\t\t},\n\t\t{\n\t\t\tname:        \"InvalidMissingMinor\",\n\t\t\tinput:       \"11\",\n\t\t\texpectedErr: true,\n\t\t},\n\t\t{\n\t\t\tname:        \"LessThanMinimum\",\n\t\t\tinput:       \"9.1\",\n\t\t\texpectedErr: true,\n\t\t},\n\t}\n\n\tfor _, tc := range testCases {\n\t\tt.Run(tc.name, func(t *testing.T) {\n\t\t\terr := validateCudaVersion(tc.input)\n\t\t\tif tc.expectedErr {\n\t\t\t\trequire.Error(t, err)\n\t\t\t} else {\n\t\t\t\trequire.NoError(t, err)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc assertMinorVersion(t *testing.T, expected, actual string) {\n\texpectedVersion, err := version.NewVersion(expected)\n\tif err != nil {\n\t\tt.Errorf(\"Error parsing version: %v\", err)\n\t\treturn\n\t}\n\tactualVersion, err := version.NewVersion(actual)\n\tif err != nil {\n\t\tt.Errorf(\"Error parsing version: %v\", err)\n\t\treturn\n\t}\n\n\t// Compare only the major and minor parts\n\tif expectedVersion.Segments()[0] != actualVersion.Segments()[0] || expectedVersion.Segments()[1] != actualVersion.Segments()[1] {\n\t\tt.Errorf(\"Expected %s but got %s\", expected, actual)\n\t}\n}\n\nfunc TestPythonPackagesAndRequirementsCantBeUsedTogether(t *testing.T) {\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"replicate==1.0.0\",\n\t\t\t},\n\t\t\tPythonRequirements: \"requirements.txt\",\n\t\t},\n\t}\n\terr := config.Complete(\"\")\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"only one of python_packages or python_requirements can be set in your cog.yaml, not both\")\n}\n\nfunc TestPythonRequirementsResolvesPythonPackagesAndCudaVersions(t *testing.T) {\n\ttmpDir := t.TempDir()\n\terr := os.WriteFile(path.Join(tmpDir, \"requirements.txt\"), []byte(`torch==1.13.1\ntorchvision==0.14.1\ntorchaudio==0.13.1\nfoo==1.0.0`), 0o644)\n\trequire.NoError(t, err)\n\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:                true,\n\t\t\tPythonVersion:      \"3.10\",\n\t\t\tPythonRequirements: \"requirements.txt\",\n\t\t},\n\t}\n\terr = config.Complete(tmpDir)\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"11.7\", config.Build.CUDA)\n\trequire.Equal(t, \"8\", config.Build.CuDNN)\n\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{})\n\trequire.NoError(t, err)\n\texpected := `--extra-index-url https://download.pytorch.org/whl/cu117\ntorch==1.13.1\ntorchvision==0.14.1\ntorchaudio==0.13.1\nfoo==1.0.0`\n\trequire.Equal(t, expected, requirements)\n}\n\nfunc TestPythonRequirementsResolvesPythonPackagesAndCudaVersionsWithExtraIndexURL(t *testing.T) {\n\ttmpDir := t.TempDir()\n\terr := os.WriteFile(path.Join(tmpDir, \"requirements.txt\"), []byte(`torch==1.12.1\ntorchvision==0.13.1\ntorchaudio==0.12.1\nfoo==1.0.0`), 0o644)\n\trequire.NoError(t, err)\n\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:                true,\n\t\t\tPythonVersion:      \"3.10\",\n\t\t\tPythonRequirements: \"requirements.txt\",\n\t\t},\n\t}\n\terr = config.Complete(tmpDir)\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"11.6\", config.Build.CUDA)\n\trequire.Equal(t, \"8\", config.Build.CuDNN)\n\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{})\n\trequire.NoError(t, err)\n\texpected := `--extra-index-url https://download.pytorch.org/whl/cu116\ntorch==1.12.1\ntorchvision==0.13.1\ntorchaudio==0.12.1\nfoo==1.0.0`\n\trequire.Equal(t, expected, requirements)\n}\n\nfunc TestPythonRequirementsWorksWithLinesCogCannotParse(t *testing.T) {\n\ttmpDir := t.TempDir()\n\terr := os.WriteFile(path.Join(tmpDir, \"requirements.txt\"), []byte(`foo==1.0.0\n# complex requirements\nfastapi>=0.6,<1\nflask>0.4\n# comments!\n# blank lines!\n\n# arguments\n-f http://example.com`), 0o644)\n\trequire.NoError(t, err)\n\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:                true,\n\t\t\tPythonVersion:      \"3.10\",\n\t\t\tPythonRequirements: \"requirements.txt\",\n\t\t},\n\t}\n\terr = config.Complete(tmpDir)\n\trequire.NoError(t, err)\n\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{})\n\trequire.NoError(t, err)\n\texpected := `foo==1.0.0\nfastapi>=0.6,<1\nflask>0.4\n-f http://example.com`\n\trequire.Equal(t, expected, requirements)\n\n}\n\nfunc TestValidateAndCompleteCUDAForAllTF(t *testing.T) {\n\tfor _, compat := range TFCompatibilityMatrix {\n\t\tconfig := &Config{\n\t\t\tBuild: &Build{\n\t\t\t\tGPU:           true,\n\t\t\t\tPythonVersion: \"3.10\",\n\t\t\t\tPythonPackages: []string{\n\t\t\t\t\t\"tensorflow==\" + compat.TF,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\terr := config.Complete(\"\")\n\t\trequire.NoError(t, err)\n\t\tassertMinorVersion(t, compat.CUDA, config.Build.CUDA)\n\t\trequire.Equal(t, compat.CuDNN, config.Build.CuDNN)\n\t}\n}\n\nfunc TestValidateAndCompleteCUDAForAllTorch(t *testing.T) {\n\tfor _, compat := range TorchCompatibilityMatrix {\n\t\tconfig := &Config{\n\t\t\tBuild: &Build{\n\t\t\t\tGPU:           compat.CUDA != nil,\n\t\t\t\tPythonVersion: \"3.10\",\n\t\t\t\tPythonPackages: []string{\n\t\t\t\t\t\"torch==\" + compat.TorchVersion(),\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\terr := config.Complete(\"\")\n\t\trequire.NoError(t, err)\n\t\tif compat.CUDA == nil {\n\t\t\trequire.Equal(t, \"\", config.Build.CUDA)\n\t\t\trequire.Equal(t, \"\", config.Build.CuDNN)\n\t\t} else {\n\t\t\trequire.NotEqual(t, \"\", config.Build.CUDA)\n\t\t\trequire.NotEqual(t, \"\", config.Build.CuDNN)\n\t\t}\n\t}\n}\n\nfunc TestValidateAndCompleteCUDAForSelectedTorch(t *testing.T) {\n\tfor _, tt := range []struct {\n\t\ttorch string\n\t\tcuda  string\n\t\tcuDNN string\n\t}{\n\t\t{\"2.0.1\", \"11.8\", \"8\"},\n\t\t{\"1.13.1\", \"11.7\", \"8\"},\n\t\t{\"1.11.0\", \"11.3\", \"8\"},\n\t} {\n\t\tconfig := &Config{\n\t\t\tBuild: &Build{\n\t\t\t\tGPU:           true,\n\t\t\t\tPythonVersion: \"3.10\",\n\t\t\t\tPythonPackages: []string{\n\t\t\t\t\t\"torch==\" + tt.torch,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\terr := config.Complete(\"\")\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, tt.cuda, config.Build.CUDA)\n\t\trequire.Equal(t, tt.cuDNN, config.Build.CuDNN)\n\t}\n}\n\nfunc TestUnsupportedTorch(t *testing.T) {\n\t// Ensure version is not known by Cog\n\tcudas, err := cudasFromTorch(\"0.4.1\")\n\trequire.NoError(t, err)\n\trequire.Empty(t, cudas)\n\n\t// Unknown versions require cuda\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==0.4.1\",\n\t\t\t},\n\t\t},\n\t}\n\terr = config.Complete(\"\")\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"cog doesn't know what CUDA version is compatible with torch==0.4.1.\")\n\n\tconfig = &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tCUDA:          \"11.8\",\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==0.4.1\",\n\t\t\t},\n\t\t},\n\t}\n\terr = config.Complete(\"\")\n\trequire.NoError(t, err)\n\tassertMinorVersion(t, \"11.8\", config.Build.CUDA)\n\trequire.Equal(t, \"8\", config.Build.CuDNN)\n}\n\nfunc TestUnsupportedTensorflow(t *testing.T) {\n\t// Ensure version is not known by Cog\n\tcuda, cudnn, err := cudaFromTF(\"0.4.1\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, cuda, \"\")\n\trequire.Equal(t, cudnn, \"\")\n\n\t// Unknown versions require cuda\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"tensorflow==0.4.1\",\n\t\t\t},\n\t\t},\n\t}\n\terr = config.Complete(\"\")\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"cog doesn't know what CUDA version is compatible with tensorflow==0.4.1.\")\n\n\tconfig = &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tCUDA:          \"11.8\",\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"tensorflow==0.4.1\",\n\t\t\t},\n\t\t},\n\t}\n\terr = config.Complete(\"\")\n\trequire.NoError(t, err)\n\tassertMinorVersion(t, \"11.8\", config.Build.CUDA)\n\trequire.Equal(t, \"8\", config.Build.CuDNN)\n}\n\nfunc TestPythonPackagesForArchTorchGPU(t *testing.T) {\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==2.0.1\",\n\t\t\t\t\"torchvision==0.15.2\",\n\t\t\t\t\"torchaudio==2.0.2\",\n\t\t\t\t\"foo==1.0.0\",\n\t\t\t},\n\t\t\tCUDA: \"11.8\",\n\t\t},\n\t}\n\terr := config.Complete(\"\")\n\trequire.NoError(t, err)\n\tassertMinorVersion(t, \"11.8\", config.Build.CUDA)\n\trequire.Equal(t, \"8\", config.Build.CuDNN)\n\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{})\n\trequire.NoError(t, err)\n\texpected := `--extra-index-url https://download.pytorch.org/whl/cu118\ntorch==2.0.1\ntorchvision==0.15.2\ntorchaudio==2.0.2\nfoo==1.0.0`\n\trequire.Equal(t, expected, requirements)\n}\n\nfunc TestPythonPackagesForArchTorchCPU(t *testing.T) {\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           false,\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==2.0.1\",\n\t\t\t\t\"torchvision==0.15.2\",\n\t\t\t\t\"torchaudio==2.0.2\",\n\t\t\t\t\"foo==1.0.0\",\n\t\t\t},\n\t\t\tCUDA: \"11.8\",\n\t\t},\n\t}\n\terr := config.Complete(\"\")\n\trequire.NoError(t, err)\n\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{})\n\trequire.NoError(t, err)\n\texpected := `--extra-index-url https://download.pytorch.org/whl/cpu\ntorch==2.0.1\ntorchvision==0.15.2\ntorchaudio==2.0.2\nfoo==1.0.0`\n\trequire.Equal(t, expected, requirements)\n}\n\nfunc TestPythonPackagesForArchTensorflowGPU(t *testing.T) {\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"tensorflow==2.12.0\",\n\t\t\t\t\"foo==1.0.0\",\n\t\t\t},\n\t\t\tCUDA: \"11.8\",\n\t\t},\n\t}\n\terr := config.Complete(\"\")\n\trequire.NoError(t, err)\n\tassertMinorVersion(t, \"11.8\", config.Build.CUDA)\n\trequire.Equal(t, \"8\", config.Build.CuDNN)\n\n\t// tensorflow and tensorflow-gpu have been the same package since TensorFlow 2.1, released in September 2019.\n\t// Although the checksums differ due to metadata,\n\t// they were built in the same way and both provide GPU support via Nvidia CUDA.\n\t// As of December 2022, tensorflow-gpu has been removed and has been replaced with\n\t// this new, empty package that generates an error upon installation.\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{})\n\trequire.NoError(t, err)\n\texpected := `tensorflow==2.12.0\nfoo==1.0.0`\n\trequire.Equal(t, expected, requirements)\n\trequire.NotContains(t, requirements, \"tensorflow_gpu\")\n}\n\nfunc TestPythonPackagesBothTorchAndTensorflow(t *testing.T) {\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.11\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"tensorflow==2.16.1\",\n\t\t\t\t\"torch==2.3.1\",\n\t\t\t},\n\t\t\tCUDA: \"12.3\",\n\t\t},\n\t}\n\terr := config.Complete(\"\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"12.3\", config.Build.CUDA)\n\trequire.Equal(t, \"8\", config.Build.CuDNN)\n\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{})\n\trequire.NoError(t, err)\n\texpected := `--extra-index-url https://download.pytorch.org/whl/cu121\ntensorflow==2.16.1\ntorch==2.3.1`\n\trequire.Equal(t, expected, requirements)\n}\n\nfunc TestCUDABaseImageTag(t *testing.T) {\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"tensorflow==2.12.0\",\n\t\t\t},\n\t\t},\n\t}\n\n\terr := config.Complete(\"\")\n\trequire.NoError(t, err)\n\n\timageTag, err := config.CUDABaseImageTag()\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04\", imageTag)\n}\n\nfunc TestBuildRunItemStringYAML(t *testing.T) {\n\ttype BuildWrapper struct {\n\t\tBuild *Build `yaml:\"build\"`\n\t}\n\n\tvar buildWrapper BuildWrapper\n\n\tyamlString := `\nbuild:\n  run:\n    - \"echo 'Hello, World!'\"\n`\n\n\terr := yaml.Unmarshal([]byte(yamlString), &buildWrapper)\n\trequire.NoError(t, err)\n\trequire.NotNil(t, buildWrapper.Build)\n\trequire.Len(t, buildWrapper.Build.Run, 1)\n\trequire.Equal(t, \"echo 'Hello, World!'\", buildWrapper.Build.Run[0].Command)\n}\n\nfunc TestBuildRunItemStringJSON(t *testing.T) {\n\ttype BuildWrapper struct {\n\t\tBuild *Build `json:\"build\"`\n\t}\n\n\tvar buildWrapper BuildWrapper\n\n\tjsonString := `{\n\t\"build\": {\n\t\t\"run\": [\n\t\t\t\"echo 'Hello, World!'\"\n\t\t]\n\t}\n}`\n\n\terr := json.Unmarshal([]byte(jsonString), &buildWrapper)\n\trequire.NoError(t, err)\n\trequire.NotNil(t, buildWrapper.Build)\n\trequire.Len(t, buildWrapper.Build.Run, 1)\n\trequire.Equal(t, \"echo 'Hello, World!'\", buildWrapper.Build.Run[0].Command)\n}\n\nfunc TestBuildRunItemDictYAML(t *testing.T) {\n\ttype BuildWrapper struct {\n\t\tBuild *Build `yaml:\"build\"`\n\t}\n\n\tvar buildWrapper BuildWrapper\n\n\tyamlString := `\nbuild:\n  run:\n  - command: \"echo 'Hello, World!'\"\n    mounts:\n    - type: bind\n      id: my-volume\n      target: /mnt/data\n`\n\n\terr := yaml.Unmarshal([]byte(yamlString), &buildWrapper)\n\trequire.NoError(t, err)\n\trequire.NotNil(t, buildWrapper.Build)\n\trequire.Len(t, buildWrapper.Build.Run, 1)\n\trequire.Equal(t, \"echo 'Hello, World!'\", buildWrapper.Build.Run[0].Command)\n\trequire.Len(t, buildWrapper.Build.Run[0].Mounts, 1)\n\trequire.Equal(t, \"bind\", buildWrapper.Build.Run[0].Mounts[0].Type)\n\trequire.Equal(t, \"my-volume\", buildWrapper.Build.Run[0].Mounts[0].ID)\n\trequire.Equal(t, \"/mnt/data\", buildWrapper.Build.Run[0].Mounts[0].Target)\n}\n\nfunc TestBuildRunItemDictJSON(t *testing.T) {\n\ttype BuildWrapper struct {\n\t\tBuild *Build `json:\"build\"`\n\t}\n\n\tvar buildWrapper BuildWrapper\n\n\tjsonString := `{\n\t\"build\": {\n\t\t\"run\": [\n\t\t\t{\n\t\t\t\t\"command\": \"echo 'Hello, World!'\",\n\t\t\t\t\"mounts\": [\n\t\t\t\t\t{\n\t\t\t\t\t\t\"type\": \"bind\",\n\t\t\t\t\t\t\"id\": \"my-volume\",\n\t\t\t\t\t\t\"target\": \"/mnt/data\"\n\t\t\t\t\t}\n\t\t\t\t]\n\t\t\t}\n\t\t]\n\t}\n}`\n\n\terr := json.Unmarshal([]byte(jsonString), &buildWrapper)\n\trequire.NoError(t, err)\n\trequire.NotNil(t, buildWrapper.Build)\n\trequire.Len(t, buildWrapper.Build.Run, 1)\n\trequire.Equal(t, \"echo 'Hello, World!'\", buildWrapper.Build.Run[0].Command)\n\trequire.Len(t, buildWrapper.Build.Run[0].Mounts, 1)\n\trequire.Equal(t, \"bind\", buildWrapper.Build.Run[0].Mounts[0].Type)\n\trequire.Equal(t, \"my-volume\", buildWrapper.Build.Run[0].Mounts[0].ID)\n\trequire.Equal(t, \"/mnt/data\", buildWrapper.Build.Run[0].Mounts[0].Target)\n}\n\nfunc TestTorchWithExistingExtraIndexURL(t *testing.T) {\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==1.12.1 --extra-index-url=https://download.pytorch.org/whl/cu116\",\n\t\t\t},\n\t\t\tCUDA: \"11.6.2\",\n\t\t},\n\t}\n\terr := config.Complete(\"\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"11.6.2\", config.Build.CUDA)\n\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{})\n\trequire.NoError(t, err)\n\texpected := `--extra-index-url https://download.pytorch.org/whl/cu116\ntorch==1.12.1`\n\trequire.Equal(t, expected, requirements)\n}\n\nfunc TestBlankBuild(t *testing.T) {\n\t// Naively, this turns into nil, so make sure it's a real build object\n\t// Write a temp file\n\tdir := t.TempDir()\n\tconfigPath := path.Join(dir, \"cog.yaml\")\n\terr := os.WriteFile(configPath, []byte(`build:`), 0o644)\n\trequire.NoError(t, err)\n\n\tcfgFile, err := parseFile(configPath)\n\trequire.NoError(t, err)\n\t// Note: `build:` by itself in YAML parses to Build: nil (empty map becomes nil pointer)\n\t// The completion step should create a default Build\n\n\tconfig, err := configFileToConfig(cfgFile)\n\trequire.NoError(t, err)\n\trequire.NoError(t, config.Complete(dir))\n\trequire.NotNil(t, config.Build)\n\trequire.Equal(t, false, config.Build.GPU)\n}\n\nfunc TestPythonRequirementsForArchWithAddedPackage(t *testing.T) {\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.10\",\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==2.4.0 --extra-index-url=https://download.pytorch.org/whl/cu116\",\n\t\t\t},\n\t\t\tCUDA: \"11.6.2\",\n\t\t},\n\t}\n\terr := config.Complete(\"\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"11.6.2\", config.Build.CUDA)\n\trequirements, err := config.PythonRequirementsForArch(\"\", \"\", []string{\n\t\t\"torchvision==2.4.0\",\n\t})\n\trequire.NoError(t, err)\n\texpected := `--extra-index-url https://download.pytorch.org/whl/cu116\ntorch==2.4.0\ntorchvision==2.4.0`\n\trequire.Equal(t, expected, requirements)\n}\n\nfunc TestParseTests(t *testing.T) {\n\tyamlString := `\nbuild:\n  run:\n  - command: \"echo 'Hello, World!'\"\n`\n\tdir := t.TempDir()\n\tconfigPath := path.Join(dir, \"cog.yaml\")\n\terr := os.WriteFile(configPath, []byte(yamlString), 0o644)\n\trequire.NoError(t, err)\n\n\t_, err = parseFile(configPath)\n\trequire.NoError(t, err)\n}\n\nfunc TestConfigMarshal(t *testing.T) {\n\tcfg := defaultConfig()\n\tdata, err := yaml.Marshal(cfg)\n\trequire.NoError(t, err)\n\t// yaml v4 uses 4-space indentation by default\n\trequire.Equal(t, `build:\n    python_version: \"3.13\"\npredict: \"\"\n`, string(data))\n}\n\nfunc TestAbsolutePathInPythonRequirements(t *testing.T) {\n\tdir := t.TempDir()\n\trequirementsFilePath := filepath.Join(dir, \"requirements.txt\")\n\terr := os.WriteFile(requirementsFilePath, []byte(\"torch==2.5.0\"), 0o644)\n\trequire.NoError(t, err)\n\tconfig := &Config{\n\t\tBuild: &Build{\n\t\t\tGPU:                true,\n\t\t\tPythonVersion:      \"3.10\",\n\t\t\tPythonRequirements: requirementsFilePath,\n\t\t},\n\t}\n\terr = config.Complete(dir)\n\trequire.NoError(t, err)\n\ttorchVersion, ok := config.TorchVersion()\n\trequire.Equal(t, torchVersion, \"2.5.0\")\n\trequire.True(t, ok)\n}\n\nfunc TestWeightsWithNameYAML(t *testing.T) {\n\tyamlString := `build:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\nweights:\n  - name: model-v1\n    source: file://./weights/model-v1.zip\n    target: \"/weights/model-v1\"\n  - name: model-v2\n    source: file://./weights/model-v2.zip\n    target: \"/weights/model-v2\"\n`\n\n\tconfig, err := FromYAML([]byte(yamlString))\n\trequire.NoError(t, err)\n\trequire.Len(t, config.Weights, 2)\n\n\trequire.Equal(t, \"model-v1\", config.Weights[0].Name)\n\trequire.Equal(t, \"file://./weights/model-v1.zip\", config.Weights[0].Source)\n\trequire.Equal(t, \"/weights/model-v1\", config.Weights[0].Target)\n\n\trequire.Equal(t, \"model-v2\", config.Weights[1].Name)\n\trequire.Equal(t, \"file://./weights/model-v2.zip\", config.Weights[1].Source)\n\trequire.Equal(t, \"/weights/model-v2\", config.Weights[1].Target)\n}\n\nfunc TestWeightsWithoutNameYAML(t *testing.T) {\n\tyamlString := `build:\n  python_version: \"3.12\"\npredict: \"predict.py:Predictor\"\n\nweights:\n  - source: file://./weights/model.zip\n    target: \"/weights/model\"\n`\n\n\tconfig, err := FromYAML([]byte(yamlString))\n\trequire.NoError(t, err)\n\trequire.Len(t, config.Weights, 1)\n\n\trequire.Equal(t, \"\", config.Weights[0].Name)\n\trequire.Equal(t, \"file://./weights/model.zip\", config.Weights[0].Source)\n\trequire.Equal(t, \"/weights/model\", config.Weights[0].Target)\n}\n\nfunc TestWeightsWithNameJSON(t *testing.T) {\n\tjsonString := `{\n\t\"build\": {\n\t\t\"python_version\": \"3.12\"\n\t},\n\t\"predict\": \"predict.py:Predictor\",\n\t\"weights\": [\n\t\t{\n\t\t\t\"name\": \"model-v1\",\n\t\t\t\"source\": \"file://./weights/model-v1.zip\",\n\t\t\t\"target\": \"/weights/model-v1\"\n\t\t},\n\t\t{\n\t\t\t\"name\": \"model-v2\",\n\t\t\t\"source\": \"file://./weights/model-v2.zip\",\n\t\t\t\"target\": \"/weights/model-v2\"\n\t\t}\n\t]\n}`\n\n\tvar config Config\n\terr := json.Unmarshal([]byte(jsonString), &config)\n\trequire.NoError(t, err)\n\trequire.Len(t, config.Weights, 2)\n\n\trequire.Equal(t, \"model-v1\", config.Weights[0].Name)\n\trequire.Equal(t, \"file://./weights/model-v1.zip\", config.Weights[0].Source)\n\trequire.Equal(t, \"/weights/model-v1\", config.Weights[0].Target)\n\n\trequire.Equal(t, \"model-v2\", config.Weights[1].Name)\n\trequire.Equal(t, \"file://./weights/model-v2.zip\", config.Weights[1].Source)\n\trequire.Equal(t, \"/weights/model-v2\", config.Weights[1].Target)\n}\n\nfunc TestSDKVersionConfig(t *testing.T) {\n\t// build.sdk_version is parsed and stored correctly\n\tconf, err := FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  sdk_version: \"0.18.0\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"0.18.0\", conf.Build.SDKVersion)\n}\n\nfunc TestSDKVersionConfigEmpty(t *testing.T) {\n\t// Omitting build.sdk_version leaves the field empty\n\tconf, err := FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"\", conf.Build.SDKVersion)\n}\n\nfunc TestSDKVersionConfigPreRelease(t *testing.T) {\n\t// Pre-release PEP 440 version is accepted and stored verbatim\n\tconf, err := FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  sdk_version: \"0.18.0a1\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"0.18.0a1\", conf.Build.SDKVersion)\n}\n\nfunc TestSDKVersionConfigBelowMinimumExplodesInGenerator(t *testing.T) {\n\t// build.sdk_version < 0.16.0 must be rejected — parsing succeeds but the\n\t// Dockerfile generator must return an error so the build never proceeds.\n\tconf, err := FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  sdk_version: \"0.15.0\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\t// Parsing itself is fine; enforcement happens at Dockerfile generation time.\n\trequire.Equal(t, \"0.15.0\", conf.Build.SDKVersion)\n}\n"
  },
  {
    "path": "pkg/config/cuda_compatibility.json",
    "content": "[\n  {\n    \"Tag\": \"11.0.3-cudnn8-devel-ubuntu16.04\",\n    \"CUDA\": \"11.0.3\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"16.04\"\n  },\n  {\n    \"Tag\": \"11.0.3-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.0.3\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.0.3-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.0.3\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.1.1-cudnn8-devel-ubuntu16.04\",\n    \"CUDA\": \"11.1.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"16.04\"\n  },\n  {\n    \"Tag\": \"11.1.1-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.1.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.1.1-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.1.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.2.2-cudnn8-devel-ubuntu16.04\",\n    \"CUDA\": \"11.2.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"16.04\"\n  },\n  {\n    \"Tag\": \"11.2.2-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.2.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.2.2-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.2.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.3.1-cudnn8-devel-ubuntu16.04\",\n    \"CUDA\": \"11.3.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"16.04\"\n  },\n  {\n    \"Tag\": \"11.3.1-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.3.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.3.1-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.3.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.4.3-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.4.3\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.4.3-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.4.3\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.5.2-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.5.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.5.2-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.5.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.6.1-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.6.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.6.2-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.6.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.6.2-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.6.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.7.1-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.7.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.7.1-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.7.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.7.1-cudnn8-devel-ubuntu22.04\",\n    \"CUDA\": \"11.7.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"11.8.0-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"11.8.0\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"11.8.0-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"11.8.0\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"11.8.0-cudnn8-devel-ubuntu22.04\",\n    \"CUDA\": \"11.8.0\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.0.0-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"12.0.0\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"12.0.0-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"12.0.0\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.0.0-cudnn8-devel-ubuntu22.04\",\n    \"CUDA\": \"12.0.0\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.0.1-cudnn8-devel-ubuntu18.04\",\n    \"CUDA\": \"12.0.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"18.04\"\n  },\n  {\n    \"Tag\": \"12.0.1-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"12.0.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.0.1-cudnn8-devel-ubuntu22.04\",\n    \"CUDA\": \"12.0.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.1.0-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"12.1.0\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.1.0-cudnn8-devel-ubuntu22.04\",\n    \"CUDA\": \"12.1.0\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.1.1-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"12.1.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.1.1-cudnn8-devel-ubuntu22.04\",\n    \"CUDA\": \"12.1.1\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.2.2-cudnn8-devel-ubuntu20.04\",\n    \"CUDA\": \"12.2.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.2.2-cudnn8-devel-ubuntu22.04\",\n    \"CUDA\": \"12.2.2\",\n    \"CuDNN\": \"8\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.3.2-cudnn9-devel-ubuntu20.04\",\n    \"CUDA\": \"12.3.2\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.3.2-cudnn9-devel-ubuntu22.04\",\n    \"CUDA\": \"12.3.2\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.4.1-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.4.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.4.1-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.4.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.5.1-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.5.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.5.1-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.5.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.6.0-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.6.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.6.0-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.6.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.6.0-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"12.6.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"12.6.1-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.6.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.6.1-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.6.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.6.1-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"12.6.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"12.6.2-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.6.2\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.6.2-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.6.2\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.6.2-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"12.6.2\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"12.6.3-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.6.3\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.6.3-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.6.3\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.6.3-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"12.6.3\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"12.8.0-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.8.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.8.0-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.8.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.8.0-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"12.8.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"12.8.1-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.8.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.8.1-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.8.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.8.1-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"12.8.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"12.9.0-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.9.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.9.0-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.9.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.9.0-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"12.9.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"12.9.1-cudnn-devel-ubuntu20.04\",\n    \"CUDA\": \"12.9.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"20.04\"\n  },\n  {\n    \"Tag\": \"12.9.1-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"12.9.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"12.9.1-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"12.9.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"13.0.0-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"13.0.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"13.0.0-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"13.0.0\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"13.0.1-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"13.0.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"13.0.1-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"13.0.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"13.0.2-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"13.0.2\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"13.0.2-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"13.0.2\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  },\n  {\n    \"Tag\": \"13.1.1-cudnn-devel-ubuntu22.04\",\n    \"CUDA\": \"13.1.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"22.04\"\n  },\n  {\n    \"Tag\": \"13.1.1-cudnn-devel-ubuntu24.04\",\n    \"CUDA\": \"13.1.1\",\n    \"CuDNN\": \"9\",\n    \"IsDevel\": true,\n    \"Ubuntu\": \"24.04\"\n  }\n]"
  },
  {
    "path": "pkg/config/data/config_schema_v1.0.json",
    "content": "{\n  \"$schema\": \"http://json-schema.org/draft-07/schema\",\n  \"type\": \"object\",\n  \"title\": \"Schema for cog.yaml\",\n  \"description\": \"Defines how to build a Docker image and how to run predictions on your model inside that image.\",\n  \"properties\": {\n    \"build\": {\n      \"$id\": \"#/properties/build\",\n      \"type\": \"object\",\n      \"description\": \"This stanza describes how to build the Docker image your model runs in.\",\n      \"properties\": {\n        \"cuda\": {\n          \"$id\": \"#/properties/build/properties/cuda\",\n          \"type\": \"string\",\n          \"description\": \"Cog automatically picks the correct version of CUDA to install, but this lets you override it for whatever reason.\"\n        },\n        \"cudnn\": {\n          \"$id\": \"#/properties/build/properties/cudnn\",\n          \"type\": \"string\",\n          \"description\": \"Cog automatically picks the correct version of cuDNN to install, but this lets you override it for whatever reason.\"\n        },\n        \"gpu\": {\n          \"$id\": \"#/properties/build/properties/gpu\",\n          \"type\": \"boolean\",\n          \"description\": \"Enable GPUs for this model. When enabled, the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image will be used, and Cog will automatically figure out what versions of CUDA and cuDNN to use based on the version of Python, PyTorch, and Tensorflow that you are using.\"\n        },\n        \"python_version\": {\n          \"$id\": \"#/properties/build/properties/python_version\",\n          \"type\": [\n            \"string\",\n            \"number\"\n          ],\n          \"description\": \"The minor (`3.13`) or patch (`3.13.1`) version of Python to use.\"\n        },\n        \"python_packages\": {\n          \"$id\": \"#/properties/build/properties/python_packages\",\n          \"type\": [\n            \"array\",\n            \"null\"\n          ],\n          \"description\": \"A list of Python packages to install, in the format `package==version`.\",\n          \"additionalItems\": true,\n          \"items\": {\n            \"$id\": \"#/properties/build/properties/python_packages/items\",\n            \"anyOf\": [\n              {\n                \"$id\": \"#/properties/build/properties/python_packages/items/anyOf/0\",\n                \"type\": \"string\"\n              }\n            ]\n          }\n        },\n        \"pre_install\": {\n          \"$id\": \"#/properties/build/properties/pre_install\",\n          \"type\": [\n            \"array\",\n            \"null\"\n          ],\n          \"description\": \"A list of setup commands to run in the environment before your Python packages are installed.\",\n          \"additionalItems\": true,\n          \"items\": {\n            \"$id\": \"#/properties/build/properties/pre_install/items\",\n            \"anyOf\": [\n              {\n                \"$id\": \"#/properties/build/properties/pre_install/items/anyOf/0\",\n                \"type\": \"string\"\n              }\n            ]\n          }\n        },\n        \"python_requirements\": {\n          \"$id\": \"#/properties/build/properties/python_requirements\",\n          \"type\": \"string\",\n          \"description\": \"A pip requirements file specifying the Python packages to install.\"\n        },\n        \"system_packages\": {\n          \"$id\": \"#/properties/build/properties/system_packages\",\n          \"type\": [\n            \"array\",\n            \"null\"\n          ],\n          \"description\": \"A list of Ubuntu APT packages to install.\",\n          \"additionalItems\": true,\n          \"items\": {\n            \"$id\": \"#/properties/build/properties/system_packages/items\",\n            \"anyOf\": [\n              {\n                \"$id\": \"#/properties/build/properties/system_packages/items/anyOf/0\",\n                \"type\": \"string\"\n              }\n            ]\n          }\n        },\n        \"sdk_version\": {\n          \"$id\": \"#/properties/build/properties/sdk_version\",\n          \"type\": \"string\",\n          \"description\": \"Pin the cog Python SDK version installed in the container (e.g. \\\"0.18.0\\\" or \\\"0.18.0a1\\\"). Defaults to latest. Overridden by the COG_SDK_WHEEL environment variable.\"\n        },\n        \"run\": {\n          \"$id\": \"#/properties/build/properties/run\",\n          \"type\": [\n            \"array\",\n            \"null\"\n          ],\n          \"description\": \"A list of setup commands to run in the environment after your system packages and Python packages have been installed. If you're familiar with Docker, it's like a `RUN` instruction in your `Dockerfile`.\",\n          \"additionalItems\": true,\n          \"items\": {\n            \"$id\": \"#/properties/build/properties/run/items\",\n            \"anyOf\": [\n              {\n                \"$id\": \"#/properties/build/properties/run/items/anyOf/0\",\n                \"type\": \"string\"\n              },\n              {\n                \"$id\": \"#/properties/build/properties/run/items/anyOf/1\",\n                \"type\": \"object\",\n                \"properties\": {\n                  \"command\": {\n                    \"type\": \"string\"\n                  },\n                  \"mounts\": {\n                    \"type\": \"array\",\n                    \"items\": {\n                      \"type\": \"object\",\n                      \"properties\": {\n                        \"type\": {\n                          \"type\": \"string\",\n                          \"enum\": [\n                            \"secret\"\n                          ]\n                        },\n                        \"id\": {\n                          \"type\": \"string\"\n                        },\n                        \"target\": {\n                          \"type\": \"string\"\n                        }\n                      },\n                      \"required\": [\n                        \"type\",\n                        \"id\",\n                        \"target\"\n                      ]\n                    }\n                  }\n                },\n                \"required\": [\n                  \"command\"\n                ]\n              }\n            ]\n          }\n        }\n      },\n      \"required\": [\"python_version\"],\n      \"additionalProperties\": false\n    },\n    \"image\": {\n      \"$id\": \"#/properties/image\",\n      \"type\": \"string\",\n      \"description\": \"The name given to built Docker images. If you want to push to a registry, this should also include the registry name.\"\n    },\n    \"predict\": {\n      \"$id\": \"#/properties/predict\",\n      \"type\": \"string\",\n      \"description\": \"The pointer to the `Predictor` object in your code, which defines how predictions are run on your model.\"\n    },\n    \"train\": {\n      \"$id\": \"#/properties/train\",\n      \"type\": \"string\",\n      \"description\": \"The pointer to the `Predictor` object in your code, which defines how predictions are run on your model.\"\n    },\n    \"concurrency\": {\n      \"$id\": \"#/properties/concurrency\",\n      \"type\": \"object\",\n      \"description\": \"The concurrency settings for the model.\",\n      \"required\": [\n        \"max\"\n      ],\n      \"additionalProperties\": false,\n      \"properties\": {\n        \"max\": {\n          \"$id\": \"#/properties/concurrency/properties/max\",\n          \"type\": \"integer\",\n          \"description\": \"The maximum number of concurrent predictions.\"\n        },\n        \"default_target\": {\n          \"$id\": \"#/properties/concurrency/properties/default_target\",\n          \"type\": \"integer\",\n          \"description\": \"The default target for number of concurrent predictions. This setting can be used by an autoscaler to determine when to scale a deployment of a model up or down.\"\n        }\n      }\n    },\n    \"environment\": {\n      \"$id\": \"#/properties/properties/environment\",\n      \"type\": [\n        \"array\",\n        \"null\"\n      ],\n      \"description\": \"A list of environment variables to make available during builds and at runtime, in the format `NAME=value`\",\n      \"additionalItems\": true,\n      \"items\": {\n        \"$id\": \"#/properties/properties/environment/items\",\n        \"type\": \"string\",\n        \"pattern\": \"^[A-Za-z_][A-Za-z0-9_]*=.*$\"\n      }\n    },\n    \"weights\": {\n      \"$id\": \"#/properties/weights\",\n      \"type\": [\n        \"array\",\n        \"null\"\n      ],\n      \"description\": \"A list of weight files or directories to include in the model.\",\n      \"items\": {\n        \"type\": \"object\",\n        \"required\": [\"source\"],\n        \"additionalProperties\": false,\n        \"properties\": {\n          \"name\": {\n            \"type\": \"string\",\n            \"description\": \"A unique identifier for this weight entry.\"\n          },\n          \"source\": {\n            \"type\": \"string\",\n            \"description\": \"Path to a weight file or directory (relative to cog.yaml).\"\n          },\n          \"target\": {\n            \"type\": \"string\",\n            \"description\": \"Target path in the container (must be under /cache/). Defaults to /cache/.\"\n          }\n        }\n      }\n    }\n  },\n  \"additionalProperties\": false\n}\n"
  },
  {
    "path": "pkg/config/env.go",
    "content": "package config\n\nimport (\n\t\"fmt\"\n\t\"strings\"\n)\n\n// environmentVariableDenyList is a list of environment variable patterns that are\n// used internally during build or runtime and thus not allowed to be set by the user.\n// There are ways around this restriction, but it's likely to cause unexpected behavior\n// and hard to debug issues. So on Cog's predict-build-push happy path, we don't allow\n// these to be set.\n// This list may change at any time. For more context, see:\n// https://github.com/replicate/cog/pull/2274/#issuecomment-2831823185\nvar environmentVariableDenyList = []string{\n\t// paths\n\t\"PATH\",\n\t\"LD_LIBRARY_PATH\",\n\t\"PYTHONPATH\",\n\t\"VIRTUAL_ENV\",\n\t\"PYTHONUNBUFFERED\",\n\t// Replicate\n\t\"R8_*\",\n\t\"REPLICATE_*\",\n\t// Nvidia\n\t\"LIBRARY_PATH\",\n\t\"CUDA_*\",\n\t\"NVIDIA_*\",\n\t\"NV_*\",\n\t// pget\n\t\"PGET_*\",\n\t\"HF_ENDPOINT\",\n\t\"HF_HUB_ENABLE_HF_TRANSFER\",\n\t// k8s\n\t\"KUBERNETES_*\",\n}\n\n// validateEnvName checks if the given environment variable name is allowed.\n// Returns an error if the name matches any of the restricted patterns.\nfunc validateEnvName(name string) error {\n\tfor _, pattern := range environmentVariableDenyList {\n\t\t// Check for exact match\n\t\tif pattern == name {\n\t\t\treturn fmt.Errorf(\"environment variable %q is not allowed\", name)\n\t\t}\n\n\t\t// Check for wildcard pattern\n\t\tif strings.HasSuffix(pattern, \"*\") {\n\t\t\tif strings.HasPrefix(name, pattern[:len(pattern)-1]) {\n\t\t\t\treturn fmt.Errorf(\"environment variable %q is not allowed\", name)\n\t\t\t}\n\t\t}\n\t}\n\treturn nil\n}\n\n// parseAndValidateEnvironment converts a slice of strings in the format of KEY=VALUE\n// to a map[string]string. An error is returned if the format is incorrect or if either\n// the variable name or value are invalid.\nfunc parseAndValidateEnvironment(input []string) (map[string]string, error) {\n\tenv := map[string]string{}\n\tfor _, input := range input {\n\t\tparts := strings.SplitN(input, \"=\", 2)\n\t\tif len(parts) != 2 || parts[0] == \"\" {\n\t\t\treturn nil, fmt.Errorf(\"environment variable %q is not in the KEY=VALUE format\", input)\n\t\t}\n\t\tif err := validateEnvName(parts[0]); err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tif _, ok := env[parts[0]]; ok {\n\t\t\treturn nil, fmt.Errorf(\"environment variable %q is already defined\", parts[0])\n\t\t}\n\t\tenv[parts[0]] = parts[1]\n\t}\n\treturn env, nil\n}\n"
  },
  {
    "path": "pkg/config/env_variables_test.go",
    "content": "package config\n\nimport (\n\t\"fmt\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestEnvironmentConfig(t *testing.T) {\n\tt.Run(\"ParsingValidInput\", func(t *testing.T) {\n\t\tcases := []struct {\n\t\t\tName     string\n\t\t\tInput    []string\n\t\t\tExpected map[string]string\n\t\t}{\n\t\t\t{\n\t\t\t\tName:     \"ValidInput\",\n\t\t\t\tInput:    []string{\"NAME=VALUE\"},\n\t\t\t\tExpected: map[string]string{\"NAME\": \"VALUE\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"ValidInputWithSpaces\",\n\t\t\t\tInput:    []string{\"NAME=VALUE WITH SPACES\"},\n\t\t\t\tExpected: map[string]string{\"NAME\": \"VALUE WITH SPACES\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"ValidInputWithQuotes\",\n\t\t\t\tInput:    []string{\"NAME=\\\"VALUE WITH QUOTES\\\"\"},\n\t\t\t\tExpected: map[string]string{\"NAME\": `\"VALUE WITH QUOTES\"`},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"DelimitedValue\",\n\t\t\t\tInput:    []string{\"NAME=VALUE1,VALUE2\"},\n\t\t\t\tExpected: map[string]string{\"NAME\": \"VALUE1,VALUE2\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"EmptyValue\",\n\t\t\t\tInput:    []string{\"NAME=\"},\n\t\t\t\tExpected: map[string]string{\"NAME\": \"\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"EmptyValueWithSpaces\",\n\t\t\t\tInput:    []string{\"NAME= \"},\n\t\t\t\tExpected: map[string]string{\"NAME\": \" \"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"LowerCaseName\",\n\t\t\t\tInput:    []string{\"name=VALUE\"},\n\t\t\t\tExpected: map[string]string{\"name\": \"VALUE\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"MixedCaseName\",\n\t\t\t\tInput:    []string{\"MiXeD_Case=VALUE\"},\n\t\t\t\tExpected: map[string]string{\"MiXeD_Case\": \"VALUE\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"EqualSignInValue\",\n\t\t\t\tInput:    []string{\"NAME=VALUE=EQUAL\"},\n\t\t\t\tExpected: map[string]string{\"NAME\": \"VALUE=EQUAL\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"EqualSignInValueWithSpaces\",\n\t\t\t\tInput:    []string{\"NAME=VALUE=EQUAL WITH SPACES\"},\n\t\t\t\tExpected: map[string]string{\"NAME\": \"VALUE=EQUAL WITH SPACES\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"MultiLineValue\",\n\t\t\t\tInput:    []string{\"NAME=VALUE1\\nVALUE2\"},\n\t\t\t\tExpected: map[string]string{\"NAME\": \"VALUE1\\nVALUE2\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"UserAgentWithSpaces\",\n\t\t\t\tInput:    []string{\"COG_USER_AGENT=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36\"},\n\t\t\t\tExpected: map[string]string{\"COG_USER_AGENT\": \"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36\"},\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:     \"MultiplePairs\",\n\t\t\t\tInput:    []string{\"NAME1=VALUE1\", \"NAME2=VALUE2\"},\n\t\t\t\tExpected: map[string]string{\"NAME1\": \"VALUE1\", \"NAME2\": \"VALUE2\"},\n\t\t\t},\n\t\t}\n\n\t\tfor _, c := range cases {\n\t\t\tt.Run(c.Name, func(t *testing.T) {\n\t\t\t\tparsed, err := parseAndValidateEnvironment(c.Input)\n\t\t\t\trequire.NoError(t, err)\n\t\t\t\trequire.Equal(t, c.Expected, parsed)\n\t\t\t})\n\t\t}\n\t})\n\n\tt.Run(\"ParsingInvalidInput\", func(t *testing.T) {\n\t\tcases := []struct {\n\t\t\tName                 string\n\t\t\tInput                []string\n\t\t\tExpectedErrorMessage string\n\t\t}{\n\t\t\t{\n\t\t\t\tName:                 \"NameWithoutValue\",\n\t\t\t\tInput:                []string{\"NAME\"},\n\t\t\t\tExpectedErrorMessage: `environment variable \"NAME\" is not in the KEY=VALUE format`,\n\t\t\t},\n\t\t\t{\n\t\t\t\tName:                 \"EmptyName\",\n\t\t\t\tInput:                []string{\"=VALUE\"},\n\t\t\t\tExpectedErrorMessage: `environment variable \"=VALUE\" is not in the KEY=VALUE format`,\n\t\t\t},\n\t\t}\n\n\t\tfor _, c := range cases {\n\t\t\tt.Run(c.Name, func(t *testing.T) {\n\t\t\t\t_, err := parseAndValidateEnvironment(c.Input)\n\t\t\t\trequire.Error(t, err)\n\t\t\t\trequire.ErrorContains(t, err, c.ExpectedErrorMessage)\n\t\t\t})\n\t\t}\n\t})\n\n\tt.Run(\"EnforceDenyList\", func(t *testing.T) {\n\t\tfor _, pattern := range environmentVariableDenyList {\n\t\t\t// test that exact matches are rejected\n\t\t\tt.Run(fmt.Sprintf(\"Rejects %q\", pattern), func(t *testing.T) {\n\t\t\t\tinput := fmt.Sprintf(\"%s=VALUE\", pattern)\n\t\t\t\t_, err := parseAndValidateEnvironment([]string{input})\n\t\t\t\trequire.Error(t, err)\n\t\t\t\trequire.ErrorContains(t, err, fmt.Sprintf(\"environment variable %q is not allowed\", pattern))\n\t\t\t})\n\n\t\t\t// test that prefix matches are rejected\n\t\t\tif before, ok := strings.CutSuffix(pattern, \"*\"); ok {\n\t\t\t\tt.Run(fmt.Sprintf(\"Rejects %q prefix\", pattern), func(t *testing.T) {\n\t\t\t\t\tname := before + \"SUFFIX\"\n\t\t\t\t\tinput := fmt.Sprintf(\"%s=VALUE\", name)\n\t\t\t\t\t_, err := parseAndValidateEnvironment([]string{input})\n\t\t\t\t\trequire.Error(t, err)\n\t\t\t\t\trequire.ErrorContains(t, err, fmt.Sprintf(\"environment variable %q is not allowed\", name))\n\t\t\t\t})\n\t\t\t}\n\t\t}\n\t})\n\n\tt.Run(\"DuplicateNamesAreRejected\", func(t *testing.T) {\n\t\tinput := []string{\"NAME=VALUE\", \"NAME=VALUE2\"}\n\t\t_, err := parseAndValidateEnvironment(input)\n\t\trequire.Error(t, err)\n\t\trequire.ErrorContains(t, err, \"environment variable \\\"NAME\\\" is already defined\")\n\t})\n}\n"
  },
  {
    "path": "pkg/config/errors.go",
    "content": "package config\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n)\n\n// ConfigError is the base interface for all config errors.\n// Allows callers to use errors.As to get config-specific details.\ntype ConfigError interface {\n\terror\n\tConfigError() // marker method\n}\n\n// ParseError indicates the YAML file could not be parsed.\ntype ParseError struct {\n\tFilename string\n\tErr      error\n}\n\nfunc (e *ParseError) Error() string {\n\treturn fmt.Sprintf(\"failed to parse %s: %v\", e.Filename, e.Err)\n}\n\nfunc (e *ParseError) Unwrap() error {\n\treturn e.Err\n}\n\nfunc (e *ParseError) ConfigError() {}\n\n// SchemaError indicates the config structure doesn't match the schema.\n// For example, wrong type for a field or unknown field.\ntype SchemaError struct {\n\tField   string\n\tMessage string\n}\n\nfunc (e *SchemaError) Error() string {\n\treturn fmt.Sprintf(\"schema error in %q: %s\", e.Field, e.Message)\n}\n\nfunc (e *SchemaError) ConfigError() {}\n\n// ValidationError indicates a semantic validation failure.\n// The config parses correctly but values are invalid.\ntype ValidationError struct {\n\tField   string\n\tValue   string\n\tMessage string\n}\n\nfunc (e *ValidationError) Error() string {\n\tif e.Value != \"\" {\n\t\treturn fmt.Sprintf(\"invalid %s %q: %s\", e.Field, e.Value, e.Message)\n\t}\n\treturn fmt.Sprintf(\"invalid %s: %s\", e.Field, e.Message)\n}\n\nfunc (e *ValidationError) ConfigError() {}\n\n// DeprecationWarning indicates use of a deprecated field.\n// This is a warning, not an error - validation still succeeds.\ntype DeprecationWarning struct {\n\tField       string\n\tReplacement string\n\tMessage     string\n}\n\nfunc (w *DeprecationWarning) Error() string {\n\tif w.Replacement != \"\" {\n\t\treturn fmt.Sprintf(\"deprecated field %q: use %q instead\", w.Field, w.Replacement)\n\t}\n\treturn fmt.Sprintf(\"deprecated field %q: %s\", w.Field, w.Message)\n}\n\nfunc (w *DeprecationWarning) ConfigError() {}\n\n// CompatibilityError indicates an incompatible version combination.\ntype CompatibilityError struct {\n\tComponent1 string\n\tVersion1   string\n\tComponent2 string\n\tVersion2   string\n\tMessage    string\n}\n\nfunc (e *CompatibilityError) Error() string {\n\treturn fmt.Sprintf(\"%s %s is incompatible with %s %s: %s\",\n\t\te.Component1, e.Version1, e.Component2, e.Version2, e.Message)\n}\n\nfunc (e *CompatibilityError) ConfigError() {}\n\n// ValidationResult holds all errors and warnings from validation.\ntype ValidationResult struct {\n\tErrors   []error\n\tWarnings []DeprecationWarning\n}\n\n// HasErrors returns true if there are any validation errors.\nfunc (r *ValidationResult) HasErrors() bool {\n\treturn len(r.Errors) > 0\n}\n\n// HasWarnings returns true if there are any deprecation warnings.\nfunc (r *ValidationResult) HasWarnings() bool {\n\treturn len(r.Warnings) > 0\n}\n\n// Err returns a combined error if there are any validation errors, nil otherwise.\nfunc (r *ValidationResult) Err() error {\n\tif !r.HasErrors() {\n\t\treturn nil\n\t}\n\treturn errors.Join(r.Errors...)\n}\n\n// AddError adds a validation error.\nfunc (r *ValidationResult) AddError(err error) {\n\tr.Errors = append(r.Errors, err)\n}\n\n// AddWarning adds a deprecation warning.\nfunc (r *ValidationResult) AddWarning(w DeprecationWarning) {\n\tr.Warnings = append(r.Warnings, w)\n}\n\n// NewValidationResult creates an empty ValidationResult.\nfunc NewValidationResult() *ValidationResult {\n\treturn &ValidationResult{\n\t\tErrors:   []error{},\n\t\tWarnings: []DeprecationWarning{},\n\t}\n}\n"
  },
  {
    "path": "pkg/config/image_name.go",
    "content": "package config\n\nimport (\n\t\"path\"\n\t\"regexp\"\n\t\"strings\"\n)\n\n// DockerImageName returns the default Docker image name for images\nfunc DockerImageName(projectDir string) string {\n\tprefix := \"cog-\"\n\tprojectName := strings.ToLower(path.Base(projectDir))\n\n\t// Convert whitespace to dashes\n\tprojectName = strings.ReplaceAll(projectName, \" \", \"-\")\n\n\t// Remove anything non-alphanumeric\n\treg := regexp.MustCompile(`[^a-z0-9\\-]+`)\n\tprojectName = reg.ReplaceAllString(projectName, \"\")\n\n\t// Limit to 30 characters (max Docker image name length)\n\tlength := 30 - len(prefix)\n\tif len(projectName) > length {\n\t\tprojectName = projectName[:length]\n\t}\n\n\tif !strings.HasPrefix(projectName, prefix) {\n\t\tprojectName = prefix + projectName\n\t}\n\n\treturn projectName\n}\n"
  },
  {
    "path": "pkg/config/image_name_test.go",
    "content": "package config\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestDockerImageName(t *testing.T) {\n\trequire.Equal(t, \"cog-foo\", DockerImageName(\"/home/joe/foo\"))\n\trequire.Equal(t, \"cog-foo\", DockerImageName(\"/home/joe/Foo\"))\n\trequire.Equal(t, \"cog-foo\", DockerImageName(\"/home/joe/cog-foo\"))\n\trequire.Equal(t, \"cog-my-great-model\", DockerImageName(\"/home/joe/my great model\"))\n\trequire.Equal(t, 30, len(DockerImageName(\"/home/joe/verylongverylongverylongverylongverylongverylongverylong\")))\n}\n"
  },
  {
    "path": "pkg/config/load.go",
    "content": "package config\n\nimport (\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n\t\"path/filepath\"\n\n\t\"github.com/replicate/cog/pkg/errors\"\n\t\"github.com/replicate/cog/pkg/util/files\"\n)\n\nconst maxSearchDepth = 100\n\n// LoadResult contains the loaded config and any warnings.\ntype LoadResult struct {\n\tConfig   *Config\n\tWarnings []DeprecationWarning\n\tRootDir  string\n}\n\n// Load parses, validates, and completes a config from an io.Reader.\n// The projectDir is used for validation (checking that referenced files exist)\n// and for completion (resolving CUDA versions, loading requirements files, etc.).\n// Always returns warnings if present, even on success.\nfunc Load(r io.Reader, projectDir string) (*LoadResult, error) {\n\t// Parse\n\tcfgFile, err := parse(r)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Validate\n\tvalidationResult := ValidateConfigFile(cfgFile, WithProjectDir(projectDir))\n\n\t// Collect warnings\n\twarnings := validationResult.Warnings\n\n\t// Check for errors\n\tif validationResult.HasErrors() {\n\t\treturn nil, validationResult.Err()\n\t}\n\n\t// Convert to Config struct\n\tconfig, err := configFileToConfig(cfgFile)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Complete (resolve CUDA, load requirements, etc.)\n\tif err := config.Complete(projectDir); err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &LoadResult{\n\t\tConfig:   config,\n\t\tWarnings: warnings,\n\t\tRootDir:  projectDir,\n\t}, nil\n}\n\n// GetProjectDir returns the project's root directory by searching for\n// the config file starting from the current working directory.\nfunc GetProjectDir(configFilename string) (string, error) {\n\tif configFilename == \"\" {\n\t\tconfigFilename = \"cog.yaml\"\n\t}\n\n\tcwd, err := os.Getwd()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn findProjectRootDir(cwd, configFilename)\n}\n\n// findConfigPathInDirectory checks if the config file exists in the given directory.\nfunc findConfigPathInDirectory(dir string, configFilename string) (configPath string, err error) {\n\tfilePath := filepath.Join(dir, configFilename)\n\texists, err := files.Exists(filePath)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to scan directory %s for %s: %w\", dir, filePath, err)\n\t} else if exists {\n\t\treturn filePath, nil\n\t}\n\n\treturn \"\", errors.ConfigNotFound(fmt.Sprintf(\"%s not found in %s\", configFilename, dir))\n}\n\n// findProjectRootDir walks up the directory tree to find the root of the project.\n// The project root is defined as the directory housing a `cog.yaml` file.\nfunc findProjectRootDir(startDir string, configFilename string) (string, error) {\n\tdir := startDir\n\tfor range maxSearchDepth {\n\t\tswitch _, err := findConfigPathInDirectory(dir, configFilename); {\n\t\tcase err != nil && !errors.IsConfigNotFound(err):\n\t\t\treturn \"\", err\n\t\tcase err == nil:\n\t\t\treturn dir, nil\n\t\tcase dir == \".\" || dir == \"/\":\n\t\t\treturn \"\", errors.ConfigNotFound(fmt.Sprintf(\"%s not found in %s (or in any parent directories)\", configFilename, startDir))\n\t\t}\n\n\t\tdir = filepath.Dir(dir)\n\t}\n\n\treturn \"\", errors.ConfigNotFound(\"No cog.yaml found in parent directories.\")\n}\n"
  },
  {
    "path": "pkg/config/load_test.go",
    "content": "package config\n\nimport (\n\t\"os\"\n\t\"path\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nconst testConfig = `\nbuild:\n  python_version: \"3.10\"\n  python_requirements: requirements.txt\n  system_packages:\n    - libgl1-mesa-glx\n    - libglib2.0-0\npredict: \"predict.py:SomePredictor\"\n`\n\nfunc TestFindProjectRootDirShouldFindParentDir(t *testing.T) {\n\tprojectDir := t.TempDir()\n\n\terr := os.WriteFile(path.Join(projectDir, \"cog.yaml\"), []byte(testConfig), 0o644)\n\trequire.NoError(t, err)\n\n\tsubdir := path.Join(projectDir, \"some/sub/dir\")\n\terr = os.MkdirAll(subdir, 0o700)\n\trequire.NoError(t, err)\n\n\tfoundDir, err := findProjectRootDir(subdir, \"cog.yaml\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, foundDir, projectDir)\n}\n\nfunc TestFindProjectRootDirShouldReturnErrIfNoConfig(t *testing.T) {\n\tprojectDir := t.TempDir()\n\n\tsubdir := path.Join(projectDir, \"some/sub/dir\")\n\terr := os.MkdirAll(subdir, 0o700)\n\trequire.NoError(t, err)\n\n\t_, err = findProjectRootDir(subdir, \"cog.yaml\")\n\trequire.Error(t, err)\n}\n"
  },
  {
    "path": "pkg/config/parse.go",
    "content": "package config\n\nimport (\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n\t\"path/filepath\"\n\n\t\"go.yaml.in/yaml/v4\"\n\n\t\"github.com/replicate/cog/pkg/util/files\"\n)\n\n// parse reads and parses YAML content from an io.Reader into a configFile.\n// This only does YAML parsing - no validation or defaults.\n// Returns ParseError if the content cannot be read or parsed.\nfunc parse(r io.Reader) (*configFile, error) {\n\tcontents, err := io.ReadAll(r)\n\tif err != nil {\n\t\treturn nil, &ParseError{Err: err}\n\t}\n\n\treturn parseBytes(contents)\n}\n\n// parseFile reads and parses a cog.yaml file into a configFile.\n// This only does YAML parsing - no validation or defaults.\n// Returns ParseError if the file cannot be read or parsed.\nfunc parseFile(filename string) (*configFile, error) {\n\texists, err := files.Exists(filename)\n\tif err != nil {\n\t\treturn nil, &ParseError{Filename: filename, Err: err}\n\t}\n\n\tif !exists {\n\t\treturn nil, &ParseError{\n\t\t\tFilename: filename,\n\t\t\tErr:      fmt.Errorf(\"%s does not exist in %s\", filepath.Base(filename), filepath.Dir(filename)),\n\t\t}\n\t}\n\n\tf, err := os.Open(filename)\n\tif err != nil {\n\t\treturn nil, &ParseError{Filename: filename, Err: err}\n\t}\n\tdefer f.Close()\n\n\tcfg, err := parse(f)\n\tif err != nil {\n\t\t// Add filename context to the error\n\t\tif parseErr, ok := err.(*ParseError); ok {\n\t\t\tparseErr.Filename = filename\n\t\t\treturn nil, parseErr\n\t\t}\n\t\treturn nil, &ParseError{Filename: filename, Err: err}\n\t}\n\n\treturn cfg, nil\n}\n\n// parseBytes parses YAML content into a configFile.\nfunc parseBytes(contents []byte) (*configFile, error) {\n\tcfg := &configFile{}\n\n\tif len(contents) == 0 {\n\t\t// Empty file is valid, returns empty config\n\t\treturn cfg, nil\n\t}\n\n\tif err := yaml.Unmarshal(contents, cfg); err != nil {\n\t\treturn nil, &ParseError{\n\t\t\tErr: fmt.Errorf(\"invalid YAML: %w\", err),\n\t\t}\n\t}\n\n\treturn cfg, nil\n}\n\n// FromYAML parses YAML content into an uncompleted Config.\n// This is a convenience function primarily for testing.\n// Callers should call Complete() on the returned config to resolve CUDA versions etc.\n// For production code, use Load() which handles validation and completion.\n//\n// Note: This function skips validation since it has no project directory context.\n// The Complete() method will validate requirements files exist when called.\nfunc FromYAML(contents []byte) (*Config, error) {\n\tcfgFile, err := parseBytes(contents)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Convert to Config struct without completion or validation\n\t// The caller should call Complete() with the appropriate project dir\n\treturn configFileToConfig(cfgFile)\n}\n\n// configFileToConfig converts a ConfigFile to a Config without running completion logic.\n// This is the minimal conversion used by FromYAML for test compatibility.\nfunc configFileToConfig(cfg *configFile) (*Config, error) {\n\tconfig := &Config{\n\t\tBuild: &Build{},\n\t}\n\n\tif cfg.Build != nil {\n\t\tif cfg.Build.GPU != nil {\n\t\t\tconfig.Build.GPU = *cfg.Build.GPU\n\t\t}\n\t\tif cfg.Build.PythonVersion != nil {\n\t\t\tconfig.Build.PythonVersion = *cfg.Build.PythonVersion\n\t\t}\n\t\tif cfg.Build.PythonRequirements != nil {\n\t\t\tconfig.Build.PythonRequirements = *cfg.Build.PythonRequirements\n\t\t}\n\t\tconfig.Build.PythonPackages = cfg.Build.PythonPackages\n\t\tconfig.Build.SystemPackages = cfg.Build.SystemPackages\n\t\tconfig.Build.PreInstall = cfg.Build.PreInstall\n\t\tif cfg.Build.CUDA != nil {\n\t\t\tconfig.Build.CUDA = *cfg.Build.CUDA\n\t\t}\n\t\tif cfg.Build.CuDNN != nil {\n\t\t\tconfig.Build.CuDNN = *cfg.Build.CuDNN\n\t\t}\n\t\tif cfg.Build.SDKVersion != nil {\n\t\t\tconfig.Build.SDKVersion = *cfg.Build.SDKVersion\n\t\t}\n\n\t\t// Convert Run items\n\t\tconfig.Build.Run = make([]RunItem, len(cfg.Build.Run))\n\t\tfor i, runFile := range cfg.Build.Run {\n\t\t\tconfig.Build.Run[i] = RunItem{\n\t\t\t\tCommand: runFile.Command,\n\t\t\t}\n\t\t\tif len(runFile.Mounts) > 0 {\n\t\t\t\tconfig.Build.Run[i].Mounts = make([]struct {\n\t\t\t\t\tType   string `json:\"type,omitempty\" yaml:\"type\"`\n\t\t\t\t\tID     string `json:\"id,omitempty\" yaml:\"id\"`\n\t\t\t\t\tTarget string `json:\"target,omitempty\" yaml:\"target\"`\n\t\t\t\t}, len(runFile.Mounts))\n\t\t\t\tfor j, mountFile := range runFile.Mounts {\n\t\t\t\t\tconfig.Build.Run[i].Mounts[j].Type = mountFile.Type\n\t\t\t\t\tconfig.Build.Run[i].Mounts[j].ID = mountFile.ID\n\t\t\t\t\tconfig.Build.Run[i].Mounts[j].Target = mountFile.Target\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tif cfg.Image != nil {\n\t\tconfig.Image = *cfg.Image\n\t}\n\tif cfg.Predict != nil {\n\t\tconfig.Predict = *cfg.Predict\n\t}\n\tif cfg.Train != nil {\n\t\tconfig.Train = *cfg.Train\n\t}\n\tif cfg.Concurrency != nil {\n\t\tconfig.Concurrency = &Concurrency{}\n\t\tif cfg.Concurrency.Max != nil {\n\t\t\tconfig.Concurrency.Max = *cfg.Concurrency.Max\n\t\t}\n\t}\n\tconfig.Environment = cfg.Environment\n\n\t// Convert weights\n\tif len(cfg.Weights) > 0 {\n\t\tconfig.Weights = make([]WeightSource, len(cfg.Weights))\n\t\tfor i, w := range cfg.Weights {\n\t\t\tconfig.Weights[i] = WeightSource(w)\n\t\t}\n\t}\n\n\treturn config, nil\n}\n"
  },
  {
    "path": "pkg/config/tf_compatibility.json",
    "content": "[\n  {\n    \"TF\": \"2.20.0\",\n    \"TFCPUPackage\": \"tensorflow==2.20.0\",\n    \"TFGPUPackage\": \"tensorflow==2.20.0\",\n    \"CUDA\": \"12.5\",\n    \"CuDNN\": \"9.3\",\n    \"Pythons\": [\n      \"3.9\",\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"TF\": \"2.19.0\",\n    \"TFCPUPackage\": \"tensorflow==2.19.0\",\n    \"TFGPUPackage\": \"tensorflow==2.19.0\",\n    \"CUDA\": \"12.5\",\n    \"CuDNN\": \"9.3\",\n    \"Pythons\": [\n      \"3.9\",\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"TF\": \"2.18.0\",\n    \"TFCPUPackage\": \"tensorflow==2.18.0\",\n    \"TFGPUPackage\": \"tensorflow==2.18.0\",\n    \"CUDA\": \"12.5\",\n    \"CuDNN\": \"9.3\",\n    \"Pythons\": [\n      \"3.9\",\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"TF\": \"2.17.0\",\n    \"TFCPUPackage\": \"tensorflow==2.17.0\",\n    \"TFGPUPackage\": \"tensorflow==2.17.0\",\n    \"CUDA\": \"12.3\",\n    \"CuDNN\": \"8.9\",\n    \"Pythons\": [\n      \"3.9\",\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"TF\": \"2.16.1\",\n    \"TFCPUPackage\": \"tensorflow==2.16.1\",\n    \"TFGPUPackage\": \"tensorflow==2.16.1\",\n    \"CUDA\": \"12.3\",\n    \"CuDNN\": \"8.9\",\n    \"Pythons\": [\n      \"3.9\",\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"TF\": \"2.15.0\",\n    \"TFCPUPackage\": \"tensorflow==2.15.0\",\n    \"TFGPUPackage\": \"tensorflow==2.15.0\",\n    \"CUDA\": \"12.2\",\n    \"CuDNN\": \"8.9\",\n    \"Pythons\": [\n      \"3.9\",\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"TF\": \"2.14.0\",\n    \"TFCPUPackage\": \"tensorflow==2.14.0\",\n    \"TFGPUPackage\": \"tensorflow==2.14.0\",\n    \"CUDA\": \"11.8\",\n    \"CuDNN\": \"8.7\",\n    \"Pythons\": [\n      \"3.9\",\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"TF\": \"2.13.0\",\n    \"TFCPUPackage\": \"tensorflow==2.13.0\",\n    \"TFGPUPackage\": \"tensorflow==2.13.0\",\n    \"CUDA\": \"11.8\",\n    \"CuDNN\": \"8.6\",\n    \"Pythons\": [\n      \"3.8\",\n      \"3.9\",\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"TF\": \"2.12.0\",\n    \"TFCPUPackage\": \"tensorflow==2.12.0\",\n    \"TFGPUPackage\": \"tensorflow==2.12.0\",\n    \"CUDA\": \"11.8\",\n    \"CuDNN\": \"8.6\",\n    \"Pythons\": [\n      \"3.8\",\n      \"3.9\",\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"TF\": \"2.11.0\",\n    \"TFCPUPackage\": \"tensorflow==2.11.0\",\n    \"TFGPUPackage\": \"tensorflow==2.11.0\",\n    \"CUDA\": \"11.2\",\n    \"CuDNN\": \"8.1\",\n    \"Pythons\": [\n      \"3.7\",\n      \"3.8\",\n      \"3.9\",\n      \"3.10\"\n    ]\n  },\n  {\n    \"TF\": \"2.10.0\",\n    \"TFCPUPackage\": \"tensorflow==2.10.0\",\n    \"TFGPUPackage\": \"tensorflow==2.10.0\",\n    \"CUDA\": \"11.2\",\n    \"CuDNN\": \"8.1\",\n    \"Pythons\": [\n      \"3.7\",\n      \"3.8\",\n      \"3.9\",\n      \"3.10\"\n    ]\n  },\n  {\n    \"TF\": \"2.9.0\",\n    \"TFCPUPackage\": \"tensorflow==2.9.0\",\n    \"TFGPUPackage\": \"tensorflow==2.9.0\",\n    \"CUDA\": \"11.2\",\n    \"CuDNN\": \"8.1\",\n    \"Pythons\": [\n      \"3.7\",\n      \"3.8\",\n      \"3.9\",\n      \"3.10\"\n    ]\n  },\n  {\n    \"TF\": \"2.8.0\",\n    \"TFCPUPackage\": \"tensorflow==2.8.0\",\n    \"TFGPUPackage\": \"tensorflow==2.8.0\",\n    \"CUDA\": \"11.2\",\n    \"CuDNN\": \"8.1\",\n    \"Pythons\": [\n      \"3.7\",\n      \"3.8\",\n      \"3.9\",\n      \"3.10\"\n    ]\n  },\n  {\n    \"TF\": \"2.7.0\",\n    \"TFCPUPackage\": \"tensorflow==2.7.0\",\n    \"TFGPUPackage\": \"tensorflow==2.7.0\",\n    \"CUDA\": \"11.2\",\n    \"CuDNN\": \"8.1\",\n    \"Pythons\": [\n      \"3.7\",\n      \"3.8\",\n      \"3.9\"\n    ]\n  },\n  {\n    \"TF\": \"2.6.0\",\n    \"TFCPUPackage\": \"tensorflow==2.6.0\",\n    \"TFGPUPackage\": \"tensorflow==2.6.0\",\n    \"CUDA\": \"11.2\",\n    \"CuDNN\": \"8.1\",\n    \"Pythons\": [\n      \"3.6\",\n      \"3.7\",\n      \"3.8\",\n      \"3.9\"\n    ]\n  },\n  {\n    \"TF\": \"2.5.0\",\n    \"TFCPUPackage\": \"tensorflow==2.5.0\",\n    \"TFGPUPackage\": \"tensorflow==2.5.0\",\n    \"CUDA\": \"11.2\",\n    \"CuDNN\": \"8.1\",\n    \"Pythons\": [\n      \"3.6\",\n      \"3.7\",\n      \"3.8\",\n      \"3.9\"\n    ]\n  },\n  {\n    \"TF\": \"2.4.0\",\n    \"TFCPUPackage\": \"tensorflow==2.4.0\",\n    \"TFGPUPackage\": \"tensorflow==2.4.0\",\n    \"CUDA\": \"11.0\",\n    \"CuDNN\": \"8.0\",\n    \"Pythons\": [\n      \"3.6\",\n      \"3.7\",\n      \"3.8\"\n    ]\n  }\n]"
  },
  {
    "path": "pkg/config/torch_compatibility.json",
    "content": "[\n  {\n    \"Torch\": \"2.10.0+cu129\",\n    \"Torchvision\": \"0.25.0\",\n    \"Torchaudio\": \"2.10.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu129/\",\n    \"CUDA\": \"12.9\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.10.0+cu130\",\n    \"Torchvision\": \"0.25.0\",\n    \"Torchaudio\": \"2.10.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu130/\",\n    \"CUDA\": \"13.0\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.10.0+cpu\",\n    \"Torchvision\": \"0.25.0\",\n    \"Torchaudio\": \"2.10.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu/\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.10.0+cu126\",\n    \"Torchvision\": \"0.25.0\",\n    \"Torchaudio\": \"2.10.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu126/\",\n    \"CUDA\": \"12.6\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.10.0+cu128\",\n    \"Torchvision\": \"0.25.0\",\n    \"Torchaudio\": \"2.10.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu128/\",\n    \"CUDA\": \"12.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.9.1\",\n    \"Torchvision\": \"0.24.1\",\n    \"Torchaudio\": \"2.9.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu126\",\n    \"CUDA\": \"12.6\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.9.1\",\n    \"Torchvision\": \"0.24.1\",\n    \"Torchaudio\": \"2.9.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu128\",\n    \"CUDA\": \"12.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.9.1\",\n    \"Torchvision\": \"0.24.1\",\n    \"Torchaudio\": \"2.9.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu130\",\n    \"CUDA\": \"13.0\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.9.1\",\n    \"Torchvision\": \"0.24.1\",\n    \"Torchaudio\": \"2.9.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.9.0\",\n    \"Torchvision\": \"0.24.0\",\n    \"Torchaudio\": \"2.9.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu126\",\n    \"CUDA\": \"12.6\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.9.0\",\n    \"Torchvision\": \"0.24.0\",\n    \"Torchaudio\": \"2.9.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu128\",\n    \"CUDA\": \"12.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.9.0\",\n    \"Torchvision\": \"0.24.0\",\n    \"Torchaudio\": \"2.9.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu130\",\n    \"CUDA\": \"13.0\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.9.0\",\n    \"Torchvision\": \"0.24.0\",\n    \"Torchaudio\": \"2.9.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\",\n      \"3.14\"\n    ]\n  },\n  {\n    \"Torch\": \"2.8.0\",\n    \"Torchvision\": \"0.23.0\",\n    \"Torchaudio\": \"2.8.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu126\",\n    \"CUDA\": \"12.6\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.8.0\",\n    \"Torchvision\": \"0.23.0\",\n    \"Torchaudio\": \"2.8.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu128\",\n    \"CUDA\": \"12.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.8.0\",\n    \"Torchvision\": \"0.23.0\",\n    \"Torchaudio\": \"2.8.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu129\",\n    \"CUDA\": \"12.9\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.8.0\",\n    \"Torchvision\": \"0.23.0\",\n    \"Torchaudio\": \"2.8.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.7.1\",\n    \"Torchvision\": \"0.22.1\",\n    \"Torchaudio\": \"2.7.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.7.1\",\n    \"Torchvision\": \"0.22.1\",\n    \"Torchaudio\": \"2.7.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu126\",\n    \"CUDA\": \"12.6\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.7.1\",\n    \"Torchvision\": \"0.22.1\",\n    \"Torchaudio\": \"2.7.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu128\",\n    \"CUDA\": \"12.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.7.1\",\n    \"Torchvision\": \"0.22.1\",\n    \"Torchaudio\": \"2.7.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.7.0\",\n    \"Torchvision\": \"0.22.0\",\n    \"Torchaudio\": \"2.7.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.7.0\",\n    \"Torchvision\": \"0.22.0\",\n    \"Torchaudio\": \"2.7.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu126\",\n    \"CUDA\": \"12.6\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.7.0\",\n    \"Torchvision\": \"0.22.0\",\n    \"Torchaudio\": \"2.7.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu128\",\n    \"CUDA\": \"12.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.7.0\",\n    \"Torchvision\": \"0.22.0\",\n    \"Torchaudio\": \"2.7.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.6.0\",\n    \"Torchvision\": \"0.21.0\",\n    \"Torchaudio\": \"2.6.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.6.0\",\n    \"Torchvision\": \"0.21.0\",\n    \"Torchaudio\": \"2.6.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu124\",\n    \"CUDA\": \"12.4\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.6.0\",\n    \"Torchvision\": \"0.21.0\",\n    \"Torchaudio\": \"2.6.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu126\",\n    \"CUDA\": \"12.6\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.6.0\",\n    \"Torchvision\": \"0.21.0\",\n    \"Torchaudio\": \"2.6.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\",\n      \"3.13\"\n    ]\n  },\n  {\n    \"Torch\": \"2.5.1\",\n    \"Torchvision\": \"0.20.1\",\n    \"Torchaudio\": \"2.5.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.5.1\",\n    \"Torchvision\": \"0.20.1\",\n    \"Torchaudio\": \"2.5.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.5.1\",\n    \"Torchvision\": \"0.20.1\",\n    \"Torchaudio\": \"2.5.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu124\",\n    \"CUDA\": \"12.4\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.5.1\",\n    \"Torchvision\": \"0.20.1\",\n    \"Torchaudio\": \"2.5.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.5.0\",\n    \"Torchvision\": \"0.20.0\",\n    \"Torchaudio\": \"2.5.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.5.0\",\n    \"Torchvision\": \"0.20.0\",\n    \"Torchaudio\": \"2.5.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.5.0\",\n    \"Torchvision\": \"0.20.0\",\n    \"Torchaudio\": \"2.5.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu124\",\n    \"CUDA\": \"12.4\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.5.0\",\n    \"Torchvision\": \"0.20.0\",\n    \"Torchaudio\": \"2.5.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.4.1\",\n    \"Torchvision\": \"0.19.1\",\n    \"Torchaudio\": \"2.4.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.4.1\",\n    \"Torchvision\": \"0.19.1\",\n    \"Torchaudio\": \"2.4.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.4.1\",\n    \"Torchvision\": \"0.19.1\",\n    \"Torchaudio\": \"2.4.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu124\",\n    \"CUDA\": \"12.4\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.4.1\",\n    \"Torchvision\": \"0.19.1\",\n    \"Torchaudio\": \"2.4.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.4.0\",\n    \"Torchvision\": \"0.19.0\",\n    \"Torchaudio\": \"2.4.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.4.0\",\n    \"Torchvision\": \"0.19.0\",\n    \"Torchaudio\": \"2.4.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.4.0\",\n    \"Torchvision\": \"0.19.0\",\n    \"Torchaudio\": \"2.4.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu124\",\n    \"CUDA\": \"12.4\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.4.0\",\n    \"Torchvision\": \"0.19.0\",\n    \"Torchaudio\": \"2.4.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.3.1\",\n    \"Torchvision\": \"0.18.1\",\n    \"Torchaudio\": \"2.3.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.3.1\",\n    \"Torchvision\": \"0.18.1\",\n    \"Torchaudio\": \"2.3.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.3.1\",\n    \"Torchvision\": \"0.18.1\",\n    \"Torchaudio\": \"2.3.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.3.0\",\n    \"Torchvision\": \"0.18.0\",\n    \"Torchaudio\": \"2.3.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.3.0\",\n    \"Torchvision\": \"0.18.0\",\n    \"Torchaudio\": \"2.3.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.3.0\",\n    \"Torchvision\": \"0.18.0\",\n    \"Torchaudio\": \"2.3.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.2\",\n    \"Torchvision\": \"0.17.2\",\n    \"Torchaudio\": \"2.2.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.2\",\n    \"Torchvision\": \"0.17.2\",\n    \"Torchaudio\": \"2.2.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.2\",\n    \"Torchvision\": \"0.17.2\",\n    \"Torchaudio\": \"2.2.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.1\",\n    \"Torchvision\": \"0.17.1\",\n    \"Torchaudio\": \"2.2.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.1\",\n    \"Torchvision\": \"0.17.1\",\n    \"Torchaudio\": \"2.2.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.1\",\n    \"Torchvision\": \"0.17.1\",\n    \"Torchaudio\": \"2.2.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.0\",\n    \"Torchvision\": \"0.17.0\",\n    \"Torchaudio\": \"2.2.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.0\",\n    \"Torchvision\": \"0.17.0\",\n    \"Torchaudio\": \"2.2.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.2.0\",\n    \"Torchvision\": \"0.17.0\",\n    \"Torchaudio\": \"2.2.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\",\n      \"3.12\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.2\",\n    \"Torchvision\": \"0.16.2\",\n    \"Torchaudio\": \"2.1.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.2\",\n    \"Torchvision\": \"0.16.2\",\n    \"Torchaudio\": \"2.1.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.2\",\n    \"Torchvision\": \"0.16.2\",\n    \"Torchaudio\": \"2.1.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.1\",\n    \"Torchvision\": \"0.16.1\",\n    \"Torchaudio\": \"2.1.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.1\",\n    \"Torchvision\": \"0.16.1\",\n    \"Torchaudio\": \"2.1.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.1\",\n    \"Torchvision\": \"0.16.1\",\n    \"Torchaudio\": \"2.1.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.0\",\n    \"Torchvision\": \"0.16.0\",\n    \"Torchaudio\": \"2.1.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.0\",\n    \"Torchvision\": \"0.16.0\",\n    \"Torchaudio\": \"2.1.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu121\",\n    \"CUDA\": \"12.1\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.1.0\",\n    \"Torchvision\": \"0.16.0\",\n    \"Torchaudio\": \"2.1.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.0.1\",\n    \"Torchvision\": \"0.15.2\",\n    \"Torchaudio\": \"2.0.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"\",\n    \"CUDA\": \"11.7\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.0.1\",\n    \"Torchvision\": \"0.15.2\",\n    \"Torchaudio\": \"2.0.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.0.1\",\n    \"Torchvision\": \"0.15.2\",\n    \"Torchaudio\": \"2.0.2\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.0.0\",\n    \"Torchvision\": \"0.15.1\",\n    \"Torchaudio\": \"2.0.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"\",\n    \"CUDA\": \"11.7\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.0.0\",\n    \"Torchvision\": \"0.15.1\",\n    \"Torchaudio\": \"2.0.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu118\",\n    \"CUDA\": \"11.8\",\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"2.0.0\",\n    \"Torchvision\": \"0.15.1\",\n    \"Torchaudio\": \"2.0.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\",\n      \"3.11\"\n    ]\n  },\n  {\n    \"Torch\": \"1.13.1+cu116\",\n    \"Torchvision\": \"0.14.1+cu116\",\n    \"Torchaudio\": \"0.13.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu116\",\n    \"CUDA\": \"11.6\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.13.1+cu117\",\n    \"Torchvision\": \"0.14.1+cu117\",\n    \"Torchaudio\": \"0.13.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu117\",\n    \"CUDA\": \"11.7\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.13.1+cpu\",\n    \"Torchvision\": \"0.14.1+cpu\",\n    \"Torchaudio\": \"0.13.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.13.0+cu116\",\n    \"Torchvision\": \"0.14.0+cu116\",\n    \"Torchaudio\": \"0.13.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu116\",\n    \"CUDA\": \"11.6\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.13.0+cu117\",\n    \"Torchvision\": \"0.14.0+cu117\",\n    \"Torchaudio\": \"0.13.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu117\",\n    \"CUDA\": \"11.7\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.13.0+cpu\",\n    \"Torchvision\": \"0.14.0+cpu\",\n    \"Torchaudio\": \"0.13.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.12.1+cu116\",\n    \"Torchvision\": \"0.13.1+cu116\",\n    \"Torchaudio\": \"0.12.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu116\",\n    \"CUDA\": \"11.6\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.12.1+cu113\",\n    \"Torchvision\": \"0.13.1+cu113\",\n    \"Torchaudio\": \"0.12.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu113\",\n    \"CUDA\": \"11.3\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.12.1+cu102\",\n    \"Torchvision\": \"0.13.1+cu102\",\n    \"Torchaudio\": \"0.12.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu102\",\n    \"CUDA\": \"10.2\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.12.1+cpu\",\n    \"Torchvision\": \"0.13.1+cpu\",\n    \"Torchaudio\": \"0.12.1\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.12.0+cu116\",\n    \"Torchvision\": \"0.13.0+cu116\",\n    \"Torchaudio\": \"0.12.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu116\",\n    \"CUDA\": \"11.6\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.12.0+cu113\",\n    \"Torchvision\": \"0.13.0+cu113\",\n    \"Torchaudio\": \"0.12.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu113\",\n    \"CUDA\": \"11.3\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.12.0+cu102\",\n    \"Torchvision\": \"0.13.0+cu102\",\n    \"Torchaudio\": \"0.12.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu102\",\n    \"CUDA\": \"10.2\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.12.0+cpu\",\n    \"Torchvision\": \"0.13.0+cpu\",\n    \"Torchaudio\": \"0.12.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.11.0+cu113\",\n    \"Torchvision\": \"0.12.0+cu113\",\n    \"Torchaudio\": \"0.11.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu113\",\n    \"CUDA\": \"11.3\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.11.0+cu102\",\n    \"Torchvision\": \"0.12.0+cu102\",\n    \"Torchaudio\": \"0.11.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cu102\",\n    \"CUDA\": \"10.2\",\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.11.0+cpu\",\n    \"Torchvision\": \"0.12.0+cpu\",\n    \"Torchaudio\": \"0.11.0\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"https://download.pytorch.org/whl/cpu\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"3.10\"\n    ]\n  },\n  {\n    \"Torch\": \"1.4.0\",\n    \"Torchvision\": \"0.5.0\",\n    \"Torchaudio\": \"\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"\",\n    \"CUDA\": \"10.1\",\n    \"Pythons\": [\n      \"2.7\"\n    ]\n  },\n  {\n    \"Torch\": \"1.4.0+cu92\",\n    \"Torchvision\": \"0.5.0+cu92\",\n    \"Torchaudio\": \"\",\n    \"FindLinks\": \"https://download.pytorch.org/whl/torch_stable.html\",\n    \"ExtraIndexURL\": \"\",\n    \"CUDA\": \"9.2\",\n    \"Pythons\": [\n      \"2.7\"\n    ]\n  },\n  {\n    \"Torch\": \"1.4.0+cpu\",\n    \"Torchvision\": \"0.5.0+cpu\",\n    \"Torchaudio\": \"\",\n    \"FindLinks\": \"https://download.pytorch.org/whl/torch_stable.html\",\n    \"ExtraIndexURL\": \"\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"2.7\"\n    ]\n  },\n  {\n    \"Torch\": \"1.2.0\",\n    \"Torchvision\": \"0.4.0\",\n    \"Torchaudio\": \"\",\n    \"FindLinks\": \"\",\n    \"ExtraIndexURL\": \"\",\n    \"CUDA\": \"10.0\",\n    \"Pythons\": [\n      \"2.7\"\n    ]\n  },\n  {\n    \"Torch\": \"1.2.0+cu92\",\n    \"Torchvision\": \"0.4.0+cu92\",\n    \"Torchaudio\": \"\",\n    \"FindLinks\": \"https://download.pytorch.org/whl/torch_stable.html\",\n    \"ExtraIndexURL\": \"\",\n    \"CUDA\": \"9.2\",\n    \"Pythons\": [\n      \"2.7\"\n    ]\n  },\n  {\n    \"Torch\": \"1.2.0+cpu\",\n    \"Torchvision\": \"0.4.0+cpu\",\n    \"Torchaudio\": \"\",\n    \"FindLinks\": \"https://download.pytorch.org/whl/torch_stable.html\",\n    \"ExtraIndexURL\": \"\",\n    \"CUDA\": null,\n    \"Pythons\": [\n      \"2.7\"\n    ]\n  }\n]"
  },
  {
    "path": "pkg/config/validate.go",
    "content": "package config\n\nimport (\n\t// blank import for embeds\n\t_ \"embed\"\n\t\"fmt\"\n\t\"io/fs\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"slices\"\n\t\"strconv\"\n\t\"strings\"\n\n\t\"github.com/xeipuuv/gojsonschema\"\n\n\t\"github.com/replicate/cog/pkg/requirements\"\n)\n\n//go:embed data/config_schema_v1.0.json\nvar schemaV1 []byte\n\n// ValidateOption configures validation behavior.\ntype ValidateOption func(*validateOptions)\n\ntype validateOptions struct {\n\tprojectDir         string\n\trequirementsFS     fs.FS\n\tstrictDeprecations bool\n}\n\n// WithProjectDir sets the project directory for resolving relative paths.\nfunc WithProjectDir(dir string) ValidateOption {\n\treturn func(o *validateOptions) {\n\t\to.projectDir = dir\n\t}\n}\n\n// WithRequirementsFS sets the filesystem for reading python_requirements file.\nfunc WithRequirementsFS(fsys fs.FS) ValidateOption {\n\treturn func(o *validateOptions) {\n\t\to.requirementsFS = fsys\n\t}\n}\n\n// WithStrictDeprecations treats deprecation warnings as errors.\nfunc WithStrictDeprecations() ValidateOption {\n\treturn func(o *validateOptions) {\n\t\to.strictDeprecations = true\n\t}\n}\n\n// ValidateConfigFile checks a configFile for errors.\n// Returns all validation errors and deprecation warnings.\n// Does not mutate the input.\nfunc ValidateConfigFile(cfg *configFile, opts ...ValidateOption) *ValidationResult {\n\toptions := &validateOptions{}\n\tfor _, opt := range opts {\n\t\topt(options)\n\t}\n\n\tresult := NewValidationResult()\n\n\t// Schema validation\n\tif err := validateSchema(cfg); err != nil {\n\t\tresult.AddError(err)\n\t}\n\n\t// Semantic validation\n\tvalidatePredict(cfg, result)\n\tvalidateTrain(cfg, result)\n\tvalidateBuild(cfg, options, result)\n\tvalidateEnvironment(cfg, result)\n\tvalidateConcurrency(cfg, result)\n\n\t// Check deprecated fields\n\tcheckDeprecatedFields(cfg, result)\n\n\t// If strict deprecations, convert warnings to errors\n\tif options.strictDeprecations && result.HasWarnings() {\n\t\tfor _, w := range result.Warnings {\n\t\t\tresult.AddError(&w)\n\t\t}\n\t\tresult.Warnings = nil\n\t}\n\n\treturn result\n}\n\n// validateSchema validates the config against the JSON schema.\nfunc validateSchema(cfg *configFile) error {\n\tschemaLoader := gojsonschema.NewStringLoader(string(schemaV1))\n\tdataLoader := gojsonschema.NewGoLoader(cfg)\n\n\tvalidationResult, err := gojsonschema.Validate(schemaLoader, dataLoader)\n\tif err != nil {\n\t\treturn &SchemaError{Field: \"(root)\", Message: err.Error()}\n\t}\n\n\tif !validationResult.Valid() {\n\t\t// Get the most specific error\n\t\terr := getMostSpecificSchemaError(validationResult.Errors())\n\t\treturn err\n\t}\n\n\treturn nil\n}\n\n// validatePredict validates the predict field.\nfunc validatePredict(cfg *configFile, result *ValidationResult) {\n\tif cfg.Predict == nil || *cfg.Predict == \"\" {\n\t\treturn\n\t}\n\n\tpredict := *cfg.Predict\n\tif len(strings.Split(predict, \".py:\")) != 2 {\n\t\tresult.AddError(&ValidationError{\n\t\t\tField:   \"predict\",\n\t\t\tValue:   predict,\n\t\t\tMessage: \"must be in the form 'predict.py:Predictor'\",\n\t\t})\n\t}\n}\n\n// validateTrain validates the train field.\nfunc validateTrain(cfg *configFile, result *ValidationResult) {\n\tif cfg.Train == nil || *cfg.Train == \"\" {\n\t\treturn\n\t}\n\n\ttrain := *cfg.Train\n\tif len(strings.Split(train, \".py:\")) != 2 {\n\t\tresult.AddError(&ValidationError{\n\t\t\tField:   \"train\",\n\t\t\tValue:   train,\n\t\t\tMessage: \"must be in the form 'train.py:Trainer'\",\n\t\t})\n\t}\n}\n\n// validateBuild validates the build configuration.\nfunc validateBuild(cfg *configFile, opts *validateOptions, result *ValidationResult) {\n\tif cfg.Build == nil {\n\t\treturn\n\t}\n\n\tbuild := cfg.Build\n\n\t// Validate Python version is set and valid\n\tif build.PythonVersion == nil || *build.PythonVersion == \"\" {\n\t\tresult.AddError(&ValidationError{\n\t\t\tField:   \"build.python_version\",\n\t\t\tMessage: \"python_version is required. Add it to the build section of your cog.yaml, e.g. `python_version: \\\"3.13\\\"`\",\n\t\t})\n\t} else {\n\t\tif err := validatePythonVersion(*build.PythonVersion); err != nil {\n\t\t\tresult.AddError(err)\n\t\t}\n\t}\n\n\t// Validate mutual exclusivity of python_packages and python_requirements\n\tif len(build.PythonPackages) > 0 && build.PythonRequirements != nil && *build.PythonRequirements != \"\" {\n\t\tresult.AddError(&ValidationError{\n\t\t\tField:   \"build\",\n\t\t\tMessage: \"only one of python_packages or python_requirements can be set, not both\",\n\t\t})\n\t}\n\n\t// Validate python_requirements file exists\n\tif build.PythonRequirements != nil && *build.PythonRequirements != \"\" {\n\t\tif err := validateRequirementsFile(*build.PythonRequirements, opts); err != nil {\n\t\t\tresult.AddError(err)\n\t\t}\n\t}\n\n\t// Validate CUDA version if specified\n\tif build.CUDA != nil && *build.CUDA != \"\" {\n\t\tif err := validateCUDAVersion(*build.CUDA); err != nil {\n\t\t\tresult.AddError(err)\n\t\t}\n\t}\n\n\t// Validate GPU-specific settings\n\tif build.GetGPU() {\n\t\tvalidateGPUConfig(cfg, opts, result)\n\t}\n}\n\n// validatePythonVersion validates the Python version string.\nfunc validatePythonVersion(version string) error {\n\tversion = strings.TrimSpace(version)\n\tparts := strings.SplitN(version, \".\", 3)\n\tif len(parts) < 2 {\n\t\treturn &ValidationError{\n\t\t\tField:   \"build.python_version\",\n\t\t\tValue:   version,\n\t\t\tMessage: \"must include major and minor version (e.g., '3.11')\",\n\t\t}\n\t}\n\n\tmajor, err := strconv.Atoi(parts[0])\n\tif err != nil {\n\t\treturn &ValidationError{\n\t\t\tField:   \"build.python_version\",\n\t\t\tValue:   version,\n\t\t\tMessage: \"invalid major version number\",\n\t\t}\n\t}\n\n\tminor, err := strconv.Atoi(parts[1])\n\tif err != nil {\n\t\treturn &ValidationError{\n\t\t\tField:   \"build.python_version\",\n\t\t\tValue:   version,\n\t\t\tMessage: \"invalid minor version number\",\n\t\t}\n\t}\n\n\tif major < MinimumMajorPythonVersion || (major == MinimumMajorPythonVersion && minor < MinimumMinorPythonVersion) {\n\t\treturn &ValidationError{\n\t\t\tField:   \"build.python_version\",\n\t\t\tValue:   version,\n\t\t\tMessage: fmt.Sprintf(\"minimum supported Python version is %d.%d\", MinimumMajorPythonVersion, MinimumMinorPythonVersion),\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// validateCUDAVersion validates the CUDA version string.\nfunc validateCUDAVersion(cudaVersion string) error {\n\tparts := strings.Split(cudaVersion, \".\")\n\tif len(parts) < 2 {\n\t\treturn &ValidationError{\n\t\t\tField:   \"build.cuda\",\n\t\t\tValue:   cudaVersion,\n\t\t\tMessage: \"must include both major and minor versions (e.g., '11.8')\",\n\t\t}\n\t}\n\n\tmajor, err := strconv.Atoi(parts[0])\n\tif err != nil {\n\t\treturn &ValidationError{\n\t\t\tField:   \"build.cuda\",\n\t\t\tValue:   cudaVersion,\n\t\t\tMessage: \"invalid major version number\",\n\t\t}\n\t}\n\n\tif major < MinimumMajorCudaVersion {\n\t\treturn &ValidationError{\n\t\t\tField:   \"build.cuda\",\n\t\t\tValue:   cudaVersion,\n\t\t\tMessage: fmt.Sprintf(\"minimum supported CUDA version is %d\", MinimumMajorCudaVersion),\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// validateRequirementsFile validates that the requirements file exists and is readable.\nfunc validateRequirementsFile(reqPath string, opts *validateOptions) error {\n\tfullPath := reqPath\n\tif !strings.HasPrefix(reqPath, \"/\") && opts.projectDir != \"\" {\n\t\tfullPath = filepath.Join(opts.projectDir, reqPath)\n\t}\n\n\tif opts.requirementsFS != nil {\n\t\t_, err := fs.ReadFile(opts.requirementsFS, reqPath)\n\t\tif err != nil {\n\t\t\treturn &ValidationError{\n\t\t\t\tField:   \"build.python_requirements\",\n\t\t\t\tValue:   reqPath,\n\t\t\t\tMessage: fmt.Sprintf(\"cannot read file: %v\", err),\n\t\t\t}\n\t\t}\n\t\treturn nil\n\t}\n\n\t// Use the real filesystem\n\tif _, err := os.Stat(fullPath); os.IsNotExist(err) {\n\t\treturn &ValidationError{\n\t\t\tField:   \"build.python_requirements\",\n\t\t\tValue:   reqPath,\n\t\t\tMessage: \"file does not exist\",\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// validateGPUConfig validates GPU-specific configuration like CUDA/CuDNN compatibility.\nfunc validateGPUConfig(cfg *configFile, opts *validateOptions, result *ValidationResult) {\n\tbuild := cfg.Build\n\tif build == nil {\n\t\treturn\n\t}\n\n\t// If both CUDA and CuDNN are specified, check compatibility\n\tif build.CUDA != nil && *build.CUDA != \"\" && build.CuDNN != nil && *build.CuDNN != \"\" {\n\t\tcuda := *build.CUDA\n\t\tcudnn := *build.CuDNN\n\t\tcompatibleCuDNNs := compatibleCuDNNsForCUDA(cuda)\n\t\tfound := slices.Contains(compatibleCuDNNs, cudnn)\n\t\tif !found && len(compatibleCuDNNs) > 0 {\n\t\t\tresult.AddError(&CompatibilityError{\n\t\t\t\tComponent1: \"CUDA\",\n\t\t\t\tVersion1:   cuda,\n\t\t\t\tComponent2: \"CuDNN\",\n\t\t\t\tVersion2:   cudnn,\n\t\t\t\tMessage:    fmt.Sprintf(\"compatible CuDNN versions are: %s\", strings.Join(compatibleCuDNNs, \", \")),\n\t\t\t})\n\t\t}\n\t}\n\n\t// Validate torch/tensorflow requirements if we can read them\n\tif build.PythonRequirements != nil && *build.PythonRequirements != \"\" {\n\t\treqs := loadRequirementsForValidation(*build.PythonRequirements, opts)\n\t\tif len(reqs) > 0 {\n\t\t\tvalidateFrameworkCompatibility(cfg, reqs, result)\n\t\t}\n\t} else if len(build.PythonPackages) > 0 {\n\t\tvalidateFrameworkCompatibility(cfg, build.PythonPackages, result)\n\t}\n}\n\n// loadRequirementsForValidation loads requirements file contents for validation.\nfunc loadRequirementsForValidation(reqPath string, opts *validateOptions) []string {\n\tfullPath := reqPath\n\tif !strings.HasPrefix(reqPath, \"/\") && opts.projectDir != \"\" {\n\t\tfullPath = filepath.Join(opts.projectDir, reqPath)\n\t}\n\n\tif opts.requirementsFS != nil {\n\t\tdata, err := fs.ReadFile(opts.requirementsFS, reqPath)\n\t\tif err != nil {\n\t\t\treturn nil\n\t\t}\n\t\treturn parseRequirementsContent(string(data))\n\t}\n\n\treqs, err := requirements.ReadRequirements(fullPath)\n\tif err != nil {\n\t\treturn nil\n\t}\n\treturn reqs\n}\n\n// parseRequirementsContent parses requirements.txt content into lines.\nfunc parseRequirementsContent(content string) []string {\n\tlines := strings.Split(content, \"\\n\")\n\tresult := make([]string, 0, len(lines))\n\tfor _, line := range lines {\n\t\tline = strings.TrimSpace(line)\n\t\tif line == \"\" || strings.HasPrefix(line, \"#\") || strings.HasPrefix(line, \"-\") {\n\t\t\tcontinue\n\t\t}\n\t\tresult = append(result, line)\n\t}\n\treturn result\n}\n\n// validateFrameworkCompatibility checks torch/tensorflow compatibility with CUDA.\nfunc validateFrameworkCompatibility(cfg *configFile, reqs []string, result *ValidationResult) {\n\t// This is a simplified version - the full logic is in Complete()\n\t// Here we just check for obvious errors.\n\t// Note: torch compatibility is checked in Complete() where it can emit warnings.\n\t// We only validate TensorFlow here since it has stricter requirements.\n\n\tbuild := cfg.Build\n\tif build == nil {\n\t\treturn\n\t}\n\n\ttfVersion := findPackageVersion(reqs, \"tensorflow\")\n\n\t// If CUDA is specified, check TensorFlow compatibility\n\tif build.CUDA != nil && *build.CUDA != \"\" {\n\t\tcuda := *build.CUDA\n\n\t\tif tfVersion != \"\" {\n\t\t\ttfCUDA, _, _ := cudaFromTF(tfVersion)\n\t\t\tif tfCUDA != \"\" && !strings.HasPrefix(cuda, strings.Split(tfCUDA, \".\")[0]) {\n\t\t\t\tresult.AddError(&CompatibilityError{\n\t\t\t\t\tComponent1: \"TensorFlow\",\n\t\t\t\t\tVersion1:   tfVersion,\n\t\t\t\t\tComponent2: \"CUDA\",\n\t\t\t\t\tVersion2:   cuda,\n\t\t\t\t\tMessage:    fmt.Sprintf(\"TensorFlow %s requires CUDA %s\", tfVersion, tfCUDA),\n\t\t\t\t})\n\t\t\t}\n\t\t}\n\t}\n}\n\n// findPackageVersion finds a package version in requirements.\nfunc findPackageVersion(reqs []string, name string) string {\n\tfor _, req := range reqs {\n\t\tpkgName := requirements.PackageName(req)\n\t\tif pkgName == name {\n\t\t\tversions := requirements.Versions(req)\n\t\t\tif len(versions) > 0 {\n\t\t\t\treturn versions[0]\n\t\t\t}\n\t\t}\n\t}\n\treturn \"\"\n}\n\n// validateEnvironment validates environment variables.\nfunc validateEnvironment(cfg *configFile, result *ValidationResult) {\n\tif len(cfg.Environment) == 0 {\n\t\treturn\n\t}\n\n\t_, err := parseAndValidateEnvironment(cfg.Environment)\n\tif err != nil {\n\t\tresult.AddError(&ValidationError{\n\t\t\tField:   \"environment\",\n\t\t\tMessage: err.Error(),\n\t\t})\n\t}\n}\n\n// validateConcurrency validates concurrency settings.\nfunc validateConcurrency(cfg *configFile, result *ValidationResult) {\n\tif cfg.Concurrency == nil || cfg.Concurrency.Max == nil {\n\t\treturn\n\t}\n\n\tmaxConcurrency := *cfg.Concurrency.Max\n\tif maxConcurrency < 1 {\n\t\tresult.AddError(&ValidationError{\n\t\t\tField:   \"concurrency.max\",\n\t\t\tValue:   fmt.Sprintf(\"%d\", maxConcurrency),\n\t\t\tMessage: \"must be at least 1\",\n\t\t})\n\t}\n\n\t// Check Python version requirement for concurrency\n\tif maxConcurrency > 1 && cfg.Build != nil && cfg.Build.PythonVersion != nil {\n\t\tpyVersion := *cfg.Build.PythonVersion\n\t\tmajor, minor, err := splitPythonVersion(pyVersion)\n\t\tif err == nil {\n\t\t\t// Only check minor version if major version is the minimum (3)\n\t\t\t// For major > 3, any minor version would be acceptable\n\t\t\tif major == MinimumMajorPythonVersion && minor < MinimumMinorPythonVersionForConcurrency {\n\t\t\t\tresult.AddError(&ValidationError{\n\t\t\t\t\tField:   \"concurrency.max\",\n\t\t\t\t\tValue:   fmt.Sprintf(\"%d\", maxConcurrency),\n\t\t\t\t\tMessage: fmt.Sprintf(\"concurrency requires Python %d.%d or higher\", MinimumMajorPythonVersion, MinimumMinorPythonVersionForConcurrency),\n\t\t\t\t})\n\t\t\t}\n\t\t}\n\t}\n}\n\n// checkDeprecatedFields checks for deprecated fields and adds warnings.\nfunc checkDeprecatedFields(cfg *configFile, result *ValidationResult) {\n\tif cfg.Build == nil {\n\t\treturn\n\t}\n\n\tif len(cfg.Build.PythonPackages) > 0 {\n\t\tresult.AddWarning(DeprecationWarning{\n\t\t\tField:       \"build.python_packages\",\n\t\t\tReplacement: \"build.python_requirements\",\n\t\t\tMessage:     \"use a requirements.txt file instead\",\n\t\t})\n\t}\n\n\tif len(cfg.Build.PreInstall) > 0 {\n\t\tresult.AddWarning(DeprecationWarning{\n\t\t\tField:       \"build.pre_install\",\n\t\t\tReplacement: \"build.run\",\n\t\t\tMessage:     \"use build.run commands instead\",\n\t\t})\n\t}\n}\n\n// getMostSpecificSchemaError extracts the most specific error from schema validation.\nfunc getMostSpecificSchemaError(errors []gojsonschema.ResultError) *SchemaError {\n\tif len(errors) == 0 {\n\t\treturn &SchemaError{Field: \"(unknown)\", Message: \"unknown schema error\"}\n\t}\n\n\tmostSpecific := 0\n\tfor i, err := range errors {\n\t\tif schemaErrorSpecificity(err) > schemaErrorSpecificity(errors[mostSpecific]) {\n\t\t\tmostSpecific = i\n\t\t} else if schemaErrorSpecificity(err) == schemaErrorSpecificity(errors[mostSpecific]) {\n\t\t\t// Invalid type errors win in a tie-breaker\n\t\t\tif err.Type() == \"invalid_type\" && errors[mostSpecific].Type() != \"invalid_type\" {\n\t\t\t\tmostSpecific = i\n\t\t\t}\n\t\t}\n\t}\n\n\terr := errors[mostSpecific]\n\tfield := err.Field()\n\tif field == \"(root)\" {\n\t\tfield = \"cog.yaml\"\n\t}\n\n\tmessage := getSchemaErrorDescription(err, errors, mostSpecific)\n\n\treturn &SchemaError{\n\t\tField:   field,\n\t\tMessage: message,\n\t}\n}\n\n// getSchemaErrorDescription generates a human-readable description for a schema error.\nfunc getSchemaErrorDescription(err gojsonschema.ResultError, allErrors []gojsonschema.ResultError, index int) string {\n\tswitch err.Type() {\n\tcase \"invalid_type\":\n\t\tif expectedType, ok := err.Details()[\"expected\"].(string); ok {\n\t\t\treturn fmt.Sprintf(\"must be a %s\", humanReadableSchemaType(expectedType))\n\t\t}\n\tcase \"number_one_of\", \"number_any_of\":\n\t\tif index+1 < len(allErrors) {\n\t\t\treturn allErrors[index+1].Description()\n\t\t}\n\t}\n\treturn err.Description()\n}\n\n// humanReadableSchemaType converts JSON schema type names to human-readable names.\nfunc humanReadableSchemaType(definition string) string {\n\tif len(definition) > 0 && definition[0] == '[' {\n\t\tallTypes := strings.Split(definition[1:len(definition)-1], \",\")\n\t\tfor i, t := range allTypes {\n\t\t\tallTypes[i] = humanReadableSchemaType(strings.TrimSpace(t))\n\t\t}\n\t\treturn fmt.Sprintf(\"%s or %s\",\n\t\t\tstrings.Join(allTypes[0:len(allTypes)-1], \", \"),\n\t\t\tallTypes[len(allTypes)-1])\n\t}\n\tswitch definition {\n\tcase \"object\":\n\t\treturn \"mapping\"\n\tcase \"array\":\n\t\treturn \"list\"\n\tdefault:\n\t\treturn definition\n\t}\n}\n\n// schemaErrorSpecificity returns how specific a schema error is based on field depth.\nfunc schemaErrorSpecificity(err gojsonschema.ResultError) int {\n\treturn len(strings.Split(err.Field(), \".\"))\n}\n\n// Note: The legacy Validate function is in validator.go for backwards compatibility\n"
  },
  {
    "path": "pkg/config/validate_test.go",
    "content": "package config\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestValidateConfigFile(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tGPU:           ptr(true),\n\t\t\tPythonVersion: ptr(\"3.10\"),\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"tensorflow==2.12.0\",\n\t\t\t\t\"foo==1.0.0\",\n\t\t\t},\n\t\t\tCUDA: ptr(\"11.8\"),\n\t\t},\n\t}\n\tresult := ValidateConfigFile(cfg)\n\trequire.False(t, result.HasErrors(), \"expected no errors, got: %v\", result.Errors)\n}\n\nfunc TestValidateConfigFileSuccess(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tGPU: ptr(true),\n\t\t\tSystemPackages: []string{\n\t\t\t\t\"libgl1-mesa-glx\",\n\t\t\t\t\"libglib2.0-0\",\n\t\t\t},\n\t\t\tPythonVersion: ptr(\"3.10\"),\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==1.8.1\",\n\t\t\t},\n\t\t},\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.False(t, result.HasErrors(), \"expected no errors, got: %v\", result.Errors)\n}\n\nfunc TestValidateConfigFilePythonVersionNumerical(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tGPU: ptr(true),\n\t\t\tSystemPackages: []string{\n\t\t\t\t\"libgl1-mesa-glx\",\n\t\t\t\t\"libglib2.0-0\",\n\t\t\t},\n\t\t\tPythonVersion: ptr(\"3.10\"),\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==1.8.1\",\n\t\t\t},\n\t\t},\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.False(t, result.HasErrors(), \"expected no errors, got: %v\", result.Errors)\n}\n\nfunc TestValidateConfigFileNullListsAllowed(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tGPU:            ptr(true),\n\t\t\tPythonVersion:  ptr(\"3.10\"),\n\t\t\tSystemPackages: nil,\n\t\t\tPythonPackages: nil,\n\t\t\tRun:            nil,\n\t\t},\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.False(t, result.HasErrors(), \"expected no errors, got: %v\", result.Errors)\n}\n\nfunc TestValidateConfigFilePredictFormat(t *testing.T) {\n\t// Valid predict format\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tPythonVersion: ptr(\"3.10\"),\n\t\t},\n\t\tPredict: ptr(\"predict.py:Predictor\"),\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.False(t, result.HasErrors(), \"expected no errors, got: %v\", result.Errors)\n\n\t// Invalid predict format\n\tcfg.Predict = ptr(\"invalid_format\")\n\tresult = ValidateConfigFile(cfg)\n\trequire.True(t, result.HasErrors())\n\trequire.Contains(t, result.Err().Error(), \"predict.py:Predictor\")\n}\n\nfunc TestValidateConfigFileConcurrencyType(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tGPU:           ptr(true),\n\t\t\tCUDA:          ptr(\"11.8\"),\n\t\t\tPythonVersion: ptr(\"3.11\"),\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==2.0.1\",\n\t\t\t},\n\t\t},\n\t\tPredict: ptr(\"predict.py:Predictor\"),\n\t\tConcurrency: &concurrencyFile{\n\t\t\tMax: ptr(5),\n\t\t},\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.False(t, result.HasErrors(), \"expected no errors, got: %v\", result.Errors)\n}\n\nfunc TestValidateConfigFileDeprecatedPythonPackages(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tPythonVersion: ptr(\"3.10\"),\n\t\t\tPythonPackages: []string{\n\t\t\t\t\"torch==1.8.1\",\n\t\t\t},\n\t\t},\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.False(t, result.HasErrors())\n\trequire.Len(t, result.Warnings, 1)\n\trequire.Contains(t, result.Warnings[0].Message, \"requirements.txt\")\n}\n\nfunc TestValidateConfigFileDeprecatedPreInstall(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tPythonVersion: ptr(\"3.10\"),\n\t\t\tPreInstall: []string{\n\t\t\t\t\"echo hello\",\n\t\t\t},\n\t\t},\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.False(t, result.HasErrors())\n\trequire.Len(t, result.Warnings, 1)\n\trequire.Contains(t, result.Warnings[0].Message, \"build.run\")\n}\n\nfunc TestValidateConfigFileMissingPythonVersion(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{\n\t\t\tGPU: ptr(true),\n\t\t},\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.True(t, result.HasErrors())\n\trequire.Contains(t, result.Err().Error(), \"python_version is required\")\n}\n\nfunc TestValidateConfigFileMissingPythonVersionEmptyBuild(t *testing.T) {\n\tcfg := &configFile{\n\t\tBuild: &buildFile{},\n\t}\n\n\tresult := ValidateConfigFile(cfg)\n\trequire.True(t, result.HasErrors())\n\trequire.Contains(t, result.Err().Error(), \"python_version is required\")\n}\n\nfunc TestValidateConfigFileNilBuildSkipsPythonVersionCheck(t *testing.T) {\n\tcfg := &configFile{}\n\n\tresult := ValidateConfigFile(cfg)\n\t// No build section at all should not error about python_version\n\trequire.False(t, result.HasErrors(), \"expected no errors for nil build, got: %v\", result.Errors)\n}\n\n// ptr returns a pointer to the given value.\nfunc ptr[T any](v T) *T { return &v }\n"
  },
  {
    "path": "pkg/config/version.go",
    "content": "package config\n\n// ArgumentType represents the type of a run argument.\ntype ArgumentType string\n\nconst (\n\tArgumentTypeString ArgumentType = \"str\"\n\tArgumentTypeInt    ArgumentType = \"int\"\n\tArgumentTypeFloat  ArgumentType = \"float\"\n\tArgumentTypeBool   ArgumentType = \"bool\"\n\tArgumentTypePath   ArgumentType = \"Path\"\n)\n\n// RunArgument describes a single argument for a prediction run.\ntype RunArgument struct {\n\tType    ArgumentType `json:\"type\"`\n\tDefault *string      `json:\"default\"`\n\tMin     *string      `json:\"min\"`\n\tMax     *string      `json:\"max\"`\n\tOptions *[]string    `json:\"options\"`\n\tHelp    *string      `json:\"help\"`\n}\n"
  },
  {
    "path": "pkg/docker/build_secrets.go",
    "content": "package docker\n\nimport (\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"strings\"\n\n\t\"github.com/moby/buildkit/session/secrets\"\n\t\"github.com/moby/buildkit/session/secrets/secretsprovider\"\n\t\"github.com/pkg/errors\"\n\t\"github.com/tonistiigi/go-csvvalue\"\n)\n\nfunc ParseSecretsFromHost(workingDir string, secrets []string) (secrets.SecretStore, error) {\n\tsources := make([]secretsprovider.Source, 0, len(secrets))\n\n\tfor _, secret := range secrets {\n\t\tsrc, err := parseSecretFromHost(workingDir, secret)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tsources = append(sources, *src)\n\t}\n\n\treturn secretsprovider.NewStore(sources)\n}\n\nfunc parseSecretFromHost(workingDir, secret string) (*secretsprovider.Source, error) {\n\tfields, err := csvvalue.Fields(secret, nil)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to parse csv secret: %w\", err)\n\t}\n\n\tsrc := secretsprovider.Source{}\n\n\tvar typ string\n\tfor _, field := range fields {\n\t\tkey, value, ok := strings.Cut(field, \"=\")\n\t\tif !ok {\n\t\t\treturn nil, errors.Errorf(\"invalid field %q must be a key=value pair\", field)\n\t\t}\n\t\tkey = strings.ToLower(key)\n\t\tswitch key {\n\t\tcase \"type\":\n\t\t\tif value != \"file\" && value != \"env\" {\n\t\t\t\treturn nil, errors.Errorf(\"unsupported secret type %q\", value)\n\t\t\t}\n\t\t\ttyp = value\n\t\tcase \"id\":\n\t\t\tsrc.ID = value\n\t\tcase \"source\", \"src\":\n\t\t\tif !filepath.IsAbs(value) {\n\t\t\t\tvalue = filepath.Join(workingDir, value)\n\t\t\t\tvalue, err = filepath.Abs(value)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn nil, fmt.Errorf(\"failed to get absolute path for %q: %w\", value, err)\n\t\t\t\t}\n\t\t\t}\n\t\t\tsrc.FilePath = value\n\t\tcase \"env\":\n\t\t\tsrc.Env = value\n\t\tdefault:\n\t\t\treturn nil, errors.Errorf(\"unexpected key '%s' in '%s'\", key, field)\n\t\t}\n\t}\n\tif typ == \"env\" && src.Env == \"\" {\n\t\tsrc.Env = src.FilePath\n\t\tsrc.FilePath = \"\"\n\t}\n\treturn &src, nil\n}\n"
  },
  {
    "path": "pkg/docker/buildkit.go",
    "content": "package docker\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"sync\"\n\n\t\"github.com/docker/docker/api/types/registry\"\n\tbuildkitclient \"github.com/moby/buildkit/client\"\n\t\"github.com/moby/buildkit/session\"\n\t\"github.com/moby/buildkit/session/auth\"\n\t\"github.com/moby/buildkit/session/secrets/secretsprovider\"\n\t\"github.com/moby/buildkit/util/progress/progressui\"\n\t\"google.golang.org/grpc\"\n\t\"google.golang.org/grpc/codes\"\n\t\"google.golang.org/grpc/status\"\n\n\tcogconfig \"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc prepareDockerfileDir(buildDir string, dockerfileContents string) (string, error) {\n\tdockerfilePath := filepath.Join(buildDir, \"Dockerfile\")\n\terr := os.WriteFile(dockerfilePath, []byte(dockerfileContents), 0o644)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn dockerfilePath, nil\n}\n\nfunc solveOptFromImageOptions(buildDir string, opts command.ImageBuildOptions) (buildkitclient.SolveOpt, error) {\n\tdockerfilePath, err := prepareDockerfileDir(buildDir, opts.DockerfileContents)\n\tif err != nil {\n\t\treturn buildkitclient.SolveOpt{}, err\n\t}\n\n\t// first, configure the frontend, in this case, dockerfile.v0\n\tfrontendAttrs := map[string]string{\n\t\t// filename is the path to the Dockerfile within the \"dockerfile\" LocalDir context\n\t\t\"filename\": filepath.Base(dockerfilePath),\n\t\t\"syntax\":   \"docker/dockerfile:1\",\n\t\t// TODO[md]: support multi-stage target\n\t\t// target is the name of a stage in a multi-stage Dockerfile\n\t\t// \"target\": opts.Target,\n\t\t// Replicate only supports linux/amd64, but local Docker Engine could be running on ARM,\n\t\t// including Apple Silicon. Force it to linux/amd64 for now.\n\t\t\"platform\": \"linux/amd64\",\n\t}\n\n\t// disable cache if requested\n\tif opts.NoCache {\n\t\tfrontendAttrs[\"no-cache\"] = \"\"\n\t}\n\n\t// add labels to the image\n\tfor k, v := range opts.Labels {\n\t\tfrontendAttrs[\"label:\"+k] = v\n\t}\n\n\t// add build args to the image\n\tfor k, v := range opts.BuildArgs {\n\t\tif v == nil {\n\t\t\tcontinue\n\t\t}\n\t\tfrontendAttrs[\"build-arg:\"+k] = *v\n\t}\n\n\t// Add SOURCE_DATE_EPOCH if Epoch is set\n\tif opts.Epoch != nil && *opts.Epoch >= 0 {\n\t\tfrontendAttrs[\"build-arg:SOURCE_DATE_EPOCH\"] = fmt.Sprintf(\"%d\", *opts.Epoch)\n\t}\n\n\t// Use WorkingDir as context if ContextDir is relative to ensure consistency with CLI client\n\tcontextDir := opts.ContextDir\n\tif opts.WorkingDir != \"\" && !filepath.IsAbs(opts.ContextDir) {\n\t\tcontextDir = filepath.Join(opts.WorkingDir, opts.ContextDir)\n\t}\n\n\tlocalDirs := map[string]string{\n\t\t\"dockerfile\": filepath.Dir(dockerfilePath),\n\t\t\"context\":    contextDir,\n\t}\n\n\t// Add user-supplied build contexts, but don't overwrite 'dockerfile' or 'context'\n\tfor name, dir := range opts.BuildContexts {\n\t\tif name == \"dockerfile\" || name == \"context\" {\n\t\t\tconsole.Warnf(\"build context name collision: %q\", name)\n\t\t\tcontinue\n\t\t}\n\t\tlocalDirs[name] = dir\n\t\t// Tell the dockerfile frontend about this build context\n\t\tfrontendAttrs[\"context:\"+name] = \"local:\" + name\n\t}\n\n\t// Set exporter attributes\n\texporterAttrs := map[string]string{\n\t\t\"name\": opts.ImageName,\n\t}\n\n\t// if SOURCE_DATE_EPOCH is present in the build args, tell the frontend to rewrite timestamps\n\tif _, ok := frontendAttrs[\"build-arg:SOURCE_DATE_EPOCH\"]; ok {\n\t\texporterAttrs[\"rewrite-timestamp\"] = \"true\"\n\t}\n\n\tsolveOpts := buildkitclient.SolveOpt{\n\t\tFrontend:      \"dockerfile.v0\",\n\t\tFrontendAttrs: frontendAttrs,\n\t\tLocalDirs:     localDirs,\n\t\t// Docker Engine's worker only supports three exporters.\n\t\t// \"moby\" exporter works best for cog, since we want to keep images in\n\t\t// Docker Engine's image store. The others are exporting images to somewhere else.\n\t\t// https://github.com/moby/moby/blob/v20.10.24/builder/builder-next/worker/worker.go#L221\n\t\tExports: []buildkitclient.ExportEntry{\n\t\t\t{Type: \"moby\", Attrs: exporterAttrs},\n\t\t},\n\t}\n\n\t// add auth provider to the session so the local engine can pull and push images\n\tsolveOpts.Session = append(\n\t\tsolveOpts.Session,\n\t\tnewBuildkitAuthProvider(\"r8.im\"),\n\t)\n\n\t// add secrets to the session\n\tif len(opts.Secrets) > 0 {\n\t\t// TODO[md]: support secrets direct from input in addition to env+file\n\t\tstore, err := ParseSecretsFromHost(opts.WorkingDir, opts.Secrets)\n\t\tif err != nil {\n\t\t\treturn buildkitclient.SolveOpt{}, fmt.Errorf(\"failed to parse secrets: %w\", err)\n\t\t}\n\t\tsolveOpts.Session = append(solveOpts.Session, secretsprovider.NewSecretProvider(store))\n\t}\n\n\t// Set cache imports/exports to match DockerCommand logic\n\t// If cogconfig.BuildXCachePath is set, use local cache; otherwise, use inline\n\tif cogconfig.BuildXCachePath != \"\" {\n\t\tsolveOpts.CacheImports = []buildkitclient.CacheOptionsEntry{\n\t\t\t{Type: \"local\", Attrs: map[string]string{\"src\": cogconfig.BuildXCachePath}},\n\t\t}\n\t\tsolveOpts.CacheExports = []buildkitclient.CacheOptionsEntry{\n\t\t\t{Type: \"local\", Attrs: map[string]string{\"dest\": cogconfig.BuildXCachePath}},\n\t\t}\n\t} else {\n\t\tsolveOpts.CacheExports = []buildkitclient.CacheOptionsEntry{\n\t\t\t{Type: \"inline\"},\n\t\t}\n\t}\n\n\treturn solveOpts, nil\n}\n\nfunc newDisplay(statusCh chan *buildkitclient.SolveStatus, displayMode string) func() error {\n\treturn func() error {\n\t\tdisplay, err := progressui.NewDisplay(\n\t\t\tos.Stderr,\n\t\t\tprogressui.DisplayMode(displayMode),\n\t\t\t// progressui.WithPhase(\"BUILDINGGGGG\"),\n\t\t\t// progressui.WithDesc(\"SOMETEXT\", \"SOMECONSOLE\"),\n\t\t)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\t// UpdateFrom must not use the incoming context.\n\t\t// Canceling this context kills the reader of statusCh which blocks buildkit.Client's Solve() indefinitely.\n\t\t// Solve() closes statusCh at the end and UpdateFrom returns by reading the closed channel.\n\t\t//\n\t\t// See https://github.com/superfly/flyctl/pull/2682 for the context.\n\t\t_, err = display.UpdateFrom(context.Background(), statusCh)\n\t\treturn err\n\t}\n}\n\nfunc newBuildkitAuthProvider(registryHosts ...string) session.Attachable {\n\treturn &buildkitAuthProvider{\n\t\tregistryHosts: sync.OnceValues(func() (map[string]registry.AuthConfig, error) {\n\t\t\treturn loadRegistryAuths(context.Background(), registryHosts...)\n\t\t}),\n\t\t// TODO[md]: here's where we'd set the token from config rather than fetching from the credentials helper\n\t\t// token: token,\n\t}\n}\n\ntype buildkitAuthProvider struct {\n\tregistryHosts func() (map[string]registry.AuthConfig, error)\n}\n\nfunc (ap *buildkitAuthProvider) Register(server *grpc.Server) {\n\tauth.RegisterAuthServer(server, ap)\n}\n\nfunc (ap *buildkitAuthProvider) Credentials(ctx context.Context, req *auth.CredentialsRequest) (*auth.CredentialsResponse, error) {\n\tauths, err := ap.registryHosts()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to load registry auth configs: %w\", err)\n\t}\n\tres := &auth.CredentialsResponse{}\n\tif a, ok := auths[req.Host]; ok {\n\t\tres.Username = a.Username\n\t\tres.Secret = a.Password\n\t}\n\n\treturn res, nil\n}\n\nfunc (ap *buildkitAuthProvider) FetchToken(ctx context.Context, req *auth.FetchTokenRequest) (*auth.FetchTokenResponse, error) {\n\treturn nil, status.Errorf(codes.Unavailable, \"client side tokens disabled\")\n}\n\nfunc (ap *buildkitAuthProvider) GetTokenAuthority(ctx context.Context, req *auth.GetTokenAuthorityRequest) (*auth.GetTokenAuthorityResponse, error) {\n\treturn nil, status.Errorf(codes.Unavailable, \"client side tokens disabled\")\n}\n\nfunc (ap *buildkitAuthProvider) VerifyTokenAuthority(ctx context.Context, req *auth.VerifyTokenAuthorityRequest) (*auth.VerifyTokenAuthorityResponse, error) {\n\treturn nil, status.Errorf(codes.Unavailable, \"client side tokens disabled\")\n}\n"
  },
  {
    "path": "pkg/docker/command/command.go",
    "content": "package command\n\nimport (\n\t\"context\"\n\t\"io\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/docker/api/types/image\"\n)\n\ntype Command interface {\n\t// Pull pulls an image from a remote registry and returns the inspect response for the local image.\n\t// If the image already exists, it will return the inspect response for the local image without pulling.\n\t// When force is true, it will always attempt to pull the image.\n\tPull(ctx context.Context, ref string, force bool) (*image.InspectResponse, error)\n\tPush(ctx context.Context, ref string) error\n\tRemoveImage(ctx context.Context, ref string) error\n\tLoadUserInformation(ctx context.Context, registryHost string) (*UserInfo, error)\n\tInspect(ctx context.Context, ref string) (*image.InspectResponse, error)\n\tImageExists(ctx context.Context, ref string) (bool, error)\n\tContainerLogs(ctx context.Context, containerID string, w io.Writer) error\n\tContainerInspect(ctx context.Context, id string) (*container.InspectResponse, error)\n\tContainerStop(ctx context.Context, containerID string) error\n\n\t// ImageBuild builds an image and returns the image ID (sha256:...) on success.\n\tImageBuild(ctx context.Context, options ImageBuildOptions) (string, error)\n\tRun(ctx context.Context, options RunOptions) error\n\tContainerStart(ctx context.Context, options RunOptions) (string, error)\n\n\t// ImageSave exports a Docker image as a tar stream.\n\t// The caller must close the returned ReadCloser.\n\tImageSave(ctx context.Context, imageRef string) (io.ReadCloser, error)\n}\n\ntype ImageBuildOptions struct {\n\tWorkingDir         string\n\tDockerfileContents string\n\t// TODO[md]: ImageName should be renamed to Tag\n\tImageName string\n\t// Secrets in the format of \"id=foo,src=/path/to/file\" or \"id=kube,env=KUBECONFIG\"\n\t// docs: https://docs.docker.com/build/building/secrets/#use-secrets-in-dockerfile\n\tSecrets        []string\n\tNoCache        bool\n\tProgressOutput string\n\tEpoch          *int64\n\tContextDir     string\n\tBuildContexts  map[string]string\n\tLabels         map[string]string\n\n\t// only supported on buildkit client, not cli client\n\tBuildArgs map[string]*string\n}\n\ntype RunOptions struct {\n\tDetach     bool\n\tArgs       []string\n\tEnv        []string\n\tGPUs       string\n\tImage      string\n\tPorts      []Port\n\tVolumes    []Volume\n\tWorkdir    string\n\tExtraHosts []string\n\tStdin      io.Reader\n\tStdout     io.Writer\n\tStderr     io.Writer\n}\n\ntype Port struct {\n\tHostPort      int\n\tContainerPort int\n}\n\ntype Volume struct {\n\tSource      string\n\tDestination string\n}\n"
  },
  {
    "path": "pkg/docker/command/errors.go",
    "content": "package command\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n)\n\n// NotFoundError represents “object  wasn’t found” inside the Docker engine.\ntype NotFoundError struct {\n\t// Ref is a unique identifier, such as an image reference, container ID, etc.\n\tRef string\n\t// Object is the ref type, such as \"container\", \"image\", \"volume\", etc.\n\tObject string\n}\n\nfunc (e *NotFoundError) Error() string {\n\tobjType := e.Object\n\tif objType == \"\" {\n\t\tobjType = \"object\"\n\t}\n\treturn fmt.Sprintf(\"%s not found: %q\", objType, e.Ref)\n}\n\nfunc (e *NotFoundError) Is(target error) bool {\n\t_, ok := target.(*NotFoundError)\n\treturn ok\n}\n\nfunc IsNotFoundError(err error) bool {\n\treturn errors.Is(err, &NotFoundError{})\n}\n\nvar ErrAuthorizationFailed = errors.New(\"authorization failed\")\n"
  },
  {
    "path": "pkg/docker/command/manifest.go",
    "content": "package command\n\nimport \"github.com/replicate/cog/pkg/global\"\n\ntype Config struct {\n\tLabels map[string]string `json:\"Labels\"`\n\tEnv    []string          `json:\"Env\"`\n}\n\ntype Manifest struct {\n\tConfig Config `json:\"Config\"`\n}\n\nconst (\n\tR8CogVersionEnvVarName    = \"R8_COG_VERSION\"\n\tR8TorchVersionEnvVarName  = \"R8_TORCH_VERSION\"\n\tR8CudaVersionEnvVarName   = \"R8_CUDA_VERSION\"\n\tR8CudnnVersionEnvVarName  = \"R8_CUDNN_VERSION\"\n\tR8PythonVersionEnvVarName = \"R8_PYTHON_VERSION\"\n)\n\nvar (\n\tCogConfigLabelKey          = global.LabelNamespace + \"config\"\n\tCogVersionLabelKey         = global.LabelNamespace + \"version\"\n\tCogOpenAPISchemaLabelKey   = global.LabelNamespace + \"openapi_schema\"\n\tCogWeightsManifestLabelKey = global.LabelNamespace + \"r8_weights_manifest\"\n)\n"
  },
  {
    "path": "pkg/docker/command/user_info.go",
    "content": "package command\n\ntype UserInfo struct {\n\tToken    string\n\tUsername string\n}\n"
  },
  {
    "path": "pkg/docker/credential_helper_input.go",
    "content": "package docker\n\ntype CredentialHelperInput struct {\n\tUsername  string\n\tSecret    string //nolint:gosec // G117: this is a Docker credential, not a hardcoded secret\n\tServerURL string\n}\n"
  },
  {
    "path": "pkg/docker/credentials.go",
    "content": "package docker\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n\t\"os/exec\"\n\t\"strings\"\n\n\t\"github.com/docker/cli/cli/config\"\n\t\"github.com/docker/cli/cli/config/configfile\"\n\t\"github.com/docker/cli/cli/config/types\"\n\t\"github.com/docker/docker/api/types/registry\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc loadUserInformation(ctx context.Context, registryHost string) (*command.UserInfo, error) {\n\tconf := config.LoadDefaultConfigFile(os.Stderr)\n\tcredsStore := conf.CredentialsStore\n\tif credsStore == \"\" {\n\t\tauthConf, err := loadAuthFromConfig(conf, registryHost)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\treturn &command.UserInfo{\n\t\t\tToken:    authConf.Password,\n\t\t\tUsername: authConf.Username,\n\t\t}, nil\n\t}\n\tcredsHelper, err := loadAuthFromCredentialsStore(ctx, credsStore, registryHost)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn &command.UserInfo{\n\t\tToken:    credsHelper.Secret,\n\t\tUsername: credsHelper.Username,\n\t}, nil\n}\n\nfunc loadAuthFromConfig(conf *configfile.ConfigFile, registryHost string) (types.AuthConfig, error) {\n\treturn conf.AuthConfigs[registryHost], nil\n}\n\nfunc loadRegistryAuths(ctx context.Context, registryHosts ...string) (map[string]registry.AuthConfig, error) {\n\tconf := config.LoadDefaultConfigFile(os.Stderr)\n\tout := make(map[string]registry.AuthConfig)\n\n\tfor _, host := range registryHosts {\n\t\t// Try loading auth for the requested host\n\t\tauth, err := tryLoadAuthForHost(ctx, conf, host)\n\t\tif err == nil && auth != nil {\n\t\t\tout[host] = *auth\n\t\t\tcontinue\n\t\t}\n\n\t\t// FALLBACK: If requesting alternate registry and no auth found,\n\t\t// try reusing r8.im credentials\n\t\tif host != global.DefaultReplicateRegistryHost {\n\t\t\tauth, err := tryLoadAuthForHost(ctx, conf, global.DefaultReplicateRegistryHost)\n\t\t\tif err == nil && auth != nil {\n\t\t\t\t// Reuse credentials for the alternate registry\n\t\t\t\tauth.ServerAddress = host // Update to new host\n\t\t\t\tout[host] = *auth\n\t\t\t\tconsole.Infof(\"Using existing %s credentials for %s\", global.DefaultReplicateRegistryHost, host)\n\t\t\t\tcontinue\n\t\t\t}\n\t\t}\n\t}\n\n\treturn out, nil\n}\n\nfunc tryLoadAuthForHost(ctx context.Context, conf *configfile.ConfigFile, host string) (*registry.AuthConfig, error) {\n\t// Try credentials store first (e.g., osxkeychain, pass)\n\tif conf.CredentialsStore != \"\" {\n\t\tcredsHelper, err := loadAuthFromCredentialsStore(ctx, conf.CredentialsStore, host)\n\t\tif err == nil {\n\t\t\treturn ®istry.AuthConfig{\n\t\t\t\tUsername:      credsHelper.Username,\n\t\t\t\tPassword:      credsHelper.Secret,\n\t\t\t\tServerAddress: host,\n\t\t\t}, nil\n\t\t}\n\t}\n\n\t// Fallback to config file\n\tif auth, ok := conf.AuthConfigs[host]; ok {\n\t\treturn ®istry.AuthConfig{\n\t\t\tUsername:      auth.Username,\n\t\t\tPassword:      auth.Password,\n\t\t\tAuth:          auth.Auth,\n\t\t\tServerAddress: host,\n\t\t\tIdentityToken: auth.IdentityToken,\n\t\t\tRegistryToken: auth.RegistryToken,\n\t\t}, nil\n\t}\n\n\treturn nil, fmt.Errorf(\"no credentials found for %s\", host)\n}\n\nfunc loadAuthFromCredentialsStore(ctx context.Context, credsStore string, registryHost string) (*CredentialHelperInput, error) {\n\tvar out strings.Builder\n\tbinary := dockerCredentialBinary(credsStore)\n\tcmd := exec.CommandContext(ctx, binary, \"get\") //nolint:gosec // G702: binary is from Docker config, not user input\n\tcmd.Env = os.Environ()\n\tcmd.Stdout = &out\n\tcmd.Stderr = &out\n\tstdin, err := cmd.StdinPipe()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdefer stdin.Close()\n\tconsole.Debug(\"$ \" + strings.Join(cmd.Args, \" \"))\n\terr = cmd.Start()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\t_, err = io.WriteString(stdin, registryHost)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\terr = stdin.Close()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\terr = cmd.Wait()\n\tif err != nil {\n\t\toutput := strings.TrimSpace(out.String())\n\t\tif output != \"\" {\n\t\t\treturn nil, fmt.Errorf(\"failed to get credentials for %q: %s\", registryHost, output)\n\t\t}\n\t\treturn nil, fmt.Errorf(\"failed to get credentials for %q: %w\", registryHost, err)\n\t}\n\n\tvar config CredentialHelperInput\n\terr = json.Unmarshal([]byte(out.String()), &config)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &config, nil\n}\n\nfunc dockerCredentialBinary(credsStore string) string {\n\treturn \"docker-credential-\" + credsStore\n}\n"
  },
  {
    "path": "pkg/docker/credentials_test.go",
    "content": "package docker\n\nimport (\n\t\"context\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/docker/cli/cli/config/configfile\"\n\t\"github.com/docker/cli/cli/config/types\"\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n)\n\nfunc TestLoadRegistryAuths_Fallback(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"uses credentials for requested host when available\", func(t *testing.T) {\n\t\t// Create a mock config with credentials for the requested host\n\t\tconf := &configfile.ConfigFile{\n\t\t\tAuthConfigs: map[string]types.AuthConfig{\n\t\t\t\t\"registry.example.com\": {\n\t\t\t\t\tUsername: \"user1\",\n\t\t\t\t\tPassword: \"pass1\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tauth, err := tryLoadAuthForHost(ctx, conf, \"registry.example.com\")\n\t\trequire.NoError(t, err)\n\t\trequire.NotNil(t, auth)\n\t\tassert.Equal(t, \"user1\", auth.Username)\n\t\tassert.Equal(t, \"pass1\", auth.Password)\n\t\tassert.Equal(t, \"registry.example.com\", auth.ServerAddress)\n\t})\n\n\tt.Run(\"falls back to default registry credentials when alternate registry has no credentials\", func(t *testing.T) {\n\t\t// Set up a temporary docker config file\n\t\ttmpDir := t.TempDir()\n\t\tdockerConfigPath := filepath.Join(tmpDir, \"config.json\")\n\n\t\t// Create a config file with credentials only for the default registry\n\t\tconf := &configfile.ConfigFile{\n\t\t\tFilename: dockerConfigPath,\n\t\t\tAuthConfigs: map[string]types.AuthConfig{\n\t\t\t\tglobal.DefaultReplicateRegistryHost: {\n\t\t\t\t\tUsername: \"defaultuser\",\n\t\t\t\t\tPassword: \"defaultpass\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\t\trequire.NoError(t, conf.Save())\n\n\t\t// Point Docker to our test config\n\t\tt.Setenv(\"DOCKER_CONFIG\", tmpDir)\n\n\t\t// Try loading auth for an alternate registry that doesn't have credentials\n\t\tauths, err := loadRegistryAuths(ctx, \"registry.example.com\")\n\t\trequire.NoError(t, err)\n\t\trequire.NotNil(t, auths)\n\n\t\t// Should have fallen back to default registry credentials\n\t\tauth, ok := auths[\"registry.example.com\"]\n\t\trequire.True(t, ok, \"should have auth for registry.example.com\")\n\t\tassert.Equal(t, \"defaultuser\", auth.Username)\n\t\tassert.Equal(t, \"defaultpass\", auth.Password)\n\t\tassert.Equal(t, \"registry.example.com\", auth.ServerAddress, \"server address should be updated to the requested host\")\n\t})\n\n\tt.Run(\"does not fallback when requesting default registry\", func(t *testing.T) {\n\t\t// This test uses tryLoadAuthForHost directly to avoid credential store issues\n\t\tconf := &configfile.ConfigFile{\n\t\t\tAuthConfigs: map[string]types.AuthConfig{},\n\t\t}\n\n\t\t// Try loading auth for the default registry\n\t\tauth, err := tryLoadAuthForHost(ctx, conf, global.DefaultReplicateRegistryHost)\n\t\trequire.Error(t, err, \"should error when no credentials found\")\n\t\tassert.Nil(t, auth)\n\t\tassert.Contains(t, err.Error(), \"no credentials found\")\n\t})\n\n\tt.Run(\"prefers direct credentials over fallback\", func(t *testing.T) {\n\t\t// Create a mock config with credentials for both registries\n\t\tconf := &configfile.ConfigFile{\n\t\t\tAuthConfigs: map[string]types.AuthConfig{\n\t\t\t\tglobal.DefaultReplicateRegistryHost: {\n\t\t\t\t\tUsername: \"defaultuser\",\n\t\t\t\t\tPassword: \"defaultpass\",\n\t\t\t\t},\n\t\t\t\t\"registry.example.com\": {\n\t\t\t\t\tUsername: \"directuser\",\n\t\t\t\t\tPassword: \"directpass\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\t// Try loading auth for the alternate registry\n\t\tauth, err := tryLoadAuthForHost(ctx, conf, \"registry.example.com\")\n\t\trequire.NoError(t, err)\n\t\trequire.NotNil(t, auth)\n\n\t\t// Should use direct credentials, not fallback\n\t\tassert.Equal(t, \"directuser\", auth.Username)\n\t\tassert.Equal(t, \"directpass\", auth.Password)\n\t\tassert.Equal(t, \"registry.example.com\", auth.ServerAddress)\n\t})\n\n\tt.Run(\"returns empty map when no credentials available\", func(t *testing.T) {\n\t\t// This test uses tryLoadAuthForHost to avoid credential store issues\n\t\t// The loadRegistryAuths function doesn't error when no credentials are found,\n\t\t// it just returns an empty map\n\t\tconf := &configfile.ConfigFile{\n\t\t\tAuthConfigs: map[string]types.AuthConfig{},\n\t\t}\n\n\t\t// Try loading auth for an alternate registry (will fail)\n\t\tauth1, err := tryLoadAuthForHost(ctx, conf, \"registry.example.com\")\n\t\trequire.Error(t, err)\n\t\tassert.Nil(t, auth1)\n\n\t\t// Try loading auth for default registry (will also fail)\n\t\tauth2, err := tryLoadAuthForHost(ctx, conf, global.DefaultReplicateRegistryHost)\n\t\trequire.Error(t, err)\n\t\tassert.Nil(t, auth2)\n\n\t\t// Since both fail, loadRegistryAuths would return an empty map\n\t\t// (it doesn't error, just silently skips hosts without credentials)\n\t})\n}\n\nfunc TestTryLoadAuthForHost(t *testing.T) {\n\tctx := context.Background()\n\n\tt.Run(\"loads auth from config file\", func(t *testing.T) {\n\t\tconf := &configfile.ConfigFile{\n\t\t\tAuthConfigs: map[string]types.AuthConfig{\n\t\t\t\t\"registry.example.com\": {\n\t\t\t\t\tUsername: \"testuser\",\n\t\t\t\t\tPassword: \"testpass\",\n\t\t\t\t\tAuth:     \"dGVzdHVzZXI6dGVzdHBhc3M=\",\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\tauth, err := tryLoadAuthForHost(ctx, conf, \"registry.example.com\")\n\t\trequire.NoError(t, err)\n\t\trequire.NotNil(t, auth)\n\t\tassert.Equal(t, \"testuser\", auth.Username)\n\t\tassert.Equal(t, \"testpass\", auth.Password)\n\t\tassert.Equal(t, \"dGVzdHVzZXI6dGVzdHBhc3M=\", auth.Auth)\n\t\tassert.Equal(t, \"registry.example.com\", auth.ServerAddress)\n\t})\n\n\tt.Run(\"returns error when no auth found\", func(t *testing.T) {\n\t\tconf := &configfile.ConfigFile{\n\t\t\tAuthConfigs: map[string]types.AuthConfig{},\n\t\t}\n\n\t\tauth, err := tryLoadAuthForHost(ctx, conf, \"registry.example.com\")\n\t\trequire.Error(t, err)\n\t\tassert.Nil(t, auth)\n\t\tassert.Contains(t, err.Error(), \"no credentials found\")\n\t})\n}\n"
  },
  {
    "path": "pkg/docker/docker.go",
    "content": "package docker\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"net\"\n\t\"os\"\n\t\"strconv\"\n\t\"strings\"\n\n\t\"github.com/containerd/errdefs\"\n\t\"github.com/docker/docker/api/types\"\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/docker/api/types/image\"\n\t\"github.com/docker/docker/api/types/network\"\n\t\"github.com/docker/docker/api/types/registry\"\n\t\"github.com/docker/docker/client\"\n\t\"github.com/docker/docker/pkg/jsonmessage\"\n\t\"github.com/docker/docker/pkg/stdcopy\"\n\t\"github.com/docker/go-connections/nat\"\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\t\"github.com/mattn/go-isatty\"\n\tbuildkitclient \"github.com/moby/buildkit/client\"\n\t\"github.com/moby/buildkit/exporter/containerimage/exptypes\"\n\t\"github.com/moby/term\"\n\tocispec \"github.com/opencontainers/image-spec/specs-go/v1\"\n\t\"golang.org/x/sync/errgroup\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc ptrVal[T any](v T) *T { return &v }\n\nfunc NewClient(ctx context.Context, opts ...Option) (*apiClient, error) {\n\tclientOptions := &clientOptions{\n\t\tauthConfigs: make(map[string]registry.AuthConfig),\n\t}\n\tfor _, opt := range opts {\n\t\topt(clientOptions)\n\t}\n\n\tif clientOptions.host == \"\" {\n\t\thost, err := determineDockerHost()\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"error determining docker host: %w\", err)\n\t\t}\n\t\tclientOptions.host = host\n\t}\n\n\t// TODO[md]: we create a client at the top of each cli invocation, the sdk client hits an api which\n\t// adds (a tiny biy of) overead. swap this with a handle that'll lazily initialize a client and ping for health.\n\t// ditto for fetching registry credentials.\n\n\tdockerClientOpts := []client.Opt{\n\t\tclient.WithTLSClientConfigFromEnv(),\n\t\tclient.WithVersionFromEnv(),\n\t\tclient.WithAPIVersionNegotiation(),\n\t\tclient.WithHost(clientOptions.host),\n\t}\n\n\tclient, err := client.NewClientWithOpts(dockerClientOpts...)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"error creating docker client: %w\", err)\n\t}\n\n\tif _, err := client.Ping(ctx); err != nil {\n\t\treturn nil, fmt.Errorf(\"error pinging docker daemon: %w\", err)\n\t}\n\n\t// Load authentication for configured registry and any other registries that might be needed\n\tauthConfig, err := loadRegistryAuths(ctx, global.ReplicateRegistryHost)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"error loading user information: %w, you may need to authenticate using cog login\", err)\n\t}\n\n\t// Add any additional auth configs passed via options\n\tfor _, opt := range clientOptions.authConfigs {\n\t\tauthConfig[opt.ServerAddress] = opt\n\t}\n\n\treturn &apiClient{client, authConfig}, nil\n}\n\ntype apiClient struct {\n\tclient     *client.Client\n\tauthConfig map[string]registry.AuthConfig\n}\n\nfunc (c *apiClient) Pull(ctx context.Context, imageRef string, force bool) (*image.InspectResponse, error) {\n\tconsole.Debugf(\"=== APIClient.Pull %s force:%t\", imageRef, force)\n\n\tif !force {\n\t\tinspect, err := c.Inspect(ctx, imageRef)\n\t\tif err == nil {\n\t\t\treturn inspect, nil\n\t\t} else if !command.IsNotFoundError(err) {\n\t\t\t// Log a warning if inspect fails for any reason other than not found.\n\t\t\t// It's likely that pull will fail as well, but it's better to return that error\n\t\t\t// so the caller can handle it appropriately than to fail silently here.\n\t\t\tconsole.Warnf(\"failed to inspect image before pulling %q: %s\", imageRef, err)\n\t\t}\n\t}\n\n\toutput, err := c.client.ImagePull(ctx, imageRef, image.PullOptions{\n\t\t// force image to linux/amd64 to match production\n\t\tPlatform: \"linux/amd64\",\n\t})\n\tif err != nil {\n\t\tif errdefs.IsNotFound(err) {\n\t\t\treturn nil, &command.NotFoundError{Ref: imageRef, Object: \"image\"}\n\t\t}\n\t\treturn nil, fmt.Errorf(\"failed to pull image %q: %w\", imageRef, err)\n\t}\n\tdefer output.Close()\n\t_, err = io.Copy(os.Stderr, output)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to copy pull output: %w\", err)\n\t}\n\n\t// pull succeeded, inspect the image again and return\n\tinspect, err := c.Inspect(ctx, imageRef)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to inspect image after pulling %q: %w\", imageRef, err)\n\t}\n\treturn inspect, nil\n}\n\nfunc (c *apiClient) ContainerStop(ctx context.Context, containerID string) error {\n\tconsole.Debugf(\"=== APIClient.ContainerStop %s\", containerID)\n\n\terr := c.client.ContainerStop(ctx, containerID, container.StopOptions{\n\t\tTimeout: ptrVal(3),\n\t})\n\tif err != nil {\n\t\tif errdefs.IsNotFound(err) {\n\t\t\treturn &command.NotFoundError{Ref: containerID, Object: \"container\"}\n\t\t}\n\t\treturn fmt.Errorf(\"failed to stop container %q: %w\", containerID, err)\n\t}\n\treturn nil\n}\n\nfunc (c *apiClient) ContainerInspect(ctx context.Context, containerID string) (*container.InspectResponse, error) {\n\tconsole.Debugf(\"=== APIClient.ContainerInspect %s\", containerID)\n\n\tresp, err := c.client.ContainerInspect(ctx, containerID)\n\tif err != nil {\n\t\tif errdefs.IsNotFound(err) {\n\t\t\treturn nil, &command.NotFoundError{Ref: containerID, Object: \"container\"}\n\t\t}\n\t\treturn nil, fmt.Errorf(\"failed to inspect container %q: %w\", containerID, err)\n\t}\n\treturn &resp, nil\n}\n\nfunc (c *apiClient) ContainerLogs(ctx context.Context, containerID string, w io.Writer) error {\n\tconsole.Debugf(\"=== APIClient.ContainerLogs %s\", containerID)\n\n\t// First inspect the container to check if it has TTY enabled\n\tinspect, err := c.ContainerInspect(ctx, containerID)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to inspect container %q: %w\", containerID, err)\n\t}\n\n\tlogs, err := c.client.ContainerLogs(ctx, containerID, container.LogsOptions{\n\t\tShowStdout: true,\n\t\tShowStderr: true,\n\t\tFollow:     true,\n\t})\n\tif err != nil {\n\t\tif errdefs.IsNotFound(err) {\n\t\t\treturn &command.NotFoundError{Ref: containerID, Object: \"container\"}\n\t\t}\n\t\treturn fmt.Errorf(\"failed to get container logs for %q: %w\", containerID, err)\n\t}\n\tdefer logs.Close()\n\n\t// If TTY is enabled, we can just copy the logs directly\n\tif inspect.Config.Tty {\n\t\tif _, err = io.Copy(w, logs); err != nil {\n\t\t\treturn fmt.Errorf(\"failed to copy logs: %w\", err)\n\t\t}\n\t\treturn nil\n\t}\n\n\t// For non-TTY containers, use StdCopy to demultiplex stdout and stderr\n\tif _, err = stdcopy.StdCopy(w, w, logs); err != nil {\n\t\treturn fmt.Errorf(\"failed to copy logs: %w\", err)\n\t}\n\treturn nil\n}\n\nfunc (c *apiClient) Push(ctx context.Context, imageRef string) error {\n\tconsole.Debugf(\"=== APIClient.Push %s\", imageRef)\n\n\tparsedName, err := name.ParseReference(imageRef)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to parse image reference: %w\", err)\n\t}\n\n\tconsole.Debugf(\"fully qualified image ref: %s\", parsedName)\n\n\t// eagerly set auth config, or do it async\n\tvar authConfig registry.AuthConfig\n\tregistryHost := parsedName.Context().RegistryStr()\n\tif auth, ok := c.authConfig[registryHost]; ok {\n\t\tauthConfig = auth\n\t} else {\n\t\t// Dynamically load authentication for this registry if not already loaded\n\t\tauthConfigs, err := loadRegistryAuths(ctx, registryHost)\n\t\tif err == nil {\n\t\t\tif auth, ok := authConfigs[registryHost]; ok {\n\t\t\t\tauthConfig = auth\n\t\t\t\t// Cache the auth config for future use\n\t\t\t\tc.authConfig[registryHost] = auth\n\t\t\t}\n\t\t}\n\t}\n\n\tvar opts image.PushOptions\n\tencodedAuth, err := registry.EncodeAuthConfig(authConfig)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to encode auth config: %w\", err)\n\t}\n\topts.RegistryAuth = encodedAuth\n\n\toutput, err := c.client.ImagePush(ctx, imageRef, opts)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to push image: %w\", err)\n\t}\n\tdefer output.Close()\n\n\t// output is a json stream, so we need to parse it, handle errors, and write progress to stderr\n\tisTTY := console.IsTTY(os.Stderr)\n\tif err := jsonmessage.DisplayJSONMessagesStream(output, os.Stderr, os.Stderr.Fd(), isTTY, nil); err != nil {\n\t\tvar streamErr *jsonmessage.JSONError\n\t\tif errors.As(err, &streamErr) {\n\t\t\tif isTagNotFoundError(err) {\n\t\t\t\treturn &command.NotFoundError{Ref: imageRef, Object: \"tag\"}\n\t\t\t}\n\t\t\tif isRepositoryNotFoundError(err) {\n\t\t\t\treturn &command.NotFoundError{Ref: imageRef, Object: \"repository\"}\n\t\t\t}\n\t\t\tif isAuthorizationFailedError(err) {\n\t\t\t\treturn command.ErrAuthorizationFailed\n\t\t\t}\n\t\t}\n\t\treturn fmt.Errorf(\"error during image push: %w\", err)\n\t}\n\n\treturn nil\n}\n\nfunc (c *apiClient) ImageSave(ctx context.Context, imageRef string) (io.ReadCloser, error) {\n\tconsole.Debugf(\"=== APIClient.ImageSave %s\", imageRef)\n\treturn c.client.ImageSave(ctx, []string{imageRef})\n}\n\n// TODO[md]: this doesn't need to be on the interface, move to auth handler\nfunc (c *apiClient) LoadUserInformation(ctx context.Context, registryHost string) (*command.UserInfo, error) {\n\tconsole.Debugf(\"=== APIClient.LoadUserInformation %s\", registryHost)\n\n\treturn loadUserInformation(ctx, registryHost)\n}\n\nfunc (c *apiClient) Inspect(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\tconsole.Debugf(\"=== APIClient.Inspect %s\", ref)\n\n\t// TODO[md]: platform requires engine 1.49+, and it's not widly available as of 2025-05.\n\t// \tplatform := ocispec.Platform{OS: \"linux\", Architecture: \"amd64\"}\n\t//  client.ImageInspectWithPlatform(&platform),\n\tinspect, err := c.client.ImageInspect(ctx, ref)\n\n\tif err != nil {\n\t\tif errdefs.IsNotFound(err) {\n\t\t\treturn nil, &command.NotFoundError{Ref: ref, Object: \"image\"}\n\t\t}\n\t\treturn nil, fmt.Errorf(\"error inspecting image: %w\", err)\n\t}\n\n\treturn &inspect, nil\n}\n\nfunc (c *apiClient) RemoveImage(ctx context.Context, ref string) error {\n\tconsole.Debugf(\"=== APIClient.RemoveImage %s\", ref)\n\n\tresp, err := c.client.ImageRemove(ctx, ref, image.RemoveOptions{})\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tif len(resp) == 0 {\n\t\treturn &command.NotFoundError{Ref: ref, Object: \"image\"}\n\t}\n\treturn nil\n}\n\nfunc (c *apiClient) ImageExists(ctx context.Context, ref string) (bool, error) {\n\tconsole.Debugf(\"=== APIClient.ImageExists %s\", ref)\n\n\t_, err := c.Inspect(ctx, ref)\n\tif err != nil {\n\t\tif command.IsNotFoundError(err) {\n\t\t\treturn false, nil\n\t\t}\n\t\treturn false, err\n\t}\n\treturn true, nil\n}\n\nfunc (c *apiClient) ImageBuild(ctx context.Context, options command.ImageBuildOptions) (string, error) {\n\tconsole.Debugf(\"=== APIClient.ImageBuild %s\", options.ImageName)\n\n\tbuildDir, err := os.MkdirTemp(\"\", \"cog-build\")\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tdefer os.RemoveAll(buildDir)\n\n\tbc, err := buildkitclient.New(ctx, \"\",\n\t\t// Connect to Docker Engine's embedded Buildkit.\n\t\tbuildkitclient.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {\n\t\t\treturn c.client.DialHijack(ctx, \"/grpc\", \"h2c\", map[string][]string{})\n\t\t}),\n\t)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tstatusCh := make(chan *buildkitclient.SolveStatus)\n\tvar res *buildkitclient.SolveResponse\n\n\t// Determine display mode: options.ProgressOutput > env > 'auto'\n\tdisplayMode := options.ProgressOutput\n\tif displayMode == \"\" {\n\t\tdisplayMode = os.Getenv(\"BUILDKIT_PROGRESS\")\n\t}\n\tif displayMode == \"\" {\n\t\tdisplayMode = \"auto\"\n\t}\n\n\t// Build the image.\n\teg, ctx := errgroup.WithContext(ctx)\n\n\t// run the build in a goroutine\n\teg.Go(func() error {\n\t\toptions, err := solveOptFromImageOptions(buildDir, options)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\t// run the display in a goroutine _after_ we've built SolveOpt\n\t\teg.Go(newDisplay(statusCh, displayMode))\n\n\t\tres, err = bc.Solve(ctx, nil, options, statusCh)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\treturn nil\n\t})\n\terr = eg.Wait()\n\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\timageID := res.ExporterResponse[exptypes.ExporterImageDigestKey]\n\tif imageID == \"\" {\n\t\treturn \"\", fmt.Errorf(\"buildkit did not return an image digest\")\n\t}\n\tconsole.Debugf(\"image digest %s\", imageID)\n\n\treturn imageID, nil\n}\n\nfunc (c *apiClient) containerRun(ctx context.Context, options command.RunOptions) (string, error) {\n\tconsole.Debugf(\"=== APIClient.containerRun %s\", options.Image)\n\n\tvar attachStdin, tty, attachStderr, attachStdout bool\n\tif !options.Detach {\n\t\t// Determine if we should attach stdin (file, pipe, interactive stdin, etc)\n\t\tattachStdin, tty = shouldAttachStdin(options.Stdin)\n\t\tattachStdout = options.Stdout != nil\n\t\tattachStderr = options.Stderr != nil\n\t}\n\n\tcontainerCfg := &container.Config{\n\t\tImage:        options.Image,\n\t\tCmd:          options.Args,\n\t\tEnv:          options.Env,\n\t\tAttachStdin:  attachStdin,\n\t\tAttachStdout: attachStdout,\n\t\tAttachStderr: attachStderr,\n\t\tOpenStdin:    attachStdin,\n\t\tStdinOnce:    attachStdin,\n\t\tTty:          tty,\n\t}\n\n\t// Set working directory if specified\n\tif options.Workdir != \"\" {\n\t\tcontainerCfg.WorkingDir = options.Workdir\n\t}\n\n\tif len(options.Ports) > 0 {\n\t\tcontainerCfg.ExposedPorts = make(nat.PortSet)\n\t\tfor _, port := range options.Ports {\n\t\t\tcontainerPort := nat.Port(fmt.Sprintf(\"%d/tcp\", port.ContainerPort))\n\t\t\tcontainerCfg.ExposedPorts[containerPort] = struct{}{}\n\t\t}\n\t}\n\n\thostCfg := &container.HostConfig{\n\t\t// always remove container after it exits\n\t\tAutoRemove: true,\n\t\t// https://github.com/pytorch/pytorch/issues/2244\n\t\t// https://github.com/replicate/cog/issues/1293\n\t\tShmSize:   6 * 1024 * 1024 * 1024, // 6GB\n\t\tResources: container.Resources{},\n\t}\n\n\tif options.GPUs != \"\" {\n\t\tdeviceRequest, err := parseGPURequest(options)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\thostCfg.DeviceRequests = []container.DeviceRequest{deviceRequest}\n\t}\n\n\t// Configure port bindings\n\tif len(options.Ports) > 0 {\n\t\thostCfg.PortBindings = make(nat.PortMap)\n\t\tfor _, port := range options.Ports {\n\t\t\tcontainerPort := nat.Port(fmt.Sprintf(\"%d/tcp\", port.ContainerPort))\n\t\t\thostCfg.PortBindings[containerPort] = []nat.PortBinding{\n\t\t\t\t{\n\t\t\t\t\tHostIP:   \"\", // use empty string to bind to all interfaces\n\t\t\t\t\tHostPort: strconv.Itoa(port.HostPort),\n\t\t\t\t},\n\t\t\t}\n\t\t}\n\t}\n\n\t// Configure volume bindings\n\tif len(options.Volumes) > 0 {\n\t\thostCfg.Binds = make([]string, len(options.Volumes))\n\t\tfor i, volume := range options.Volumes {\n\t\t\thostCfg.Binds[i] = fmt.Sprintf(\"%s:%s\", volume.Source, volume.Destination)\n\t\t}\n\t}\n\n\t// Configure extra hosts (e.g. host.docker.internal on Linux)\n\tif len(options.ExtraHosts) > 0 {\n\t\thostCfg.ExtraHosts = options.ExtraHosts\n\t}\n\n\tnetworkingCfg := &network.NetworkingConfig{\n\t\tEndpointsConfig: map[string]*network.EndpointSettings{},\n\t}\n\n\tplatform := &ocispec.Platform{\n\t\t// force platform to linux/amd64\n\t\tArchitecture: \"amd64\",\n\t\tOS:           \"linux\",\n\t}\n\n\trunContainer, err := c.client.ContainerCreate(ctx,\n\t\tcontainerCfg,\n\t\thostCfg,\n\t\tnetworkingCfg,\n\t\tplatform,\n\t\t\"\")\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to create container: %w\", err)\n\t}\n\t// TODO[md]: ensure the container is removed if start & auto-remove fails\n\n\tconsole.Debugf(\"container id: %s\", runContainer.ID)\n\n\t// Create error group for stream copying\n\tvar eg *errgroup.Group\n\tvar stream types.HijackedResponse\n\n\t// Attach to container streams if we have any writers and not detached\n\tif attachStderr || attachStdout || attachStdin {\n\t\tattachOpts := container.AttachOptions{\n\t\t\tStream: true,\n\t\t\tStdin:  attachStdin,\n\t\t\tStdout: attachStdout,\n\t\t\tStderr: attachStderr,\n\t\t}\n\n\t\tvar err error\n\t\tstream, err = c.client.ContainerAttach(ctx, runContainer.ID, attachOpts)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to attach to container: %w\", err)\n\t\t}\n\t\tdefer stream.Close()\n\n\t\t// Start copying streams in the background\n\t\teg, _ = errgroup.WithContext(ctx)\n\t\tif attachStdout || attachStderr {\n\t\t\teg.Go(func() (err error) {\n\t\t\t\tif containerCfg.Tty {\n\t\t\t\t\tw := options.Stdout\n\t\t\t\t\tif w == nil {\n\t\t\t\t\t\tw = options.Stderr\n\t\t\t\t\t}\n\t\t\t\t\t_, err = io.Copy(w, stream.Reader)\n\t\t\t\t} else {\n\t\t\t\t\t_, err = stdcopy.StdCopy(options.Stdout, options.Stderr, stream.Reader)\n\t\t\t\t}\n\t\t\t\treturn err\n\t\t\t})\n\t\t}\n\t\tif attachStdin {\n\t\t\t// if we're in a TTY we need to set the terminal to raw mode, and restore it when we're done\n\t\t\tif tty {\n\t\t\t\t// TODO[md]: handle terminal resize events, see: github.com/containerd/console\n\t\t\t\tstate, err := term.SetRawTerminal(os.Stdin.Fd())\n\t\t\t\tif err != nil {\n\t\t\t\t\tconsole.Warnf(\"error setting raw terminal on stdin: %s\", err)\n\t\t\t\t}\n\t\t\t\tdefer func() {\n\t\t\t\t\tif err := term.RestoreTerminal(os.Stdin.Fd(), state); err != nil {\n\t\t\t\t\t\tconsole.Warnf(\"error restoring terminal on stdin: %s\", err)\n\t\t\t\t\t}\n\t\t\t\t}()\n\t\t\t}\n\n\t\t\tgo func() {\n\t\t\t\t_, err := io.Copy(stream.Conn, options.Stdin)\n\t\t\t\t// Close the stdin stream to signal EOF to the container\n\t\t\t\tif err := errors.Join(err, stream.CloseWrite()); err != nil {\n\t\t\t\t\tconsole.Errorf(\"error copying and closing stdin stream: %s\", err)\n\t\t\t\t}\n\t\t\t}()\n\t\t}\n\t}\n\n\t// Start the container\n\tif err := c.client.ContainerStart(ctx, runContainer.ID, container.StartOptions{}); err != nil {\n\t\tif isMissingDeviceDriverError(err) {\n\t\t\treturn \"\", ErrMissingDeviceDriver\n\t\t}\n\t\treturn \"\", fmt.Errorf(\"failed to start container: %w\", err)\n\t}\n\n\t// If detached, wait for container to be running before returning\n\tif options.Detach {\n\t\treturn runContainer.ID, nil\n\t}\n\n\t// Wait for the container to exit\n\tstatusCh, errCh := c.client.ContainerWait(ctx, runContainer.ID, container.WaitConditionNotRunning)\n\tselect {\n\tcase err := <-errCh:\n\t\treturn \"\", fmt.Errorf(\"error waiting for container: %w\", err)\n\tcase status := <-statusCh:\n\t\tif status.StatusCode != 0 {\n\t\t\treturn \"\", fmt.Errorf(\"container exited with status %d\", status.StatusCode)\n\t\t}\n\t}\n\n\t// container is gone, close the attached streams so stdin is released, ignore the error\n\t_ = stream.CloseWrite()\n\n\t// Wait for stream copying to complete\n\tif eg != nil {\n\t\tif err := eg.Wait(); err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"error copying streams: %w\", err)\n\t\t}\n\t}\n\n\treturn runContainer.ID, nil\n}\n\nfunc (c *apiClient) Run(ctx context.Context, options command.RunOptions) error {\n\tconsole.Debugf(\"=== APIClient.Run %s\", options.Image)\n\n\tif options.Stdout == nil {\n\t\toptions.Stdout = os.Stdout\n\t}\n\tif options.Stderr == nil {\n\t\toptions.Stderr = os.Stderr\n\t}\n\n\t_, err := c.containerRun(ctx, options)\n\treturn err\n}\n\nfunc (c *apiClient) ContainerStart(ctx context.Context, options command.RunOptions) (string, error) {\n\tconsole.Debugf(\"=== APIClient.ContainerStart %s\", options.Image)\n\n\toptions.Detach = true\n\tid, err := c.containerRun(ctx, options)\n\treturn id, err\n}\n\n// parseGPURequest converts a Docker CLI --gpus string into a DeviceRequest slice\nfunc parseGPURequest(opts command.RunOptions) (container.DeviceRequest, error) {\n\tif opts.GPUs == \"\" {\n\t\treturn container.DeviceRequest{}, nil\n\t}\n\n\tdeviceRequest := container.DeviceRequest{\n\t\tDriver:       \"nvidia\",\n\t\tCapabilities: [][]string{{\"gpu\"}},\n\t}\n\n\t// Parse the GPUs string\n\tswitch opts.GPUs {\n\tcase \"all\":\n\t\tdeviceRequest.Count = -1 // Use all available GPUs\n\tdefault:\n\t\t// Check if it's a number\n\t\tif count, err := strconv.Atoi(opts.GPUs); err == nil {\n\t\t\tdeviceRequest.Count = count\n\t\t} else if after, ok := strings.CutPrefix(opts.GPUs, \"device=\"); ok {\n\t\t\t// Handle device=0,1 format\n\t\t\tdevices := after\n\t\t\tdeviceRequest.DeviceIDs = strings.Split(devices, \",\")\n\t\t} else {\n\t\t\t// Invalid GPU specification, return nil to indicate no GPU access\n\t\t\treturn container.DeviceRequest{}, fmt.Errorf(\"invalid GPU specification: %q\", opts.GPUs)\n\t\t}\n\t}\n\n\treturn deviceRequest, nil\n}\n\n// shouldAttachStdin determines if we should attach stdin to the container\n// We should attach stdin only if:\n//   - stdin is not os.Stdin (explicit input like pipe/file/buffer)\n//   - OR stdin is os.Stdin but it's not a TTY (piped input)\nfunc shouldAttachStdin(stdin io.Reader) (attach bool, tty bool) {\n\tif stdin == nil {\n\t\treturn false, false\n\t}\n\n\t// If it's not a file, it's probably a buffer/pipe with actual data\n\tf, ok := stdin.(*os.File)\n\tif !ok {\n\t\treturn true, false\n\t}\n\n\ttty = isatty.IsTerminal(f.Fd())\n\n\t// If it's not os.Stdin, it's an explicit file, so attach it\n\tif f != os.Stdin {\n\t\treturn true, tty\n\t}\n\n\t// If it's os.Stdin but not a TTY, it's probably piped input\n\tif !tty {\n\t\treturn true, false\n\t}\n\n\t// If it's os.Stdin and a TTY, attach by default. if this becomes a problem for some\n\t// reason we need to add a flag to the run command similar to `docker run -i` that instructs\n\t// the container to attach stdin and keep open\n\treturn true, true\n}\n"
  },
  {
    "path": "pkg/docker/docker_client_test.go",
    "content": "package docker\n\nimport (\n\t\"bytes\"\n\t\"net\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/docker/api/types/registry\"\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/testcontainers/testcontainers-go\"\n\t\"github.com/testcontainers/testcontainers-go/wait\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n\t\"github.com/replicate/cog/pkg/registry_testhelpers\"\n)\n\nfunc TestDockerClient(t *testing.T) {\n\tif testing.Short() {\n\t\tt.Skip(\"skipping docker client tests in short mode\")\n\t}\n\n\tdockerClient, err := NewClient(t.Context())\n\trequire.NoError(t, err, \"Failed to create docker client\")\n\tdockerHelper := dockertest.NewHelperClient(t)\n\ttestRegistry := registry_testhelpers.StartTestRegistry(t)\n\n\tdockerHelper.CleanupImages(t)\n\n\tt.Run(\"ImageInspect\", func(t *testing.T) {\n\t\tt.Parallel()\n\n\t\tt.Run(\"ExistingLocalImage\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tref := dockertest.NewRef(t)\n\t\t\tdockerHelper.ImageFixture(t, \"alpine\", ref.String())\n\n\t\t\texpectedImage := dockerHelper.InspectImage(t, ref.String())\n\t\t\tresp, err := dockerClient.Inspect(t.Context(), ref.String())\n\t\t\trequire.NoError(t, err, \"Failed to inspect image %q\", ref.String())\n\t\t\tassert.Equal(t, expectedImage.ID, resp.ID)\n\t\t})\n\n\t\tt.Run(\"MissingLocalImage\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\timage := \"not-a-valid-image\"\n\t\t\t_, err := dockerClient.Inspect(t.Context(), image)\n\t\t\tassert.ErrorIs(t, err, &command.NotFoundError{})\n\t\t\tassert.ErrorContains(t, err, \"image not found\")\n\t\t})\n\t})\n\n\tt.Run(\"Pull\", func(t *testing.T) {\n\t\tt.Parallel()\n\n\t\t// TODO[md]: add tests for the following permutations:\n\t\t// - remote reference exists/not exists\n\t\t// - local reference exists/not exists\n\t\t// - force pull true/false\n\n\t\tt.Run(\"RemoteImageExists\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\t\t\trepo := testRegistry.CloneRepoForTest(t, \"alpine\")\n\t\t\timageRef := repo + \":latest\"\n\n\t\t\tassertNoImageExists(t, dockerClient, imageRef)\n\n\t\t\tresp, err := dockerClient.Pull(t.Context(), imageRef, false)\n\t\t\trequire.NoError(t, err, \"Failed to pull image %q\", imageRef)\n\t\t\tdockerHelper.CleanupImage(t, imageRef)\n\n\t\t\tassertImageExists(t, dockerClient, imageRef)\n\t\t\texpectedResp := dockerHelper.InspectImage(t, imageRef)\n\t\t\t// TODO[md]: we should check that the responsees are actually equal beyond the IDs. but atm\n\t\t\t// the CLI and api are slightly different. The CLI leaves the descriptor field nil while the\n\t\t\t// API response is populated. These should be identical on the new client, so we can change to EqualValues\n\t\t\tassert.Equal(t, expectedResp.ID, resp.ID, \"inspect response should match expected\")\n\t\t})\n\n\t\tt.Run(\"RemoteReferenceNotFound\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\t\t\timageRef := testRegistry.ImageRefForTest(t, \"\")\n\n\t\t\tassertNoImageExists(t, dockerClient, imageRef)\n\n\t\t\tresp, err := dockerClient.Pull(t.Context(), imageRef, false)\n\t\t\t// TODO[md]: this might not be the right check. we probably want to wrap the error from the registry\n\t\t\t// so we handle other failure cases, like failed auth, unknown tag, and unknown repo\n\t\t\trequire.Error(t, err, \"Failed to pull image %q\", imageRef)\n\t\t\tassert.ErrorIs(t, err, &command.NotFoundError{Object: \"manifest\", Ref: imageRef})\n\t\t\tassert.Nil(t, resp, \"inspect response should be nil\")\n\t\t})\n\n\t\tt.Run(\"InvalidAuth\", func(t *testing.T) {\n\t\t\tt.Skip(\"skip auth tests until we're using the docker engine since we can't set auth on the host without side effects\")\n\t\t\timageRef := testRegistry.ImageRefForTest(t, \"\")\n\n\t\t\tassertNoImageExists(t, dockerClient, imageRef)\n\n\t\t\tresp, err := dockerClient.Pull(t.Context(), imageRef, false)\n\t\t\t// TODO[md]: this might not be the right check. we probably want to wrap the error from the registry\n\t\t\t// so we handle other failure cases, like failed auth, unknown tag, and unknown repo\n\t\t\trequire.Error(t, err, \"Failed to pull image %q\", imageRef)\n\t\t\tassert.ErrorContains(t, err, \"failed to resolve reference\")\n\t\t\tassert.Nil(t, resp, \"inspect response should be nil\")\n\t\t})\n\t})\n\n\tt.Run(\"ContainerStop\", func(t *testing.T) {\n\t\tt.Parallel()\n\n\t\tt.Run(\"ContainerExistsAndIsRunning\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tcontainer, err := testcontainers.Run(\n\t\t\t\tt.Context(),\n\t\t\t\ttestRegistry.ImageRef(\"alpine:latest\"),\n\t\t\t\ttestcontainers.WithCmd(\"sleep\", \"5000\"),\n\t\t\t)\n\t\t\tdefer dockerHelper.CleanupImages(t)\n\t\t\tdefer testcontainers.CleanupContainer(t, container)\n\t\t\trequire.NoError(t, err, \"Failed to run container\")\n\n\t\t\terr = dockerClient.ContainerStop(t.Context(), container.ID)\n\t\t\trequire.NoError(t, err, \"Failed to stop container %q\", container.ID)\n\n\t\t\tstate, err := container.State(t.Context())\n\t\t\trequire.NoError(t, err, \"Failed to get container state\")\n\t\t\tassert.Equal(t, state.Running, false)\n\t\t})\n\n\t\tt.Run(\"ContainerExistsAndIsNotRunning\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tcontainer, err := testcontainers.GenericContainer(t.Context(),\n\t\t\t\ttestcontainers.GenericContainerRequest{\n\t\t\t\t\tContainerRequest: testcontainers.ContainerRequest{\n\t\t\t\t\t\tImage: testRegistry.ImageRef(\"alpine:latest\"),\n\t\t\t\t\t\tCmd:   []string{\"sleep\", \"5000\"},\n\t\t\t\t\t},\n\t\t\t\t\tStarted: false,\n\t\t\t\t},\n\t\t\t)\n\t\t\tdefer testcontainers.CleanupContainer(t, container)\n\t\t\tcontainerID := container.GetContainerID()\n\t\t\trequire.NoError(t, err, \"Failed to create container\")\n\n\t\t\terr = dockerClient.ContainerStop(t.Context(), containerID)\n\t\t\trequire.NoError(t, err, \"Failed to stop container %q\", containerID)\n\n\t\t\tstate, err := container.State(t.Context())\n\t\t\trequire.NoError(t, err, \"Failed to get container state\")\n\t\t\tassert.Equal(t, state.Running, false)\n\t\t})\n\n\t\tt.Run(\"ContainerDoesNotExist\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\terr := dockerClient.ContainerStop(t.Context(), \"containerid-that-does-not-exist\")\n\t\t\trequire.ErrorIs(t, err, &command.NotFoundError{})\n\t\t\trequire.ErrorContains(t, err, \"container not found\")\n\t\t})\n\t})\n\n\tt.Run(\"ContainerInspect\", func(t *testing.T) {\n\t\tt.Parallel()\n\n\t\tt.Run(\"ContainerExists\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tcontainer, err := testcontainers.Run(\n\t\t\t\tt.Context(),\n\t\t\t\ttestRegistry.ImageRef(\"alpine:latest\"),\n\t\t\t\ttestcontainers.WithCmd(\"sleep\", \"5000\"),\n\t\t\t)\n\t\t\tdefer testcontainers.CleanupContainer(t, container)\n\t\t\trequire.NoError(t, err, \"Failed to run container\")\n\n\t\t\texpected, err := container.Inspect(t.Context())\n\t\t\trequire.NoError(t, err, \"Failed to inspect container for expected response\")\n\n\t\t\tresp, err := dockerClient.ContainerInspect(t.Context(), container.ID)\n\t\t\trequire.NoError(t, err, \"Failed to inspect container\")\n\t\t\trequire.Equal(t, expected, resp)\n\t\t})\n\n\t\tt.Run(\"ContainerDoesNotExist\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\t_, err := dockerClient.ContainerInspect(t.Context(), \"containerid-that-does-not-exist\")\n\t\t\trequire.ErrorIs(t, err, &command.NotFoundError{})\n\t\t})\n\t})\n\n\tt.Run(\"ContainerLogs\", func(t *testing.T) {\n\t\tt.Parallel()\n\n\t\tt.Run(\"ContainerExistsAndIsRunning\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tcontainer, err := testcontainers.Run(\n\t\t\t\tt.Context(),\n\t\t\t\ttestRegistry.ImageRef(\"alpine:latest\"),\n\t\t\t\t// print \"line $i\" N times then exit, where $i is the line number\n\t\t\t\ttestcontainers.WithCmd(\"sh\", \"-c\", \"for i in $(seq 1 5); do echo \\\"$i\\\"; sleep 1; done\"),\n\t\t\t\t// testcontainers.WithConfigModifier(func(config *container.Config) {\n\t\t\t\t// \tconfig.Tty = true\n\t\t\t\t// }),\n\t\t\t)\n\t\t\trequire.NoError(t, err, \"Failed to run container\")\n\t\t\tdefer testcontainers.CleanupContainer(t, container)\n\n\t\t\tvar buf bytes.Buffer\n\t\t\terr = dockerClient.ContainerLogs(t.Context(), container.ID, &buf)\n\t\t\trequire.NoError(t, err, \"Failed to get container logs\")\n\n\t\t\tassert.Equal(t, \"1\\n2\\n3\\n4\\n5\\n\", buf.String())\n\t\t})\n\n\t\tt.Run(\"ContainerAlreadyStopped\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tcontainer, err := testcontainers.Run(\n\t\t\t\tt.Context(),\n\t\t\t\ttestRegistry.ImageRef(\"alpine:latest\"),\n\t\t\t\ttestcontainers.WithCmd(\"sh\", \"-c\", \"for i in $(seq 1 3); do echo \\\"$i\\\"; sleep 0.1; done\"),\n\t\t\t\ttestcontainers.WithWaitStrategy(wait.ForExit()),\n\t\t\t)\n\t\t\trequire.NoError(t, err, \"Failed to run container\")\n\t\t\tdefer testcontainers.CleanupContainer(t, container)\n\n\t\t\tstate, err := container.State(t.Context())\n\t\t\trequire.NoError(t, err, \"Failed to get container state\")\n\t\t\tassert.Equal(t, state.Running, false)\n\n\t\t\tvar buf bytes.Buffer\n\t\t\terr = dockerClient.ContainerLogs(t.Context(), container.ID, &buf)\n\t\t\trequire.NoError(t, err, \"Failed to get container logs\")\n\n\t\t\tassert.Equal(t, \"1\\n2\\n3\\n\", buf.String())\n\t\t})\n\n\t\tt.Run(\"TTY and non-TTY streams match\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\trunContainer := func(tty bool) string {\n\t\t\t\tcontainer, err := testcontainers.Run(\n\t\t\t\t\tt.Context(),\n\t\t\t\t\ttestRegistry.ImageRef(\"alpine:latest\"),\n\t\t\t\t\t// print \"line $i\" N times then exit, where $i is the line number\n\t\t\t\t\ttestcontainers.WithCmd(\"sh\", \"-c\", \"for i in $(seq 1 5); do echo \\\"$i\\\"; sleep 0.1; done\"),\n\t\t\t\t\ttestcontainers.WithConfigModifier(func(config *container.Config) {\n\t\t\t\t\t\tconfig.Tty = tty\n\t\t\t\t\t}),\n\t\t\t\t)\n\t\t\t\trequire.NoError(t, err, \"Failed to run container\")\n\t\t\t\tdefer testcontainers.CleanupContainer(t, container)\n\n\t\t\t\tvar buf bytes.Buffer\n\t\t\t\terr = dockerClient.ContainerLogs(t.Context(), container.ID, &buf)\n\t\t\t\trequire.NoError(t, err, \"Failed to get container logs\")\n\t\t\t\treturn buf.String()\n\t\t\t}\n\n\t\t\tttyOutput := runContainer(true)\n\t\t\tnonTtyOutput := runContainer(false)\n\n\t\t\t// TTY uses CRLF for line endings, non-TTY uses LF. replace \\r\\n with \\n so they match\n\t\t\tttyOutput = strings.ReplaceAll(ttyOutput, \"\\r\\n\", \"\\n\")\n\n\t\t\tassert.Equal(t, ttyOutput, nonTtyOutput, \"TTY and non-TTY streams should match after normalizing line endings\")\n\t\t})\n\n\t\tt.Run(\"ContainerDoesNotExist\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\terr := dockerClient.ContainerLogs(t.Context(), \"containerid-that-does-not-exist\", &bytes.Buffer{})\n\t\t\trequire.ErrorIs(t, err, &command.NotFoundError{})\n\t\t})\n\t})\n\n\tt.Run(\"Push\", func(t *testing.T) {\n\t\tt.Parallel()\n\n\t\tt.Run(\"valid image, valid registry\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tref := dockertest.NewRef(t).WithRegistry(testRegistry.RegistryHost())\n\n\t\t\tdockerHelper.ImageFixture(t, \"alpine\", ref.String())\n\n\t\t\terr := dockerClient.Push(t.Context(), ref.String())\n\t\t\trequire.NoError(t, err)\n\t\t\tassert.NoError(t, testRegistry.ImageExists(t, ref.String()))\n\t\t})\n\n\t\tt.Run(\"non-existent registry\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\t// start a local tcp server that immediately closes connections\n\t\t\tlistener, err := net.Listen(\"tcp\", \"127.0.0.1:0\")\n\t\t\trequire.NoError(t, err)\n\t\t\tdefer listener.Close()\n\n\t\t\tgo func() {\n\t\t\t\tfor {\n\t\t\t\t\tconn, err := listener.Accept()\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn\n\t\t\t\t\t}\n\t\t\t\t\tconn.Close()\n\t\t\t\t}\n\t\t\t}()\n\n\t\t\t// Create a reference to the mock registry\n\t\t\tref := dockertest.NewRef(t).WithRegistry(listener.Addr().String())\n\t\t\tdockerHelper.ImageFixture(t, \"alpine\", ref.String())\n\n\t\t\t// Try to push to the mock registry\n\t\t\terr = dockerClient.Push(t.Context(), ref.String())\n\t\t\trequire.Error(t, err, \"Push should fail with unreachable registry\")\n\n\t\t\tassert.True(t, isNetworkError(err), \"Error should be a network error, got: %q\", err.Error())\n\t\t})\n\n\t\tt.Run(\"missing image\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tref := dockertest.NewRef(t).WithRegistry(testRegistry.RegistryHost())\n\n\t\t\terr := dockerClient.Push(t.Context(), ref.String())\n\t\t\tassertNotFoundError(t, err, ref.String(), \"tag\")\n\t\t})\n\n\t\tt.Run(\"registry with authentication\", func(t *testing.T) {\n\t\t\tt.Parallel()\n\n\t\t\tauthReg := registry_testhelpers.StartTestRegistry(t, registry_testhelpers.WithAuth(\"testuser\", \"testpass\"))\n\n\t\t\tt.Run(\"correct credentials\", func(t *testing.T) {\n\t\t\t\tt.Parallel()\n\n\t\t\t\tref := dockertest.NewRef(t).WithRegistry(authReg.RegistryHost())\n\t\t\t\tdockerHelper.ImageFixture(t, \"alpine\", ref.String())\n\n\t\t\t\t// create a new client with the correct auth config\n\t\t\t\tauthClient, err := NewClient(t.Context(), WithAuthConfig(registry.AuthConfig{\n\t\t\t\t\tUsername:      \"testuser\",\n\t\t\t\t\tPassword:      \"testpass\",\n\t\t\t\t\tServerAddress: authReg.RegistryHost(),\n\t\t\t\t}))\n\t\t\t\trequire.NoError(t, err)\n\n\t\t\t\terr = authClient.Push(t.Context(), ref.String())\n\t\t\t\trequire.NoError(t, err, \"Failed to push image to auth registry\")\n\t\t\t\tassert.NoError(t, authReg.ImageExists(t, ref.String()))\n\t\t\t})\n\n\t\t\tt.Run(\"missing auth\", func(t *testing.T) {\n\t\t\t\tt.Parallel()\n\n\t\t\t\tref := dockertest.NewRef(t).WithRegistry(authReg.RegistryHost())\n\t\t\t\tdockerHelper.ImageFixture(t, \"alpine\", ref.String())\n\n\t\t\t\t// use root client which doesn't have auth setup\n\t\t\t\terr := dockerClient.Push(t.Context(), ref.String())\n\t\t\t\trequire.ErrorIs(t, err, command.ErrAuthorizationFailed)\n\t\t\t})\n\n\t\t\tt.Run(\"incorrect auth\", func(t *testing.T) {\n\t\t\t\tt.Parallel()\n\n\t\t\t\tref := dockertest.NewRef(t).WithRegistry(authReg.RegistryHost())\n\t\t\t\tdockerHelper.ImageFixture(t, \"alpine\", ref.String())\n\n\t\t\t\tauthClient, err := NewClient(t.Context(), WithAuthConfig(registry.AuthConfig{\n\t\t\t\t\tUsername:      \"testuser\",\n\t\t\t\t\tPassword:      \"wrongpass\",\n\t\t\t\t\tServerAddress: authReg.RegistryHost(),\n\t\t\t\t}))\n\t\t\t\trequire.NoError(t, err)\n\n\t\t\t\terr = authClient.Push(t.Context(), ref.String())\n\t\t\t\trequire.ErrorIs(t, err, command.ErrAuthorizationFailed)\n\t\t\t})\n\n\t\t\tt.Run(\"correct credentials, not authorized\", func(t *testing.T) {\n\t\t\t\tt.Skip(\"skipping until the registry supports repo authorizations\")\n\t\t\t})\n\t\t})\n\t})\n}\n\nfunc assertImageExists(t *testing.T, dockerClient command.Command, imageRef string) {\n\tt.Helper()\n\n\tinspect, err := dockerClient.Inspect(t.Context(), imageRef)\n\tassert.NoError(t, err, \"Failed to inspect image %q\", imageRef)\n\tassert.NotNil(t, inspect, \"Image should exist\")\n}\n\nfunc assertNoImageExists(t *testing.T, dockerClient command.Command, imageRef string) {\n\tt.Helper()\n\n\tinspect, err := dockerClient.Inspect(t.Context(), imageRef)\n\tassert.ErrorIs(t, err, &command.NotFoundError{}, \"Image should not exist\")\n\tassert.Nil(t, inspect, \"Image should not exist\")\n}\n\nfunc assertNotFoundError(t *testing.T, err error, ref string, object string) {\n\tt.Helper()\n\n\tvar notFoundErr *command.NotFoundError\n\trequire.ErrorAs(t, err, ¬FoundErr, \"should be a not found error\")\n\trequire.Equal(t, ref, notFoundErr.Ref, \"ref should match\")\n\trequire.Equal(t, object, notFoundErr.Object, \"object should match\")\n}\n"
  },
  {
    "path": "pkg/docker/dockertest/command_mocks.go",
    "content": "// Code generated by mockery; DO NOT EDIT.\n// github.com/vektra/mockery\n// template: testify\n\npackage dockertest\n\nimport (\n\t\"context\"\n\t\"io\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/docker/api/types/image\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\tmock \"github.com/stretchr/testify/mock\"\n)\n\n// NewMockCommand2 creates a new instance of MockCommand2. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.\n// The first argument is typically a *testing.T value.\nfunc NewMockCommand2(t interface {\n\tmock.TestingT\n\tCleanup(func())\n}) *MockCommand2 {\n\tmock := &MockCommand2{}\n\tmock.Mock.Test(t)\n\n\tt.Cleanup(func() { mock.AssertExpectations(t) })\n\n\treturn mock\n}\n\n// MockCommand2 is an autogenerated mock type for the Command type\ntype MockCommand2 struct {\n\tmock.Mock\n}\n\ntype MockCommand2_Expecter struct {\n\tmock *mock.Mock\n}\n\nfunc (_m *MockCommand2) EXPECT() *MockCommand2_Expecter {\n\treturn &MockCommand2_Expecter{mock: &_m.Mock}\n}\n\n// ContainerInspect provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) ContainerInspect(ctx context.Context, id string) (*container.InspectResponse, error) {\n\tret := _mock.Called(ctx, id)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for ContainerInspect\")\n\t}\n\n\tvar r0 *container.InspectResponse\n\tvar r1 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) (*container.InspectResponse, error)); ok {\n\t\treturn returnFunc(ctx, id)\n\t}\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) *container.InspectResponse); ok {\n\t\tr0 = returnFunc(ctx, id)\n\t} else {\n\t\tif ret.Get(0) != nil {\n\t\t\tr0 = ret.Get(0).(*container.InspectResponse)\n\t\t}\n\t}\n\tif returnFunc, ok := ret.Get(1).(func(context.Context, string) error); ok {\n\t\tr1 = returnFunc(ctx, id)\n\t} else {\n\t\tr1 = ret.Error(1)\n\t}\n\treturn r0, r1\n}\n\n// MockCommand2_ContainerInspect_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ContainerInspect'\ntype MockCommand2_ContainerInspect_Call struct {\n\t*mock.Call\n}\n\n// ContainerInspect is a helper method to define mock.On call\n//   - ctx context.Context\n//   - id string\nfunc (_e *MockCommand2_Expecter) ContainerInspect(ctx interface{}, id interface{}) *MockCommand2_ContainerInspect_Call {\n\treturn &MockCommand2_ContainerInspect_Call{Call: _e.mock.On(\"ContainerInspect\", ctx, id)}\n}\n\nfunc (_c *MockCommand2_ContainerInspect_Call) Run(run func(ctx context.Context, id string)) *MockCommand2_ContainerInspect_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ContainerInspect_Call) Return(inspectResponse *container.InspectResponse, err error) *MockCommand2_ContainerInspect_Call {\n\t_c.Call.Return(inspectResponse, err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ContainerInspect_Call) RunAndReturn(run func(ctx context.Context, id string) (*container.InspectResponse, error)) *MockCommand2_ContainerInspect_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// ContainerLogs provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) ContainerLogs(ctx context.Context, containerID string, w io.Writer) error {\n\tret := _mock.Called(ctx, containerID, w)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for ContainerLogs\")\n\t}\n\n\tvar r0 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string, io.Writer) error); ok {\n\t\tr0 = returnFunc(ctx, containerID, w)\n\t} else {\n\t\tr0 = ret.Error(0)\n\t}\n\treturn r0\n}\n\n// MockCommand2_ContainerLogs_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ContainerLogs'\ntype MockCommand2_ContainerLogs_Call struct {\n\t*mock.Call\n}\n\n// ContainerLogs is a helper method to define mock.On call\n//   - ctx context.Context\n//   - containerID string\n//   - w io.Writer\nfunc (_e *MockCommand2_Expecter) ContainerLogs(ctx interface{}, containerID interface{}, w interface{}) *MockCommand2_ContainerLogs_Call {\n\treturn &MockCommand2_ContainerLogs_Call{Call: _e.mock.On(\"ContainerLogs\", ctx, containerID, w)}\n}\n\nfunc (_c *MockCommand2_ContainerLogs_Call) Run(run func(ctx context.Context, containerID string, w io.Writer)) *MockCommand2_ContainerLogs_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\tvar arg2 io.Writer\n\t\tif args[2] != nil {\n\t\t\targ2 = args[2].(io.Writer)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t\targ2,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ContainerLogs_Call) Return(err error) *MockCommand2_ContainerLogs_Call {\n\t_c.Call.Return(err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ContainerLogs_Call) RunAndReturn(run func(ctx context.Context, containerID string, w io.Writer) error) *MockCommand2_ContainerLogs_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// ContainerStart provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) ContainerStart(ctx context.Context, options command.RunOptions) (string, error) {\n\tret := _mock.Called(ctx, options)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for ContainerStart\")\n\t}\n\n\tvar r0 string\n\tvar r1 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, command.RunOptions) (string, error)); ok {\n\t\treturn returnFunc(ctx, options)\n\t}\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, command.RunOptions) string); ok {\n\t\tr0 = returnFunc(ctx, options)\n\t} else {\n\t\tr0 = ret.Get(0).(string)\n\t}\n\tif returnFunc, ok := ret.Get(1).(func(context.Context, command.RunOptions) error); ok {\n\t\tr1 = returnFunc(ctx, options)\n\t} else {\n\t\tr1 = ret.Error(1)\n\t}\n\treturn r0, r1\n}\n\n// MockCommand2_ContainerStart_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ContainerStart'\ntype MockCommand2_ContainerStart_Call struct {\n\t*mock.Call\n}\n\n// ContainerStart is a helper method to define mock.On call\n//   - ctx context.Context\n//   - options command.RunOptions\nfunc (_e *MockCommand2_Expecter) ContainerStart(ctx interface{}, options interface{}) *MockCommand2_ContainerStart_Call {\n\treturn &MockCommand2_ContainerStart_Call{Call: _e.mock.On(\"ContainerStart\", ctx, options)}\n}\n\nfunc (_c *MockCommand2_ContainerStart_Call) Run(run func(ctx context.Context, options command.RunOptions)) *MockCommand2_ContainerStart_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 command.RunOptions\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(command.RunOptions)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ContainerStart_Call) Return(s string, err error) *MockCommand2_ContainerStart_Call {\n\t_c.Call.Return(s, err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ContainerStart_Call) RunAndReturn(run func(ctx context.Context, options command.RunOptions) (string, error)) *MockCommand2_ContainerStart_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// ContainerStop provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) ContainerStop(ctx context.Context, containerID string) error {\n\tret := _mock.Called(ctx, containerID)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for ContainerStop\")\n\t}\n\n\tvar r0 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) error); ok {\n\t\tr0 = returnFunc(ctx, containerID)\n\t} else {\n\t\tr0 = ret.Error(0)\n\t}\n\treturn r0\n}\n\n// MockCommand2_ContainerStop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ContainerStop'\ntype MockCommand2_ContainerStop_Call struct {\n\t*mock.Call\n}\n\n// ContainerStop is a helper method to define mock.On call\n//   - ctx context.Context\n//   - containerID string\nfunc (_e *MockCommand2_Expecter) ContainerStop(ctx interface{}, containerID interface{}) *MockCommand2_ContainerStop_Call {\n\treturn &MockCommand2_ContainerStop_Call{Call: _e.mock.On(\"ContainerStop\", ctx, containerID)}\n}\n\nfunc (_c *MockCommand2_ContainerStop_Call) Run(run func(ctx context.Context, containerID string)) *MockCommand2_ContainerStop_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ContainerStop_Call) Return(err error) *MockCommand2_ContainerStop_Call {\n\t_c.Call.Return(err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ContainerStop_Call) RunAndReturn(run func(ctx context.Context, containerID string) error) *MockCommand2_ContainerStop_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// ImageBuild provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) ImageBuild(ctx context.Context, options command.ImageBuildOptions) (string, error) {\n\tret := _mock.Called(ctx, options)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for ImageBuild\")\n\t}\n\n\tvar r0 string\n\tvar r1 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, command.ImageBuildOptions) (string, error)); ok {\n\t\treturn returnFunc(ctx, options)\n\t}\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, command.ImageBuildOptions) string); ok {\n\t\tr0 = returnFunc(ctx, options)\n\t} else {\n\t\tr0 = ret.Get(0).(string)\n\t}\n\tif returnFunc, ok := ret.Get(1).(func(context.Context, command.ImageBuildOptions) error); ok {\n\t\tr1 = returnFunc(ctx, options)\n\t} else {\n\t\tr1 = ret.Error(1)\n\t}\n\treturn r0, r1\n}\n\n// MockCommand2_ImageBuild_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ImageBuild'\ntype MockCommand2_ImageBuild_Call struct {\n\t*mock.Call\n}\n\n// ImageBuild is a helper method to define mock.On call\n//   - ctx context.Context\n//   - options command.ImageBuildOptions\nfunc (_e *MockCommand2_Expecter) ImageBuild(ctx interface{}, options interface{}) *MockCommand2_ImageBuild_Call {\n\treturn &MockCommand2_ImageBuild_Call{Call: _e.mock.On(\"ImageBuild\", ctx, options)}\n}\n\nfunc (_c *MockCommand2_ImageBuild_Call) Run(run func(ctx context.Context, options command.ImageBuildOptions)) *MockCommand2_ImageBuild_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 command.ImageBuildOptions\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(command.ImageBuildOptions)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ImageBuild_Call) Return(imageID string, err error) *MockCommand2_ImageBuild_Call {\n\t_c.Call.Return(imageID, err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ImageBuild_Call) RunAndReturn(run func(ctx context.Context, options command.ImageBuildOptions) (string, error)) *MockCommand2_ImageBuild_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// ImageSave provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) ImageSave(ctx context.Context, imageRef string) (io.ReadCloser, error) {\n\tret := _mock.Called(ctx, imageRef)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for ImageSave\")\n\t}\n\n\tvar r0 io.ReadCloser\n\tvar r1 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) (io.ReadCloser, error)); ok {\n\t\treturn returnFunc(ctx, imageRef)\n\t}\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) io.ReadCloser); ok {\n\t\tr0 = returnFunc(ctx, imageRef)\n\t} else {\n\t\tif ret.Get(0) != nil {\n\t\t\tr0 = ret.Get(0).(io.ReadCloser)\n\t\t}\n\t}\n\tif returnFunc, ok := ret.Get(1).(func(context.Context, string) error); ok {\n\t\tr1 = returnFunc(ctx, imageRef)\n\t} else {\n\t\tr1 = ret.Error(1)\n\t}\n\treturn r0, r1\n}\n\n// MockCommand2_ImageSave_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ImageSave'\ntype MockCommand2_ImageSave_Call struct {\n\t*mock.Call\n}\n\n// ImageSave is a helper method to define mock.On call\n//   - ctx context.Context\n//   - imageRef string\nfunc (_e *MockCommand2_Expecter) ImageSave(ctx interface{}, imageRef interface{}) *MockCommand2_ImageSave_Call {\n\treturn &MockCommand2_ImageSave_Call{Call: _e.mock.On(\"ImageSave\", ctx, imageRef)}\n}\n\nfunc (_c *MockCommand2_ImageSave_Call) Run(run func(ctx context.Context, imageRef string)) *MockCommand2_ImageSave_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ImageSave_Call) Return(rc io.ReadCloser, err error) *MockCommand2_ImageSave_Call {\n\t_c.Call.Return(rc, err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ImageSave_Call) RunAndReturn(run func(ctx context.Context, imageRef string) (io.ReadCloser, error)) *MockCommand2_ImageSave_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// ImageExists provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) ImageExists(ctx context.Context, ref string) (bool, error) {\n\tret := _mock.Called(ctx, ref)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for ImageExists\")\n\t}\n\n\tvar r0 bool\n\tvar r1 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) (bool, error)); ok {\n\t\treturn returnFunc(ctx, ref)\n\t}\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) bool); ok {\n\t\tr0 = returnFunc(ctx, ref)\n\t} else {\n\t\tr0 = ret.Get(0).(bool)\n\t}\n\tif returnFunc, ok := ret.Get(1).(func(context.Context, string) error); ok {\n\t\tr1 = returnFunc(ctx, ref)\n\t} else {\n\t\tr1 = ret.Error(1)\n\t}\n\treturn r0, r1\n}\n\n// MockCommand2_ImageExists_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ImageExists'\ntype MockCommand2_ImageExists_Call struct {\n\t*mock.Call\n}\n\n// ImageExists is a helper method to define mock.On call\n//   - ctx context.Context\n//   - ref string\nfunc (_e *MockCommand2_Expecter) ImageExists(ctx interface{}, ref interface{}) *MockCommand2_ImageExists_Call {\n\treturn &MockCommand2_ImageExists_Call{Call: _e.mock.On(\"ImageExists\", ctx, ref)}\n}\n\nfunc (_c *MockCommand2_ImageExists_Call) Run(run func(ctx context.Context, ref string)) *MockCommand2_ImageExists_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ImageExists_Call) Return(b bool, err error) *MockCommand2_ImageExists_Call {\n\t_c.Call.Return(b, err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_ImageExists_Call) RunAndReturn(run func(ctx context.Context, ref string) (bool, error)) *MockCommand2_ImageExists_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// Inspect provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) Inspect(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\tret := _mock.Called(ctx, ref)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for Inspect\")\n\t}\n\n\tvar r0 *image.InspectResponse\n\tvar r1 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) (*image.InspectResponse, error)); ok {\n\t\treturn returnFunc(ctx, ref)\n\t}\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) *image.InspectResponse); ok {\n\t\tr0 = returnFunc(ctx, ref)\n\t} else {\n\t\tif ret.Get(0) != nil {\n\t\t\tr0 = ret.Get(0).(*image.InspectResponse)\n\t\t}\n\t}\n\tif returnFunc, ok := ret.Get(1).(func(context.Context, string) error); ok {\n\t\tr1 = returnFunc(ctx, ref)\n\t} else {\n\t\tr1 = ret.Error(1)\n\t}\n\treturn r0, r1\n}\n\n// MockCommand2_Inspect_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Inspect'\ntype MockCommand2_Inspect_Call struct {\n\t*mock.Call\n}\n\n// Inspect is a helper method to define mock.On call\n//   - ctx context.Context\n//   - ref string\nfunc (_e *MockCommand2_Expecter) Inspect(ctx interface{}, ref interface{}) *MockCommand2_Inspect_Call {\n\treturn &MockCommand2_Inspect_Call{Call: _e.mock.On(\"Inspect\", ctx, ref)}\n}\n\nfunc (_c *MockCommand2_Inspect_Call) Run(run func(ctx context.Context, ref string)) *MockCommand2_Inspect_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_Inspect_Call) Return(inspectResponse *image.InspectResponse, err error) *MockCommand2_Inspect_Call {\n\t_c.Call.Return(inspectResponse, err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_Inspect_Call) RunAndReturn(run func(ctx context.Context, ref string) (*image.InspectResponse, error)) *MockCommand2_Inspect_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// LoadUserInformation provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) LoadUserInformation(ctx context.Context, registryHost string) (*command.UserInfo, error) {\n\tret := _mock.Called(ctx, registryHost)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for LoadUserInformation\")\n\t}\n\n\tvar r0 *command.UserInfo\n\tvar r1 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) (*command.UserInfo, error)); ok {\n\t\treturn returnFunc(ctx, registryHost)\n\t}\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) *command.UserInfo); ok {\n\t\tr0 = returnFunc(ctx, registryHost)\n\t} else {\n\t\tif ret.Get(0) != nil {\n\t\t\tr0 = ret.Get(0).(*command.UserInfo)\n\t\t}\n\t}\n\tif returnFunc, ok := ret.Get(1).(func(context.Context, string) error); ok {\n\t\tr1 = returnFunc(ctx, registryHost)\n\t} else {\n\t\tr1 = ret.Error(1)\n\t}\n\treturn r0, r1\n}\n\n// MockCommand2_LoadUserInformation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadUserInformation'\ntype MockCommand2_LoadUserInformation_Call struct {\n\t*mock.Call\n}\n\n// LoadUserInformation is a helper method to define mock.On call\n//   - ctx context.Context\n//   - registryHost string\nfunc (_e *MockCommand2_Expecter) LoadUserInformation(ctx interface{}, registryHost interface{}) *MockCommand2_LoadUserInformation_Call {\n\treturn &MockCommand2_LoadUserInformation_Call{Call: _e.mock.On(\"LoadUserInformation\", ctx, registryHost)}\n}\n\nfunc (_c *MockCommand2_LoadUserInformation_Call) Run(run func(ctx context.Context, registryHost string)) *MockCommand2_LoadUserInformation_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_LoadUserInformation_Call) Return(userInfo *command.UserInfo, err error) *MockCommand2_LoadUserInformation_Call {\n\t_c.Call.Return(userInfo, err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_LoadUserInformation_Call) RunAndReturn(run func(ctx context.Context, registryHost string) (*command.UserInfo, error)) *MockCommand2_LoadUserInformation_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// Pull provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) Pull(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) {\n\tret := _mock.Called(ctx, ref, force)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for Pull\")\n\t}\n\n\tvar r0 *image.InspectResponse\n\tvar r1 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string, bool) (*image.InspectResponse, error)); ok {\n\t\treturn returnFunc(ctx, ref, force)\n\t}\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string, bool) *image.InspectResponse); ok {\n\t\tr0 = returnFunc(ctx, ref, force)\n\t} else {\n\t\tif ret.Get(0) != nil {\n\t\t\tr0 = ret.Get(0).(*image.InspectResponse)\n\t\t}\n\t}\n\tif returnFunc, ok := ret.Get(1).(func(context.Context, string, bool) error); ok {\n\t\tr1 = returnFunc(ctx, ref, force)\n\t} else {\n\t\tr1 = ret.Error(1)\n\t}\n\treturn r0, r1\n}\n\n// MockCommand2_Pull_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Pull'\ntype MockCommand2_Pull_Call struct {\n\t*mock.Call\n}\n\n// Pull is a helper method to define mock.On call\n//   - ctx context.Context\n//   - ref string\n//   - force bool\nfunc (_e *MockCommand2_Expecter) Pull(ctx interface{}, ref interface{}, force interface{}) *MockCommand2_Pull_Call {\n\treturn &MockCommand2_Pull_Call{Call: _e.mock.On(\"Pull\", ctx, ref, force)}\n}\n\nfunc (_c *MockCommand2_Pull_Call) Run(run func(ctx context.Context, ref string, force bool)) *MockCommand2_Pull_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\tvar arg2 bool\n\t\tif args[2] != nil {\n\t\t\targ2 = args[2].(bool)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t\targ2,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_Pull_Call) Return(inspectResponse *image.InspectResponse, err error) *MockCommand2_Pull_Call {\n\t_c.Call.Return(inspectResponse, err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_Pull_Call) RunAndReturn(run func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error)) *MockCommand2_Pull_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// Push provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) Push(ctx context.Context, ref string) error {\n\tret := _mock.Called(ctx, ref)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for Push\")\n\t}\n\n\tvar r0 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) error); ok {\n\t\tr0 = returnFunc(ctx, ref)\n\t} else {\n\t\tr0 = ret.Error(0)\n\t}\n\treturn r0\n}\n\n// MockCommand2_Push_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Push'\ntype MockCommand2_Push_Call struct {\n\t*mock.Call\n}\n\n// Push is a helper method to define mock.On call\n//   - ctx context.Context\n//   - ref string\nfunc (_e *MockCommand2_Expecter) Push(ctx interface{}, ref interface{}) *MockCommand2_Push_Call {\n\treturn &MockCommand2_Push_Call{Call: _e.mock.On(\"Push\", ctx, ref)}\n}\n\nfunc (_c *MockCommand2_Push_Call) Run(run func(ctx context.Context, ref string)) *MockCommand2_Push_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_Push_Call) Return(err error) *MockCommand2_Push_Call {\n\t_c.Call.Return(err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_Push_Call) RunAndReturn(run func(ctx context.Context, ref string) error) *MockCommand2_Push_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// RemoveImage provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) RemoveImage(ctx context.Context, ref string) error {\n\tret := _mock.Called(ctx, ref)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for RemoveImage\")\n\t}\n\n\tvar r0 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, string) error); ok {\n\t\tr0 = returnFunc(ctx, ref)\n\t} else {\n\t\tr0 = ret.Error(0)\n\t}\n\treturn r0\n}\n\n// MockCommand2_RemoveImage_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveImage'\ntype MockCommand2_RemoveImage_Call struct {\n\t*mock.Call\n}\n\n// RemoveImage is a helper method to define mock.On call\n//   - ctx context.Context\n//   - ref string\nfunc (_e *MockCommand2_Expecter) RemoveImage(ctx interface{}, ref interface{}) *MockCommand2_RemoveImage_Call {\n\treturn &MockCommand2_RemoveImage_Call{Call: _e.mock.On(\"RemoveImage\", ctx, ref)}\n}\n\nfunc (_c *MockCommand2_RemoveImage_Call) Run(run func(ctx context.Context, ref string)) *MockCommand2_RemoveImage_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 string\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(string)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_RemoveImage_Call) Return(err error) *MockCommand2_RemoveImage_Call {\n\t_c.Call.Return(err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_RemoveImage_Call) RunAndReturn(run func(ctx context.Context, ref string) error) *MockCommand2_RemoveImage_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n\n// Run provides a mock function for the type MockCommand2\nfunc (_mock *MockCommand2) Run(ctx context.Context, options command.RunOptions) error {\n\tret := _mock.Called(ctx, options)\n\n\tif len(ret) == 0 {\n\t\tpanic(\"no return value specified for Run\")\n\t}\n\n\tvar r0 error\n\tif returnFunc, ok := ret.Get(0).(func(context.Context, command.RunOptions) error); ok {\n\t\tr0 = returnFunc(ctx, options)\n\t} else {\n\t\tr0 = ret.Error(0)\n\t}\n\treturn r0\n}\n\n// MockCommand2_Run_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Run'\ntype MockCommand2_Run_Call struct {\n\t*mock.Call\n}\n\n// Run is a helper method to define mock.On call\n//   - ctx context.Context\n//   - options command.RunOptions\nfunc (_e *MockCommand2_Expecter) Run(ctx interface{}, options interface{}) *MockCommand2_Run_Call {\n\treturn &MockCommand2_Run_Call{Call: _e.mock.On(\"Run\", ctx, options)}\n}\n\nfunc (_c *MockCommand2_Run_Call) Run(run func(ctx context.Context, options command.RunOptions)) *MockCommand2_Run_Call {\n\t_c.Call.Run(func(args mock.Arguments) {\n\t\tvar arg0 context.Context\n\t\tif args[0] != nil {\n\t\t\targ0 = args[0].(context.Context)\n\t\t}\n\t\tvar arg1 command.RunOptions\n\t\tif args[1] != nil {\n\t\t\targ1 = args[1].(command.RunOptions)\n\t\t}\n\t\trun(\n\t\t\targ0,\n\t\t\targ1,\n\t\t)\n\t})\n\treturn _c\n}\n\nfunc (_c *MockCommand2_Run_Call) Return(err error) *MockCommand2_Run_Call {\n\t_c.Call.Return(err)\n\treturn _c\n}\n\nfunc (_c *MockCommand2_Run_Call) RunAndReturn(run func(ctx context.Context, options command.RunOptions) error) *MockCommand2_Run_Call {\n\t_c.Call.Return(run)\n\treturn _c\n}\n"
  },
  {
    "path": "pkg/docker/dockertest/helper_client.go",
    "content": "package dockertest\n\nimport (\n\t\"context\"\n\t\"encoding/base64\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"runtime\"\n\t\"slices\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/docker/api/types/image\"\n\t\"github.com/docker/docker/api/types/registry\"\n\t\"github.com/docker/docker/client\"\n\t\"github.com/stretchr/testify/require\"\n)\n\n// NewHelperClient returns a Docker client for testing.\n// It skips the test if Docker is not available.\nfunc NewHelperClient(t testing.TB) *HelperClient {\n\tt.Helper()\n\n\t// Check if we should skip integration tests\n\tif os.Getenv(\"SKIP_INTEGRATION_TESTS\") == \"1\" {\n\t\tt.Skip(\"Skipping integration tests\")\n\t}\n\n\t// Create Docker client\n\tcli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())\n\tif err != nil {\n\t\tt.Fatalf(\"Failed to create Docker client: %v\", err)\n\t}\n\n\t// Verify Docker daemon is running\n\t_, err = cli.Ping(t.Context())\n\tif err != nil {\n\t\tt.Skip(\"Docker daemon is not running\")\n\t}\n\n\thelper := &HelperClient{\n\t\tClient:   cli,\n\t\tfixtures: make(map[string]*imageFixture),\n\t\tmu:       &sync.Mutex{},\n\t}\n\n\tt.Cleanup(func() {\n\t\tfor _, img := range helper.fixtures {\n\t\t\t_, err := helper.Client.ImageRemove(context.Background(), img.imageID, image.RemoveOptions{Force: true, PruneChildren: true})\n\t\t\tif err != nil {\n\t\t\t\tt.Logf(\"Warning: Failed to remove image %q: %v\", img.imageID, err)\n\t\t\t}\n\t\t}\n\n\t\tif err := cli.Close(); err != nil {\n\t\t\tt.Fatalf(\"Failed to close Docker client: %v\", err)\n\t\t}\n\t})\n\n\treturn helper\n}\n\ntype HelperClient struct {\n\tClient *client.Client\n\n\tmu       *sync.Mutex\n\tfixtures map[string]*imageFixture\n}\n\nfunc (c *HelperClient) Close() error {\n\treturn c.Client.Close()\n}\n\nfunc (c *HelperClient) PullImage(t testing.TB, ref string) error {\n\tt.Helper()\n\tout, err := c.Client.ImagePull(t.Context(), ref, image.PullOptions{})\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer out.Close()\n\n\tt.Cleanup(func() {\n\t\tt.Logf(\"Removing image %q\", ref)\n\n\t\topts := image.RemoveOptions{\n\t\t\tForce: true,\n\t\t}\n\n\t\t// use a background context because t.Context() is already closed when cleanup functions are called\n\t\tif _, err := c.Client.ImageRemove(context.Background(), ref, opts); err != nil {\n\t\t\tt.Logf(\"Warning: Failed to remove image %q: %v\", ref, err)\n\t\t}\n\t})\n\n\t_, err = io.Copy(os.Stderr, out)\n\treturn err\n}\n\nfunc (c *HelperClient) MustPullImage(t testing.TB, ref string) {\n\tt.Helper()\n\trequire.NoError(t, c.PullImage(t, ref), \"Failed to pull image %q\", ref)\n}\n\nfunc (c *HelperClient) PushImage(t testing.TB, ref string) error {\n\tt.Helper()\n\n\t// Create auth config for the registry\n\tauthConfig := registry.AuthConfig{\n\t\t// \"username\": \"testuser\",\n\t\t// \"password\": \"testpassword\",\n\t}\n\tauthBytes, err := json.Marshal(authConfig)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to marshal auth config: %w\", err)\n\t}\n\tauthStr := base64.URLEncoding.EncodeToString(authBytes)\n\n\tout, err := c.Client.ImagePush(t.Context(), ref, image.PushOptions{\n\t\tRegistryAuth: authStr,\n\t})\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer out.Close()\n\n\t_, err = io.Copy(os.Stdout, out)\n\treturn err\n}\n\nfunc (c *HelperClient) MustPushImage(t testing.TB, ref string) {\n\tt.Helper()\n\trequire.NoError(t, c.PushImage(t, ref), \"Failed to push image %q\", ref)\n}\n\nfunc (c *HelperClient) RunContainer(t testing.TB, imageName string) string {\n\tt.Helper()\n\n\tcontainerConfig := &container.Config{\n\t\tImage: imageName,\n\t\tCmd:   []string{\"sleep\", \"60\"}, // Run a long sleep to keep container alive\n\t}\n\thostConfig := &container.HostConfig{\n\t\tAutoRemove: true,\n\t}\n\n\tresp, err := c.Client.ContainerCreate(t.Context(), containerConfig, hostConfig, nil, nil, \"\")\n\trequire.NoError(t, err, \"Failed to create container\")\n\tcontainerID := resp.ID\n\tt.Cleanup(func() {\n\t\tt.Logf(\"Removing container %q\", containerID)\n\t\t_ = c.Client.ContainerRemove(context.Background(), containerID, container.RemoveOptions{\n\t\t\tRemoveVolumes: true,\n\t\t\tRemoveLinks:   false,\n\t\t\tForce:         true,\n\t\t})\n\t})\n\n\tt.Logf(\"Created container %q\", containerID)\n\tif len(resp.Warnings) > 0 {\n\t\tt.Logf(\"Warnings: %v\", resp.Warnings)\n\t}\n\n\tif err := c.Client.ContainerStart(t.Context(), containerID, container.StartOptions{}); err != nil {\n\t\trequire.NoErrorf(t, err, \"Failed to start container\")\n\t\tt.Cleanup(func() {\n\t\t\tt.Logf(\"Stopping container %q\", containerID)\n\t\t\t_ = c.Client.ContainerStop(context.Background(), containerID, container.StopOptions{\n\t\t\t\tTimeout: new(int),\n\t\t\t})\n\t\t})\n\t}\n\n\treturn resp.ID\n}\n\nfunc (c *HelperClient) StopContainer(t testing.TB, containerID string) {\n\tt.Helper()\n\n\terr := c.Client.ContainerStop(t.Context(), containerID, container.StopOptions{\n\t\t// set timeout to 0 to force immediate stop\n\t\tTimeout: new(int),\n\t})\n\trequire.NoErrorf(t, err, \"Failed to stop container %q\", containerID)\n}\n\nfunc (c *HelperClient) InspectImage(t testing.TB, imageRef string) *image.InspectResponse {\n\tt.Helper()\n\n\timg, err := c.Client.ImageInspect(t.Context(), imageRef)\n\trequire.NoError(t, err, \"Failed to inspect image %q\", imageRef)\n\n\treturn &img\n}\n\nfunc (c *HelperClient) ImageExists(t testing.TB, imageRef string) bool {\n\tt.Helper()\n\n\t_, err := c.Client.ImageInspect(t.Context(), imageRef)\n\treturn err == nil\n}\n\nfunc (c *HelperClient) DeleteImage(t testing.TB, imageRef string) error {\n\tt.Helper()\n\n\t_, err := c.Client.ImageRemove(t.Context(), imageRef, image.RemoveOptions{\n\t\tForce:         true,\n\t\tPruneChildren: true,\n\t})\n\treturn err\n}\n\nfunc (c *HelperClient) MustDeleteImage(t testing.TB, imageRef string) {\n\tt.Helper()\n\n\t_, err := c.Client.ImageRemove(t.Context(), imageRef, image.RemoveOptions{\n\t\tForce:         true,\n\t\tPruneChildren: true,\n\t})\n\trequire.NoError(t, err, \"Failed to delete image %q\", imageRef)\n}\n\nfunc (c *HelperClient) CleanupImage(t testing.TB, imageRef string) {\n\tt.Helper()\n\n\tt.Cleanup(func() {\n\t\t_, err := c.Client.ImageRemove(context.Background(), imageRef, image.RemoveOptions{\n\t\t\tForce:         true,\n\t\t\tPruneChildren: true,\n\t\t})\n\t\tif err != nil {\n\t\t\tt.Logf(\"Warning: Failed to remove image %q: %v\", imageRef, err)\n\t\t}\n\t})\n}\n\nfunc (c *HelperClient) CleanupImages(t testing.TB) {\n\tt.Helper()\n\n\texistingImages, err := c.Client.ImageList(t.Context(), image.ListOptions{})\n\trequire.NoError(t, err, \"Failed to list images\")\n\n\timageIDs := make([]string, len(existingImages))\n\tfor i, image := range existingImages {\n\t\timageIDs[i] = image.ID\n\t}\n\n\tfmt.Println(\"existing imageIDs\", imageIDs)\n\n\tt.Cleanup(func() {\n\t\tnewImages, err := c.Client.ImageList(context.Background(), image.ListOptions{})\n\t\tif err != nil {\n\t\t\tt.Logf(\"Warning: Failed to list images: %v\", err)\n\t\t\treturn\n\t\t}\n\n\t\tfor _, image := range newImages {\n\t\t\tfmt.Println(\"new image\", image.ID)\n\t\t\tif !slices.Contains(imageIDs, image.ID) {\n\t\t\t\tc.CleanupImage(t, image.ID)\n\t\t\t}\n\t\t}\n\t})\n}\n\nfunc (c *HelperClient) InspectContainer(t testing.TB, containerID string) *container.InspectResponse {\n\tt.Helper()\n\n\tinspect, err := c.Client.ContainerInspect(t.Context(), containerID)\n\trequire.NoError(t, err, \"Failed to inspect container %q\", containerID)\n\n\treturn &inspect\n}\n\nfunc (c *HelperClient) ImageFixture(t testing.TB, name string, tag string) {\n\tt.Helper()\n\tfixture := c.loadImageFixture(t, name)\n\n\tt.Logf(\"Tagging image fixture %q with %q\", fixture.ref, tag)\n\tif err := c.Client.ImageTag(t.Context(), fixture.imageID, tag); err != nil {\n\t\trequire.NoError(t, err, \"Failed to tag image %q with %q: %v\", fixture.ref, tag, err)\n\t}\n\t// remove the image when the test is done\n\tt.Cleanup(func() {\n\t\t_, _ = c.Client.ImageRemove(context.Background(), tag, image.RemoveOptions{Force: true})\n\t})\n}\n\nfunc (c *HelperClient) loadImageFixture(t testing.TB, name string) *imageFixture {\n\tt.Helper()\n\n\tc.mu.Lock()\n\tdefer c.mu.Unlock()\n\n\tref := fmt.Sprintf(\"cog-test-fixture:%s\", name)\n\n\tif fixture, ok := c.fixtures[ref]; ok {\n\t\treturn fixture\n\t}\n\n\t// Get the path of the current file\n\t_, filename, _, ok := runtime.Caller(0)\n\tif !ok {\n\t\tt.Fatal(\"Could not get current file path\")\n\t}\n\n\t// Get the directory of the current file\n\tdir := filepath.Dir(filename)\n\n\t// Construct the path to the fixture\n\tfixturePath := filepath.Join(dir, \"testdata\", name+\".tar\")\n\n\tt.Logf(\"Loading image fixture %q from %s\", ref, fixturePath)\n\n\tf, err := os.Open(fixturePath)\n\trequire.NoError(t, err, \"Failed to open fixture %q\", name)\n\tdefer f.Close()\n\n\tl, err := c.Client.ImageLoad(t.Context(), f)\n\trequire.NoError(t, err, \"Failed to load fixture %q\", name)\n\tdefer l.Body.Close()\n\t_, err = io.Copy(os.Stderr, l.Body)\n\trequire.NoError(t, err, \"Failed to copy fixture %q\", name)\n\n\tinspect, err := c.Client.ImageInspect(t.Context(), ref)\n\trequire.NoError(t, err, \"Failed to inspect image %q\", ref)\n\n\tfixture := &imageFixture{\n\t\tref:     ref,\n\t\timageID: inspect.ID,\n\t}\n\n\tc.fixtures[ref] = fixture\n\n\treturn fixture\n}\n\ntype imageFixture struct {\n\timageID string\n\tref     string\n}\n"
  },
  {
    "path": "pkg/docker/dockertest/image.go",
    "content": "package dockertest\n\nimport (\n\t\"fmt\"\n\t\"path\"\n\t\"strings\"\n\t\"testing\"\n\t\"time\"\n)\n\n// ImageRef returns an reference based on the unique test name and label.\n// If the label is empty, it will default to \"test-\" followed by the current unix epoch time.\nfunc ImageRef(t *testing.T, label string) string {\n\tif label == \"\" {\n\t\tlabel = fmt.Sprintf(\"test-%d\", time.Now().Unix())\n\t}\n\n\treturn fmt.Sprintf(\"cog-test/%s:%s\", strings.ToLower(t.Name()), label)\n}\n\nfunc ImageRefWithRegistry(t *testing.T, registryAddr string, label string) string {\n\treturn path.Join(registryAddr, ImageRef(t, label))\n}\n"
  },
  {
    "path": "pkg/docker/dockertest/mock_command.go",
    "content": "package dockertest\n\nimport (\n\t\"context\"\n\t\"io\"\n\t\"os\"\n\t\"path/filepath\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/docker/api/types/image\"\n\tdockerspec \"github.com/moby/docker-image-spec/specs-go/v1\"\n\tocispec \"github.com/opencontainers/image-spec/specs-go/v1\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n)\n\nvar PushError error = nil\nvar MockCogConfig string = \"{\\\"build\\\":{\\\"python_version\\\":\\\"3.12\\\",\\\"python_packages\\\":[\\\"torch==2.5.0\\\",\\\"beautifulsoup4==4.12.3\\\"],\\\"system_packages\\\":[\\\"git\\\"]},\\\"image\\\":\\\"test\\\",\\\"predict\\\":\\\"predict.py:Predictor\\\"}\"\nvar MockOpenAPISchema string = \"{}\"\n\ntype MockCommand struct{}\n\nfunc NewMockCommand() *MockCommand {\n\treturn &MockCommand{}\n}\n\nfunc (c *MockCommand) Pull(ctx context.Context, image string, force bool) (*image.InspectResponse, error) {\n\treturn nil, nil\n}\n\nfunc (c *MockCommand) Push(ctx context.Context, image string) error {\n\treturn PushError\n}\n\nfunc (c *MockCommand) LoadUserInformation(ctx context.Context, registryHost string) (*command.UserInfo, error) {\n\tuserInfo := command.UserInfo{\n\t\tToken:    \"test-token\",\n\t\tUsername: \"test-user\",\n\t}\n\treturn &userInfo, nil\n}\n\nfunc (c *MockCommand) CreateTarFile(ctx context.Context, image string, tmpDir string, tarFile string, folder string) (string, error) {\n\tpath := filepath.Join(tmpDir, tarFile)\n\td1 := []byte(\"hello\\ngo\\n\")\n\terr := os.WriteFile(path, d1, 0o644)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn path, nil\n}\n\nfunc (c *MockCommand) CreateAptTarFile(ctx context.Context, tmpDir string, aptTarFile string, packages ...string) (string, error) {\n\tpath := filepath.Join(tmpDir, aptTarFile)\n\td1 := []byte(\"hello\\ngo\\n\")\n\terr := os.WriteFile(path, d1, 0o644)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn path, nil\n}\n\nfunc (c *MockCommand) Inspect(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\tresp := &image.InspectResponse{\n\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tcommand.CogConfigLabelKey:        MockCogConfig,\n\t\t\t\t\tcommand.CogOpenAPISchemaLabelKey: MockOpenAPISchema,\n\t\t\t\t\tcommand.CogVersionLabelKey:       \"0.11.3\",\n\t\t\t\t},\n\t\t\t\tEnv: []string{\n\t\t\t\t\tcommand.R8TorchVersionEnvVarName + \"=2.5.0\",\n\t\t\t\t\tcommand.R8CudaVersionEnvVarName + \"=2.4\",\n\t\t\t\t\tcommand.R8CudnnVersionEnvVarName + \"=1.0\",\n\t\t\t\t\tcommand.R8PythonVersionEnvVarName + \"=3.12\",\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t}\n\n\treturn resp, nil\n}\n\nfunc (c *MockCommand) ImageExists(ctx context.Context, ref string) (bool, error) {\n\tpanic(\"not implemented\")\n}\n\nfunc (c *MockCommand) ContainerLogs(ctx context.Context, containerID string, w io.Writer) error {\n\tpanic(\"not implemented\")\n}\n\nfunc (c *MockCommand) ContainerInspect(ctx context.Context, id string) (*container.InspectResponse, error) {\n\tpanic(\"not implemented\")\n}\n\nfunc (c *MockCommand) ContainerStop(ctx context.Context, containerID string) error {\n\tpanic(\"not implemented\")\n}\n\nfunc (c *MockCommand) RemoveImage(ctx context.Context, ref string) error {\n\tpanic(\"not implemented\")\n}\n\nfunc (c *MockCommand) ImageBuild(ctx context.Context, options command.ImageBuildOptions) (string, error) {\n\tpanic(\"not implemented\")\n}\n\nfunc (c *MockCommand) Run(ctx context.Context, options command.RunOptions) error {\n\tpanic(\"not implemented\")\n}\n\nfunc (c *MockCommand) ContainerStart(ctx context.Context, options command.RunOptions) (string, error) {\n\tpanic(\"not implemented\")\n}\n\nfunc (c *MockCommand) ImageSave(ctx context.Context, imageRef string) (io.ReadCloser, error) {\n\tpanic(\"not implemented\")\n}\n"
  },
  {
    "path": "pkg/docker/dockertest/ref.go",
    "content": "package dockertest\n\nimport (\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\t\"github.com/stretchr/testify/require\"\n)\n\ntype Ref struct {\n\tt   *testing.T\n\tref name.Reference\n}\n\nfunc NewRef(t *testing.T) Ref {\n\tt.Helper()\n\n\trepoName := strings.ToLower(t.Name())\n\t// Replace any characters that aren't valid in a docker image repo name with underscore\n\t// Valid characters are: a-z, 0-9, ., _, -, /\n\trepoName = strings.Map(func(r rune) rune {\n\t\tif (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '.' || r == '_' || r == '-' || r == '/' {\n\t\t\treturn r\n\t\t}\n\t\treturn '_'\n\t}, repoName)\n\n\tref, err := name.ParseReference(repoName, name.WithDefaultRegistry(\"\"))\n\trequire.NoError(t, err, \"Failed to create reference for test\")\n\n\treturn Ref{t: t, ref: ref}\n}\n\nfunc (r Ref) WithTag(tagName string) Ref {\n\ttagRef := r.ref.Context().Tag(tagName)\n\treturn Ref{t: r.t, ref: tagRef}\n}\n\nfunc (r Ref) WithDigest(digest string) Ref {\n\tdigestRef := r.ref.Context().Digest(digest)\n\treturn Ref{t: r.t, ref: digestRef}\n}\n\nfunc (r Ref) WithRegistry(registry string) Ref {\n\treg, err := name.NewRegistry(registry)\n\trequire.NoError(r.t, err, \"Failed to create registry for test\")\n\n\trepo := r.ref.Context()\n\trepo.Registry = reg\n\tvar newRef name.Reference\n\tswitch r.ref.(type) {\n\tcase name.Tag:\n\t\tnewRef = repo.Tag(r.ref.Identifier())\n\tcase name.Digest:\n\t\tnewRef = repo.Digest(r.ref.Identifier())\n\tdefault:\n\t\trequire.Fail(r.t, \"Unsupported reference type\")\n\t}\n\n\treturn Ref{t: r.t, ref: newRef}\n}\n\nfunc (r Ref) WithoutRegistry() Ref {\n\trepo := r.ref.Context()\n\trepo.Registry = name.Registry{}\n\tvar newRef name.Reference\n\tswitch r.ref.(type) {\n\tcase name.Tag:\n\t\tnewRef = repo.Tag(r.ref.Identifier())\n\tcase name.Digest:\n\t\tnewRef = repo.Digest(r.ref.Identifier())\n\tdefault:\n\t\trequire.Fail(r.t, \"Unsupported reference type\")\n\t}\n\n\treturn Ref{t: r.t, ref: newRef}\n}\n\nfunc (r Ref) String() string {\n\treturn r.ref.Name()\n}\n"
  },
  {
    "path": "pkg/docker/dockertest/ref_test.go",
    "content": "package dockertest\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestRef(t *testing.T) {\n\tref := NewRef(t)\n\tassert.Equal(t, \"testref:latest\", ref.String())\n\n\tref = ref.WithTag(\"v2\")\n\tassert.Equal(t, \"testref:v2\", ref.String())\n\n\tref = ref.WithRegistry(\"r8.im\")\n\tassert.Equal(t, \"r8.im/testref:v2\", ref.String())\n\n\tref = ref.WithoutRegistry()\n\tassert.Equal(t, \"testref:v2\", ref.String())\n\n\tref = ref.WithDigest(\"sha256:71859b0c62df47efaeae4f93698b56a8dddafbf041778fd668bbd1ab45a864f8\")\n\tassert.Equal(t, \"testref@sha256:71859b0c62df47efaeae4f93698b56a8dddafbf041778fd668bbd1ab45a864f8\", ref.String())\n}\n"
  },
  {
    "path": "pkg/docker/dockertest/testdata/create-image-fixtures.sh",
    "content": "#!/usr/bin/env bash\nset -euo pipefail\n\nSRC=\"amd64/alpine:3.14\"\nTAG=\"cog-test-fixture:alpine\"\n\necho \"Creating test image fixtures\"\n\ndocker pull $SRC\n\ndocker tag $SRC $TAG\n\ndocker save -o alpine.tar $TAG\n\ndocker rmi $TAG\n\necho \"Test fixtures created\"\n"
  },
  {
    "path": "pkg/docker/env.go",
    "content": "package docker\n\nimport \"os\"\n\nconst DockerCommandEnvVarName = \"R8_DOCKER_COMMAND\"\n\nfunc DockerCommandFromEnvironment() string {\n\tcommand := os.Getenv(DockerCommandEnvVarName)\n\tif command == \"\" {\n\t\tcommand = \"docker\"\n\t}\n\treturn command\n}\n"
  },
  {
    "path": "pkg/docker/errors.go",
    "content": "package docker\n\nimport (\n\t\"errors\"\n\t\"strings\"\n)\n\n// Error messages vary between different backends (dockerd, containerd, podman, orbstack, etc) or even versions of docker.\n// These helpers normalize the check so callers can handle situations without worrying about the underlying implementation.\n// Yes, it's gross, but whattaya gonna do\n\nfunc isTagNotFoundError(err error) bool {\n\tmsg := err.Error()\n\treturn strings.Contains(msg, \"tag does not exist\") ||\n\t\tstrings.Contains(msg, \"An image does not exist locally with the tag\")\n}\n\nfunc isAuthorizationFailedError(err error) bool {\n\tmsg := err.Error()\n\n\t// registry requires auth and none were provided\n\tif strings.Contains(msg, \"no basic auth credentials\") {\n\t\treturn true\n\t}\n\n\t// registry rejected the provided auth\n\tif strings.Contains(msg, \"authorization failed\") ||\n\t\tstrings.Contains(msg, \"401 Unauthorized\") ||\n\t\tstrings.Contains(msg, \"unauthorized: authentication required\") {\n\t\treturn true\n\t}\n\n\treturn false\n}\n\n// isRepositoryNotFoundError checks if the error indicates that the repository\n// doesn't exist on the registry. This typically means the model hasn't been\n// created on Replicate yet.\nfunc isRepositoryNotFoundError(err error) bool {\n\tmsg := err.Error()\n\t// NAME_UNKNOWN is an OCI registry error code meaning \"repository name not known to registry\"\n\treturn strings.Contains(msg, \"NAME_UNKNOWN\")\n}\n\nfunc isMissingDeviceDriverError(err error) bool {\n\tmsg := err.Error()\n\treturn strings.Contains(msg, \"could not select device driver\") ||\n\t\tstrings.Contains(msg, \"nvidia-container-cli: initialization error\")\n}\n\n// isNetworkError checks if the error is a network error. This is janky and intended for use in tests only\nfunc isNetworkError(err error) bool {\n\t// for both CLI and API clients, network errors are wrapped and lose the net.Error interface\n\t// CLI client: wrapped by exec.Command as exec.ExitError\n\t// API client: wrapped by JSON message stream processing\n\t// Sad as it may be, we rely on string matching for common network error messages\n\n\tmsg := err.Error()\n\tnetworkErrorStrings := []string{\n\t\t\"connection refused\",\n\t\t\"connection reset by peer\",\n\t\t\"dial tcp\",\n\t\t\"EOF\",\n\t\t\"no route to host\",\n\t\t\"network is unreachable\",\n\t\t\"server closed\",\n\t}\n\n\tfor _, errStr := range networkErrorStrings {\n\t\tif strings.Contains(msg, errStr) {\n\t\t\treturn true\n\t\t}\n\t}\n\n\t// also check wrapped errors\n\tif unwrapped := errors.Unwrap(err); unwrapped != nil {\n\t\treturn isNetworkError(unwrapped)\n\t}\n\n\treturn false\n}\n"
  },
  {
    "path": "pkg/docker/host.go",
    "content": "package docker\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\n\tdconfig \"github.com/docker/cli/cli/config\"\n\tdctxdocker \"github.com/docker/cli/cli/context/docker\"\n\tdctxstore \"github.com/docker/cli/cli/context/store\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// determineDockerHost returns the host to use for the docker client.\n// It first checks the DOCKER_HOST environment variable, then the docker context, and finally the system default.\nfunc determineDockerHost() (string, error) {\n\t// 1) if DOCKER_HOST is set, use it\n\tif host := os.Getenv(\"DOCKER_HOST\"); host != \"\" {\n\t\tconsole.Debug(\"using docker host from DOCKER_HOST\")\n\n\t\treturn host, nil\n\t}\n\n\t// 2) try to get a host from the docker context. Use DOCKER_CONTEXT if set, otherwise check the current context\n\tif host, contextName, err := dockerHostFromContext(os.Getenv(\"DOCKER_CONTEXT\")); err != nil {\n\t\tconsole.Debugf(\"could not find docker host from context %q: %v\", contextName, err)\n\n\t\t// if DOCKER_CONTEXT was explicitly set, return an error since the user probably expects that context to be used\n\t\tif os.Getenv(\"DOCKER_CONTEXT\") != \"\" {\n\t\t\treturn \"\", err\n\t\t}\n\t} else if host != \"\" {\n\t\tconsole.Debugf(\"using docker host from context %q\", contextName)\n\n\t\treturn host, nil\n\t}\n\n\tconsole.Debug(\"using system default docker host\")\n\n\t// 3) if we couldn't get a host from env or context, fallback to the system default\n\treturn defaultDockerHost, nil\n}\n\nfunc dockerHostFromContext(contextName string) (string, string, error) {\n\tif contextName == \"\" {\n\t\tcf, err := dconfig.Load(dconfig.Dir())\n\t\tif err != nil {\n\t\t\treturn \"\", \"\", fmt.Errorf(\"error loading docker config: %w\", err)\n\t\t}\n\t\tcontextName = cf.CurrentContext\n\t}\n\n\ttypeGetter := func() any { return &dctxdocker.EndpointMeta{} }\n\tstoreConfig := dctxstore.NewConfig(typeGetter, dctxstore.EndpointTypeGetter(dctxdocker.DockerEndpoint, typeGetter))\n\n\tstore := dctxstore.New(dconfig.ContextStoreDir(), storeConfig)\n\tmeta, err := store.GetMetadata(contextName)\n\tif err != nil {\n\t\treturn \"\", contextName, fmt.Errorf(\"error getting metadata for context %q: %w\", contextName, err)\n\t}\n\n\tendpoint, ok := meta.Endpoints[dctxdocker.DockerEndpoint]\n\tif !ok {\n\t\treturn \"\", contextName, fmt.Errorf(\"no docker endpoints found for context %q\", contextName)\n\t}\n\n\tdockerEPMeta, ok := endpoint.(dctxdocker.EndpointMeta)\n\tif !ok {\n\t\treturn \"\", contextName, fmt.Errorf(\"invalid context config: %v\", endpoint)\n\t}\n\n\tif dockerEPMeta.Host == \"\" {\n\t\treturn \"\", contextName, fmt.Errorf(\"no host found for context %q\", contextName)\n\t}\n\n\treturn dockerEPMeta.Host, contextName, nil\n}\n"
  },
  {
    "path": "pkg/docker/host_unix.go",
    "content": "//go:build !windows\n\npackage docker\n\nconst (\n\tdefaultDockerHost = \"unix:///var/run/docker.sock\"\n)\n"
  },
  {
    "path": "pkg/docker/host_windows.go",
    "content": "package docker\n\nconst (\n\tdefaultDockerHost = \"npipe:////.pipe/docker_engine\"\n)\n"
  },
  {
    "path": "pkg/docker/login.go",
    "content": "package docker\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"os/exec\"\n\t\"strings\"\n\n\t\"github.com/docker/cli/cli/config\"\n\t\"github.com/docker/cli/cli/config/configfile\"\n\t\"github.com/docker/cli/cli/config/types\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc SaveLoginToken(ctx context.Context, registryHost string, username string, token string) error {\n\tconf := config.LoadDefaultConfigFile(os.Stderr)\n\tcredsStore := conf.CredentialsStore\n\tif credsStore == \"\" {\n\t\treturn saveAuthToConfig(conf, registryHost, username, token)\n\t}\n\treturn saveAuthToCredentialsStore(ctx, credsStore, registryHost, username, token)\n}\n\nfunc saveAuthToConfig(conf *configfile.ConfigFile, registryHost string, username string, token string) error {\n\t// conf.Save() will base64 encode username and password\n\tconf.AuthConfigs[registryHost] = types.AuthConfig{\n\t\tUsername: username,\n\t\tPassword: token,\n\t}\n\tif err := conf.Save(); err != nil {\n\t\treturn fmt.Errorf(\"Failed to save Docker config.json: %w\", err)\n\t}\n\treturn nil\n}\n\nfunc saveAuthToCredentialsStore(ctx context.Context, credsStore string, registryHost string, username string, token string) error {\n\tbinary := dockerCredentialBinary(credsStore)\n\tinput := CredentialHelperInput{\n\t\tUsername:  username,\n\t\tSecret:    token,\n\t\tServerURL: registryHost,\n\t}\n\tcmd := exec.CommandContext(ctx, binary, \"store\") //nolint:gosec // G702: binary is from Docker config, not user input\n\tcmd.Env = os.Environ()\n\tcmd.Stderr = os.Stderr\n\tstdin, err := cmd.StdinPipe()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Failed to connect stdin to %s: %w\", binary, err)\n\t}\n\tconsole.Debug(\"$ \" + strings.Join(cmd.Args, \" \"))\n\tif err := cmd.Start(); err != nil {\n\t\treturn fmt.Errorf(\"Failed to start %s: %w\", binary, err)\n\t}\n\tif err := json.NewEncoder(stdin).Encode(input); err != nil {\n\t\treturn fmt.Errorf(\"Failed to write to %s: %w\", binary, err)\n\t}\n\tif err := stdin.Close(); err != nil {\n\t\treturn fmt.Errorf(\"Failed to close stdin to %s: %w\", binary, err)\n\t}\n\tif err := cmd.Wait(); err != nil {\n\t\treturn fmt.Errorf(\"Failed to run %s: %w\", binary, err)\n\t}\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/docker/options.go",
    "content": "package docker\n\nimport \"github.com/docker/docker/api/types/registry\"\n\ntype clientOptions struct {\n\tauthConfigs map[string]registry.AuthConfig\n\thost        string\n}\n\ntype Option func(*clientOptions)\n\nfunc WithAuthConfig(authConfig registry.AuthConfig) Option {\n\treturn func(o *clientOptions) {\n\t\to.authConfigs[authConfig.ServerAddress] = authConfig\n\t}\n}\n\nfunc WithHost(host string) Option {\n\treturn func(o *clientOptions) {\n\t\to.host = host\n\t}\n}\n"
  },
  {
    "path": "pkg/docker/progress.go",
    "content": "package docker\n\nimport (\n\t\"encoding/json\"\n\t\"io\"\n\t\"os\"\n\t\"sync\"\n\n\t\"github.com/docker/docker/pkg/jsonmessage\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// ProgressWriter adapts push progress callbacks to Docker's jsonmessage rendering.\n//\n// This uses the same ANSI cursor movement and progress display as `docker push`,\n// which handles terminal resizing correctly: each line is erased and rewritten\n// individually (ESC[2K + cursor up/down per line), rather than relying on a\n// bulk cursor-up count that can desync when lines wrap after a terminal resize.\ntype ProgressWriter struct {\n\tmu   sync.Mutex\n\tpw   *io.PipeWriter\n\tdone chan error\n\tonce sync.Once\n}\n\n// NewProgressWriter creates a ProgressWriter that renders push progress to stderr\n// using Docker's jsonmessage format, matching the output of `docker push`.\nfunc NewProgressWriter() *ProgressWriter {\n\tpr, pw := io.Pipe()\n\tisTTY := console.IsTTY(os.Stderr)\n\tdone := make(chan error, 1)\n\n\tgo func() {\n\t\tdone <- jsonmessage.DisplayJSONMessagesStream(pr, os.Stderr, os.Stderr.Fd(), isTTY, nil)\n\t}()\n\n\treturn &ProgressWriter{\n\t\tpw:   pw,\n\t\tdone: done,\n\t}\n}\n\n// Write sends a progress update for a specific layer/artifact.\n// id is a unique identifier for the item (layer digest, artifact name).\n// status is the current operation (e.g. \"Pushing\").\n// current and total are the byte counts for the progress bar.\nfunc (p *ProgressWriter) Write(id, status string, current, total int64) {\n\tmsg := jsonmessage.JSONMessage{\n\t\tID:     id,\n\t\tStatus: status,\n\t\tProgress: &jsonmessage.JSONProgress{\n\t\t\tCurrent: current,\n\t\t\tTotal:   total,\n\t\t},\n\t}\n\tp.writeMessage(msg)\n}\n\n// WriteStatus sends a status-only message for a specific layer/artifact\n// (no progress bar), e.g. \"Pushed\", \"FAILED\", or retry messages.\nfunc (p *ProgressWriter) WriteStatus(id, status string) {\n\tmsg := jsonmessage.JSONMessage{\n\t\tID:     id,\n\t\tStatus: status,\n\t}\n\tp.writeMessage(msg)\n}\n\nfunc (p *ProgressWriter) writeMessage(msg jsonmessage.JSONMessage) {\n\tp.mu.Lock()\n\tdefer p.mu.Unlock()\n\n\tif p.pw == nil {\n\t\treturn\n\t}\n\n\tdata, err := json.Marshal(msg)\n\tif err != nil {\n\t\treturn\n\t}\n\tdata = append(data, '\\n')\n\t_, _ = p.pw.Write(data)\n}\n\n// Close shuts down the progress display. Safe to call multiple times.\nfunc (p *ProgressWriter) Close() {\n\tp.once.Do(func() {\n\t\tp.mu.Lock()\n\t\tpw := p.pw\n\t\tp.pw = nil\n\t\tp.mu.Unlock()\n\n\t\tif pw != nil {\n\t\t\t_ = pw.Close()\n\t\t\t<-p.done\n\t\t}\n\t})\n}\n"
  },
  {
    "path": "pkg/docker/push.go",
    "content": "package docker\n\nimport (\n\t\"context\"\n\t\"net/http\"\n\t\"time\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n\t\"github.com/replicate/cog/pkg/web\"\n)\n\ntype BuildInfo struct {\n\tBuildTime time.Duration\n\tBuildID   string\n}\n\nfunc Push(ctx context.Context, image string, projectDir string, command command.Command, buildInfo BuildInfo, client *http.Client) error {\n\twebClient := web.NewClient(command, client)\n\n\tif err := webClient.PostPushStart(ctx, buildInfo.BuildID, buildInfo.BuildTime); err != nil {\n\t\tconsole.Warnf(\"Failed to send build timings to server: %v\", err)\n\t}\n\n\treturn StandardPush(ctx, image, command)\n}\n"
  },
  {
    "path": "pkg/docker/run.go",
    "content": "package docker\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n\t\"strconv\"\n\n\t\"github.com/docker/go-connections/nat\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nvar ErrMissingDeviceDriver = errors.New(\"Docker is missing required device driver\")\n\nfunc Run(ctx context.Context, dockerClient command.Command, options command.RunOptions) error {\n\treturn RunWithIO(ctx, dockerClient, options, os.Stdin, os.Stdout, os.Stderr)\n}\n\nfunc RunWithIO(ctx context.Context, dockerClient command.Command, options command.RunOptions, stdin io.Reader, stdout, stderr io.Writer) error {\n\toptions.Stdin = stdin\n\toptions.Stdout = stdout\n\toptions.Stderr = stderr\n\t// TODO[md]: we're gonna stop passing the entire host env to the container by default, if users indeed rely on that behavior we can uncomment this line:\n\t// options.Env = append(os.Environ(), options.Env...)\n\treturn dockerClient.Run(ctx, options)\n}\n\nfunc RunDaemon(ctx context.Context, dockerClient command.Command, options command.RunOptions, stderr io.Writer) (string, error) {\n\toptions.Stderr = stderr\n\treturn dockerClient.ContainerStart(ctx, options)\n}\n\nfunc GetHostPortForContainer(ctx context.Context, dockerCommand command.Command, containerID string, containerPort int) (int, error) {\n\tconsole.Debugf(\"=== DockerCommand.GetPort %s/%d\", containerID, containerPort)\n\n\tinspect, err := dockerCommand.ContainerInspect(ctx, containerID)\n\tif err != nil {\n\t\treturn 0, fmt.Errorf(\"failed to inspect container %q: %w\", containerID, err)\n\t}\n\n\tif inspect.ContainerJSONBase == nil || inspect.State == nil || !inspect.State.Running {\n\t\treturn 0, fmt.Errorf(\"container %s is not running\", containerID)\n\t}\n\n\ttargetPort, err := nat.NewPort(\"tcp\", strconv.Itoa(containerPort))\n\tif err != nil {\n\t\treturn 0, fmt.Errorf(\"failed to create target port: %w\", err)\n\t}\n\n\tif inspect.NetworkSettings == nil || inspect.NetworkSettings.Ports == nil {\n\t\treturn 0, fmt.Errorf(\"container %s does not have expected network configuration\", containerID)\n\t}\n\n\tfor _, portBinding := range inspect.NetworkSettings.Ports[targetPort] {\n\t\t// TODO[md]: this should not be hardcoded since docker may be bound to a different address\n\t\tif portBinding.HostIP != \"0.0.0.0\" {\n\t\t\tcontinue\n\t\t}\n\t\thostPort, err := nat.ParsePort(portBinding.HostPort)\n\t\tif err != nil {\n\t\t\treturn 0, fmt.Errorf(\"failed to parse host port: %w\", err)\n\t\t}\n\t\treturn hostPort, nil\n\t}\n\n\treturn 0, fmt.Errorf(\"container %s does not have a port bound to 0.0.0.0\", containerID)\n}\n"
  },
  {
    "path": "pkg/docker/run_test.go",
    "content": "//nolint:staticcheck // container.NetworkSettingsBase deprecated but Ports field moving to NetworkSettings in docker v29\npackage docker\n\nimport (\n\t\"testing\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/go-connections/nat\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n)\n\nfunc TestGetHostPortForContainer(t *testing.T) {\n\tt.Run(\"WithExposedPort\", func(t *testing.T) {\n\t\ttestClient := dockertest.NewMockCommand2(t)\n\t\ttestClient.EXPECT().ContainerInspect(t.Context(), \"container123\").Return(&container.InspectResponse{\n\t\t\tContainerJSONBase: &container.ContainerJSONBase{\n\t\t\t\tState: &container.State{\n\t\t\t\t\tStatus:  \"running\",\n\t\t\t\t\tRunning: true,\n\t\t\t\t},\n\t\t\t},\n\t\t\tNetworkSettings: &container.NetworkSettings{\n\t\t\t\tNetworkSettingsBase: container.NetworkSettingsBase{\n\t\t\t\t\tPorts: nat.PortMap{\n\t\t\t\t\t\tnat.Port(\"5678/tcp\"): []nat.PortBinding{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tHostIP:   \"0.0.0.0\",\n\t\t\t\t\t\t\t\tHostPort: \"12345\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, nil)\n\n\t\thostPort, err := GetHostPortForContainer(t.Context(), testClient, \"container123\", 5678)\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, 12345, hostPort)\n\t})\n\n\tt.Run(\"WithMultipleExposedPorts\", func(t *testing.T) {\n\t\ttestClient := dockertest.NewMockCommand2(t)\n\t\ttestClient.EXPECT().ContainerInspect(t.Context(), \"container123\").Return(&container.InspectResponse{\n\t\t\tContainerJSONBase: &container.ContainerJSONBase{\n\t\t\t\tState: &container.State{\n\t\t\t\t\tStatus:  \"running\",\n\t\t\t\t\tRunning: true,\n\t\t\t\t},\n\t\t\t},\n\t\t\tNetworkSettings: &container.NetworkSettings{\n\t\t\t\tNetworkSettingsBase: container.NetworkSettingsBase{\n\t\t\t\t\tPorts: nat.PortMap{\n\t\t\t\t\t\tnat.Port(\"5678/tcp\"): []nat.PortBinding{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tHostIP:   \"0.0.0.0\",\n\t\t\t\t\t\t\t\tHostPort: \"12345\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tHostIP:   \"0.0.0.0\",\n\t\t\t\t\t\t\t\tHostPort: \"54321\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, nil)\n\n\t\thostPort, err := GetHostPortForContainer(t.Context(), testClient, \"container123\", 5678)\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, 12345, hostPort)\n\t})\n\n\tt.Run(\"WithExposedPortOnDifferentAddress\", func(t *testing.T) {\n\t\ttestClient := dockertest.NewMockCommand2(t)\n\t\ttestClient.EXPECT().ContainerInspect(t.Context(), \"container123\").Return(&container.InspectResponse{\n\t\t\tContainerJSONBase: &container.ContainerJSONBase{\n\t\t\t\tState: &container.State{\n\t\t\t\t\tStatus:  \"running\",\n\t\t\t\t\tRunning: true,\n\t\t\t\t},\n\t\t\t},\n\t\t\tNetworkSettings: &container.NetworkSettings{\n\t\t\t\tNetworkSettingsBase: container.NetworkSettingsBase{\n\t\t\t\t\tPorts: nat.PortMap{\n\t\t\t\t\t\tnat.Port(\"5678/tcp\"): []nat.PortBinding{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tHostIP:   \"127.0.0.1\",\n\t\t\t\t\t\t\t\tHostPort: \"12345\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, nil)\n\n\t\t_, err := GetHostPortForContainer(t.Context(), testClient, \"container123\", 5678)\n\t\trequire.ErrorContains(t, err, \"does not have a port bound to 0.0.0.0\")\n\t})\n\n\tt.Run(\"WithDifferentPortExposed\", func(t *testing.T) {\n\t\ttestClient := dockertest.NewMockCommand2(t)\n\t\ttestClient.EXPECT().ContainerInspect(t.Context(), \"container123\").Return(&container.InspectResponse{\n\t\t\tContainerJSONBase: &container.ContainerJSONBase{\n\t\t\t\tState: &container.State{\n\t\t\t\t\tStatus:  \"running\",\n\t\t\t\t\tRunning: true,\n\t\t\t\t},\n\t\t\t},\n\t\t\tNetworkSettings: &container.NetworkSettings{\n\t\t\t\tNetworkSettingsBase: container.NetworkSettingsBase{\n\t\t\t\t\tPorts: nat.PortMap{\n\t\t\t\t\t\tnat.Port(\"1234/tcp\"): []nat.PortBinding{\n\t\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\tHostIP:   \"0.0.0.0\",\n\t\t\t\t\t\t\t\tHostPort: \"12345\",\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t}, nil)\n\n\t\t_, err := GetHostPortForContainer(t.Context(), testClient, \"container123\", 5678)\n\t\trequire.ErrorContains(t, err, \"does not have a port bound to 0.0.0.0\")\n\t})\n\n\tt.Run(\"WithNoExposedPort\", func(t *testing.T) {\n\t\ttestClient := dockertest.NewMockCommand2(t)\n\t\ttestClient.EXPECT().ContainerInspect(t.Context(), \"container123\").Return(&container.InspectResponse{\n\t\t\tContainerJSONBase: &container.ContainerJSONBase{\n\t\t\t\tState: &container.State{\n\t\t\t\t\tStatus:  \"running\",\n\t\t\t\t\tRunning: true,\n\t\t\t\t},\n\t\t\t},\n\t\t}, nil)\n\n\t\t_, err := GetHostPortForContainer(t.Context(), testClient, \"container123\", 5678)\n\t\trequire.ErrorContains(t, err, \"does not have expected network configuration\")\n\t})\n\n\tt.Run(\"ContainerNotRunning\", func(t *testing.T) {\n\t\ttestClient := dockertest.NewMockCommand2(t)\n\t\ttestClient.EXPECT().ContainerInspect(t.Context(), \"container123\").Return(&container.InspectResponse{\n\t\t\tContainerJSONBase: &container.ContainerJSONBase{\n\t\t\t\tState: &container.State{\n\t\t\t\t\tStatus: \"dead\",\n\t\t\t\t\tDead:   true,\n\t\t\t\t},\n\t\t\t},\n\t\t}, nil)\n\n\t\t_, err := GetHostPortForContainer(t.Context(), testClient, \"container123\", 5678)\n\t\trequire.ErrorContains(t, err, \"is not running\")\n\t})\n}\n"
  },
  {
    "path": "pkg/docker/standard_push.go",
    "content": "package docker\n\nimport (\n\t\"context\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n)\n\nfunc StandardPush(ctx context.Context, image string, command command.Command) error {\n\treturn command.Push(ctx, image)\n}\n"
  },
  {
    "path": "pkg/docker/standard_push_test.go",
    "content": "package docker\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n)\n\nfunc TestStandardPush(t *testing.T) {\n\tcommand := dockertest.NewMockCommand()\n\tdockertest.PushError = nil\n\terr := StandardPush(t.Context(), \"test\", command)\n\trequire.NoError(t, err)\n}\n"
  },
  {
    "path": "pkg/dockercontext/build_tempdir.go",
    "content": "package dockercontext\n\nimport (\n\t\"os\"\n\t\"path\"\n\t\"time\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n)\n\nfunc CogBuildArtifactsDirPath(dir string) (string, error) {\n\ttmpDir := path.Join(dir, global.CogBuildArtifactsFolder)\n\terr := os.MkdirAll(tmpDir, 0o777)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn tmpDir, nil\n}\n\nfunc CogTempDir(dir string, contextDir string) (string, error) {\n\ttmpDir, err := CogBuildArtifactsDirPath(dir)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn path.Join(tmpDir, \"tmp\", contextDir), nil\n}\n\nfunc BuildCogTempDir(dir string, subDir string) (string, error) {\n\trootTmp, err := CogTempDir(dir, subDir)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tif err := os.MkdirAll(rootTmp, 0o777); err != nil {\n\t\treturn \"\", err\n\t}\n\treturn rootTmp, nil\n}\n\nfunc BuildTempDir(dir string) (string, error) {\n\t// tmpDir ends up being something like dir/.cog/tmp/build20240620123456.000000\n\tnow := time.Now().Format(\"20060102150405.000000\")\n\treturn BuildCogTempDir(dir, \"build\"+now)\n}\n"
  },
  {
    "path": "pkg/dockercontext/build_tempdir_test.go",
    "content": "package dockercontext\n\nimport (\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestBuildCogTempDir(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tcogTmpDir, err := BuildCogTempDir(tmpDir, \"weights\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, filepath.Join(tmpDir, \".cog/tmp/weights\"), cogTmpDir)\n}\n"
  },
  {
    "path": "pkg/dockercontext/directories.go",
    "content": "package dockercontext\n\nimport \"path/filepath\"\n\nconst StandardBuildDirectory = \".\"\n\nconst ContextBuildDir = \"context\"\nconst AptBuildContextName = \"apt\"\nconst RequirementsBuildContextName = \"requirements\"\nconst SrcBuildContextName = \"src\"\n\nvar SrcBuildDir = filepath.Join(ContextBuildDir, \"src\")\nvar AptBuildDir = filepath.Join(ContextBuildDir, \"apt\")\nvar RequirementsBuildDir = filepath.Join(ContextBuildDir, \"requirements\")\n"
  },
  {
    "path": "pkg/dockerfile/base.go",
    "content": "package dockerfile\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"strings\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/version\"\n)\n\nconst MinimumCUDAVersion = \"11.6\"\nconst MinimumPythonVersion = \"3.10\"\nconst MinimumTorchVersion = \"1.13.1\"\nconst CogBaseImageName = \"cog-base\"\n\nvar (\n\tbaseImageSystemPackages = []string{\n\t\t\"build-essential\",\n\t\t\"cmake\",\n\t\t\"curl\",\n\t\t\"ffmpeg\",\n\t\t\"findutils\",\n\t\t\"g++\",\n\t\t\"gcc\",\n\t\t\"git\",\n\t\t\"libavcodec-dev\",\n\t\t\"libcairo2-dev\",\n\t\t\"libfontconfig1\",\n\t\t\"libgirepository1.0-dev\",\n\t\t\"libgl1\",\n\t\t\"libglx-mesa0\",\n\t\t\"libglib2.0-0\",\n\t\t\"libopencv-dev\",\n\t\t\"libsm6\",\n\t\t\"libsndfile1\",\n\t\t\"libssl-dev\",\n\t\t\"libunistring-dev\",\n\t\t\"libxext6\",\n\t\t\"libxrender1\",\n\t\t\"sox\",\n\t\t\"unzip\",\n\t\t\"wget\",\n\t\t\"zip\",\n\t\t\"zstd\",\n\t}\n)\n\ntype CUDAVersion struct {\n\tVersion string `json:\"versions\"`\n}\n\ntype PyTorchVersion struct {\n\tVersion string `json:\"version\"`\n}\n\ntype PythonVersion struct {\n\tVersion string           `json:\"version\"`\n\tPyTorch []PyTorchVersion `json:\"pytorch\"`\n\tCUDA    []CUDAVersion    `json:\"cuda\"`\n}\n\ntype AvailableBaseImageConfigurations struct {\n\tPythonVersions []PythonVersion `json:\"python_versions\"`\n}\n\ntype BaseImageConfiguration struct {\n\tCUDAVersion   string `json:\"cuda_version\" yaml:\"cuda_version\"`\n\tPythonVersion string `json:\"python_version\" yaml:\"python_version\"`\n\tTorchVersion  string `json:\"torch_version\" yaml:\"torch_version\"`\n}\n\ntype BaseImageGenerator struct {\n\tcudaVersion   string\n\tpythonVersion string\n\ttorchVersion  string\n\tcommand       command.Command\n\tclient        registry.Client\n}\n\nfunc (b BaseImageConfiguration) MarshalJSON() ([]byte, error) {\n\ttype Alias BaseImageConfiguration\n\ttype BaseImageConfigWithImageName struct {\n\t\tAlias\n\t\tImageName string `json:\"image_name,omitempty\" yaml:\"image_name,omitempty\"`\n\t\tTag       string `json:\"image_tag,omitempty\" yaml:\"image_tag,omitempty\"`\n\t}\n\n\trawName := BaseImageName(b.CUDAVersion, b.PythonVersion, b.TorchVersion)\n\trawName = strings.TrimPrefix(rawName, global.ReplicateRegistryHost+\"/\")\n\tsplit := strings.Split(rawName, \":\")\n\tif len(split) != 2 {\n\t\treturn nil, fmt.Errorf(\"invalid base image name and tag: %s\", rawName)\n\t}\n\timageName, tag := split[0], split[1]\n\n\talias := &BaseImageConfigWithImageName{\n\t\tAlias:     Alias(b),\n\t\tImageName: imageName,\n\t\tTag:       tag,\n\t}\n\treturn json.Marshal(alias)\n}\n\n// BaseImageConfigurations returns a list of CUDA/Python/Torch versions\nfunc BaseImageConfigurations() []BaseImageConfiguration {\n\tconfigs := []BaseImageConfiguration{}\n\n\t// Assuming that the Torch versions cover all Python and CUDA versions to avoid\n\t// having to hard-code a list of Python versions here.\n\tpythonVersionsSet := make(map[string]bool)\n\tcudaVersionsSet := make(map[string]bool)\n\n\t// Torch configs\n\tfor _, compat := range config.TorchCompatibilityMatrix {\n\t\tfor _, python := range compat.Pythons {\n\t\t\tif !version.GreaterOrEqual(python, MinimumPythonVersion) || !version.GreaterOrEqual(compat.Torch, MinimumTorchVersion) {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tif compat.CUDA == nil {\n\t\t\t\tconfigs = append(configs, BaseImageConfiguration{\n\t\t\t\t\tPythonVersion: python,\n\t\t\t\t\tTorchVersion:  compat.Torch,\n\t\t\t\t})\n\t\t\t} else {\n\t\t\t\tcuda := *compat.CUDA\n\t\t\t\ttorch := compat.Torch\n\t\t\t\tconf := BaseImageConfiguration{\n\t\t\t\t\tCUDAVersion:   cuda,\n\t\t\t\t\tPythonVersion: python,\n\t\t\t\t\tTorchVersion:  torch,\n\t\t\t\t}\n\t\t\t\tif version.GreaterOrEqual(cuda, MinimumCUDAVersion) {\n\t\t\t\t\tconfigs = append(configs, conf)\n\t\t\t\t\tpythonVersionsSet[python] = true\n\t\t\t\t\tcudaVersionsSet[cuda] = true\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\t// Python and CUDA-only configs\n\tfor python := range pythonVersionsSet {\n\t\tfor cuda := range cudaVersionsSet {\n\t\t\tconfigs = append(configs, BaseImageConfiguration{\n\t\t\t\tCUDAVersion:   cuda,\n\t\t\t\tPythonVersion: python,\n\t\t\t})\n\t\t}\n\t}\n\n\t// Python-only configs\n\tfor python := range pythonVersionsSet {\n\t\tconfigs = append(configs, BaseImageConfiguration{\n\t\t\tPythonVersion: python,\n\t\t})\n\t}\n\n\treturn configs\n}\n\nfunc NewBaseImageGenerator(ctx context.Context, client registry.Client, cudaVersion string, pythonVersion string, torchVersion string, command command.Command, generate bool) (*BaseImageGenerator, error) {\n\tvalid, cudaVersion, pythonVersion, torchVersion, err := BaseImageConfigurationExists(ctx, client, cudaVersion, pythonVersion, torchVersion, generate)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif valid {\n\t\treturn &BaseImageGenerator{cudaVersion, pythonVersion, torchVersion, command, client}, nil\n\t}\n\tprintNone := func(s string) string {\n\t\tif s == \"\" {\n\t\t\treturn \"(none)\"\n\t\t}\n\t\treturn s\n\t}\n\treturn nil, fmt.Errorf(\"unsupported base image configuration: CUDA: %s / Python: %s / Torch: %s\", printNone(cudaVersion), printNone(pythonVersion), printNone(torchVersion))\n}\n\nfunc (g *BaseImageGenerator) GenerateDockerfile(ctx context.Context) (string, error) {\n\tconf, err := g.makeConfig()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tgenerator, err := NewGenerator(conf, \"\", \"\", g.command, g.client, false)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tuseCogBaseImage := false\n\tgenerator.SetUseCogBaseImagePtr(&useCogBaseImage)\n\n\tdockerfile, err := generator.GenerateInitialSteps(ctx)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn dockerfile, nil\n}\n\nfunc (g *BaseImageGenerator) makeConfig() (*config.Config, error) {\n\tconf := &config.Config{\n\t\tBuild: &config.Build{\n\t\t\tGPU:            g.cudaVersion != \"\",\n\t\t\tPythonVersion:  g.pythonVersion,\n\t\t\tPythonPackages: g.pythonPackages(),\n\t\t\tRun:            g.runStatements(),\n\t\t\tSystemPackages: baseImageSystemPackages,\n\t\t\tCUDA:           g.cudaVersion,\n\t\t},\n\t}\n\tif err := conf.Complete(\"\"); err != nil {\n\t\treturn nil, err\n\t}\n\treturn conf, nil\n}\n\nfunc (g *BaseImageGenerator) pythonPackages() []string {\n\tif g.torchVersion != \"\" {\n\t\tpkgs := []string{\n\t\t\t\"torch==\" + g.torchVersion,\n\t\t\t\"opencv-python==4.12.0.88\",\n\t\t}\n\n\t\t// Find torchvision compatibility.\n\t\tfor _, compat := range config.TorchCompatibilityMatrix {\n\t\t\tif len(compat.Torchvision) == 0 {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif !version.Matches(g.torchVersion, compat.TorchVersion()) {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tpkgs = append(pkgs, \"torchvision==\"+compat.Torchvision)\n\t\t\tbreak\n\t\t}\n\n\t\t// Find torchaudio compatibility.\n\t\tfor _, compat := range config.TorchCompatibilityMatrix {\n\t\t\tif len(compat.Torchaudio) == 0 {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tif !version.Matches(g.torchVersion, compat.TorchVersion()) {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tpkgs = append(pkgs, \"torchaudio==\"+compat.Torchaudio)\n\t\t\tbreak\n\t\t}\n\n\t\treturn pkgs\n\t}\n\treturn []string{}\n}\n\nfunc (g *BaseImageGenerator) runStatements() []config.RunItem {\n\treturn []config.RunItem{}\n}\n\nfunc baseImageComponentNormalisation(cudaVersion string, pythonVersion string, torchVersion string) (string, string, string) {\n\tcompatibleTorchVersion := \"\"\n\tfor _, conf := range BaseImageConfigurations() {\n\t\t// Check CUDA version compatibility\n\t\tif !isVersionCompatible(conf.CUDAVersion, cudaVersion) {\n\t\t\tcontinue\n\t\t}\n\n\t\t// Check Python version compatibility\n\t\tif !isVersionCompatible(conf.PythonVersion, pythonVersion) {\n\t\t\tcontinue\n\t\t}\n\n\t\t// Check Torch version compatibility\n\t\tif !isVersionCompatible(conf.TorchVersion, torchVersion) {\n\t\t\tcontinue\n\t\t}\n\n\t\tif compatibleTorchVersion == \"\" || version.Greater(conf.TorchVersion, compatibleTorchVersion) {\n\t\t\tcompatibleTorchVersion = version.StripModifier(conf.TorchVersion)\n\t\t}\n\t}\n\n\treturn cudaVersion, pythonVersion, compatibleTorchVersion\n}\n\nfunc BaseImageName(cudaVersion string, pythonVersion string, torchVersion string) string {\n\tcudaVersion, pythonVersion, torchVersion = baseImageComponentNormalisation(cudaVersion, pythonVersion, torchVersion)\n\n\tcomponents := []string{}\n\tif cudaVersion != \"\" {\n\t\tcomponents = append(components, \"cuda\"+version.StripPatch(cudaVersion))\n\t}\n\tif pythonVersion != \"\" {\n\t\tcomponents = append(components, \"python\"+version.StripPatch(pythonVersion))\n\t}\n\tif torchVersion != \"\" {\n\t\tcomponents = append(components, \"torch\"+version.StripModifier(torchVersion))\n\t}\n\n\ttag := strings.Join(components, \"-\")\n\tif tag == \"\" {\n\t\ttag = \"latest\"\n\t}\n\n\treturn global.ReplicateRegistryHost + \"/\" + CogBaseImageName + \":\" + tag\n}\n\nfunc BaseImageConfigurationExists(ctx context.Context, client registry.Client, cudaVersion, pythonVersion, torchVersion string, generate bool) (bool, string, string, string, error) {\n\tcudaVersion, pythonVersion, torchVersion = baseImageComponentNormalisation(cudaVersion, pythonVersion, torchVersion)\n\n\tvalid := false\n\tfor _, conf := range BaseImageConfigurations() {\n\t\t// Check CUDA version compatibility\n\t\tif !isVersionCompatible(conf.CUDAVersion, cudaVersion) {\n\t\t\tcontinue\n\t\t}\n\n\t\t// Check Python version compatibility\n\t\tif !isVersionCompatible(conf.PythonVersion, pythonVersion) {\n\t\t\tcontinue\n\t\t}\n\n\t\t// Check Torch version compatibility\n\t\tif !isVersionCompatible(conf.TorchVersion, torchVersion) {\n\t\t\tcontinue\n\t\t}\n\n\t\tvalid = true\n\t}\n\n\tvar err error\n\tif valid && !generate {\n\t\tvalid, err = client.Exists(ctx, BaseImageName(cudaVersion, pythonVersion, torchVersion))\n\t}\n\n\treturn valid, cudaVersion, pythonVersion, torchVersion, err\n}\n\nfunc isVersionCompatible(confVersion, requestedVersion string) bool {\n\tif confVersion == \"\" || requestedVersion == \"\" {\n\t\treturn confVersion == requestedVersion\n\t}\n\treturn version.Matches(requestedVersion, confVersion)\n}\n"
  },
  {
    "path": "pkg/dockerfile/base_test.go",
    "content": "package dockerfile\n\nimport (\n\t\"reflect\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n\t\"github.com/replicate/cog/pkg/registry/registrytest\"\n)\n\nfunc TestBaseImageName(t *testing.T) {\n\tfor _, tt := range []struct {\n\t\tcuda     string\n\t\tpython   string\n\t\ttorch    string\n\t\texpected string\n\t}{\n\t\t{\"\", \"3.10\", \"\",\n\t\t\t\"r8.im/cog-base:python3.10\"},\n\t\t{\"\", \"3.10\", \"2.1\",\n\t\t\t\"r8.im/cog-base:python3.10-torch2.1.2\"},\n\t\t{\"12.1\", \"3.10\", \"\",\n\t\t\t\"r8.im/cog-base:cuda12.1-python3.10\"},\n\t\t{\"12.1\", \"3.10\", \"2.1\",\n\t\t\t\"r8.im/cog-base:cuda12.1-python3.10-torch2.1.2\"},\n\t\t{\"12.1\", \"3.10\", \"2.1\",\n\t\t\t\"r8.im/cog-base:cuda12.1-python3.10-torch2.1.2\"},\n\t} {\n\t\tactual := BaseImageName(tt.cuda, tt.python, tt.torch)\n\t\trequire.Equal(t, tt.expected, actual)\n\t}\n}\n\nfunc TestGenerateDockerfile(t *testing.T) {\n\tcudaVersion := \"12.1\"\n\tpythonVersion := \"3.10\"\n\ttorchVersion := \"2.1.0\"\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(cudaVersion, pythonVersion, torchVersion))\n\tcommand := dockertest.NewMockCommand()\n\tgenerator, err := NewBaseImageGenerator(\n\t\tt.Context(),\n\t\tclient,\n\t\tcudaVersion,\n\t\tpythonVersion,\n\t\ttorchVersion,\n\t\tcommand,\n\t\tfalse,\n\t)\n\trequire.NoError(t, err)\n\tdockerfile, err := generator.GenerateDockerfile(t.Context())\n\trequire.NoError(t, err)\n\trequire.True(t, strings.Contains(dockerfile, \"FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04\"))\n}\n\nfunc TestBaseImageNameWithVersionModifier(t *testing.T) {\n\tactual := BaseImageName(\"11.8\", \"3.10\", \"2.0.1+cu118\")\n\trequire.Equal(t, \"r8.im/cog-base:cuda11.8-python3.10-torch2.0.1\", actual)\n}\n\nfunc TestBaseImageConfigurationExists(t *testing.T) {\n\tcudaVersion := \"12.1\"\n\tpythonVersion := \"3.10\"\n\ttorchVersion := \"2.3\"\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(cudaVersion, pythonVersion, torchVersion))\n\texists, _, _, torchVersion, err := BaseImageConfigurationExists(t.Context(), client, cudaVersion, pythonVersion, torchVersion, false)\n\trequire.NoError(t, err)\n\trequire.True(t, exists)\n\trequire.Equal(t, \"2.3.1\", torchVersion)\n}\n\nfunc TestBaseImageConfigurationExistsNoTorch(t *testing.T) {\n\tcudaVersion := \"\"\n\tpythonVersion := \"3.12\"\n\ttorchVersion := \"\"\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(cudaVersion, pythonVersion, torchVersion))\n\texists, _, _, _, err := BaseImageConfigurationExists(t.Context(), client, cudaVersion, pythonVersion, torchVersion, false)\n\trequire.NoError(t, err)\n\trequire.True(t, exists)\n}\n\nfunc TestBaseImageConfigurationExistsNoCUDA(t *testing.T) {\n\tcudaVersion := \"\"\n\tpythonVersion := \"3.10\"\n\ttorchVersion := \"2.1\"\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(cudaVersion, pythonVersion, torchVersion))\n\texists, _, _, torchVersion, err := BaseImageConfigurationExists(t.Context(), client, cudaVersion, pythonVersion, torchVersion, false)\n\trequire.NoError(t, err)\n\trequire.True(t, exists)\n\trequire.Equal(t, \"2.1.2\", torchVersion)\n}\n\nfunc TestIsVersionCompatible(t *testing.T) {\n\tcompatible := isVersionCompatible(\"2.3.1+cu121\", \"2.3\")\n\trequire.True(t, compatible)\n}\n\nfunc TestPythonPackages(t *testing.T) {\n\tcudaVersion := \"12.1\"\n\tpythonVersion := \"3.10\"\n\ttorchVersion := \"2.1.0\"\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(cudaVersion, pythonVersion, torchVersion))\n\tgenerator, err := NewBaseImageGenerator(t.Context(), client, cudaVersion, pythonVersion, torchVersion, command, false)\n\trequire.NoError(t, err)\n\tpkgs := generator.pythonPackages()\n\trequire.Truef(t, reflect.DeepEqual(pkgs, []string{\n\t\t\"torch==\" + torchVersion,\n\t\t\"opencv-python==4.12.0.88\",\n\t\t\"torchvision==0.16.0\",\n\t\t\"torchaudio==2.1.0\",\n\t}), \"expected %v\", pkgs)\n}\n\nfunc TestInvalidBaseImage(t *testing.T) {\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\t_, err := NewBaseImageGenerator(t.Context(), client, \"12.78\", \"3.10\", \"2.1.0\", command, false)\n\trequire.Error(t, err)\n}\n\nfunc TestBaseImageConfigurationNoTorchPythonVersionDoesNotExist(t *testing.T) {\n\tclient := registrytest.NewMockRegistryClient()\n\texists, _, _, _, err := BaseImageConfigurationExists(t.Context(), client, \"\", \"3.99\", \"\", false)\n\trequire.NoError(t, err)\n\trequire.False(t, exists)\n}\n"
  },
  {
    "path": "pkg/dockerfile/cacert.go",
    "content": "package dockerfile\n\nimport (\n\t\"encoding/base64\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"strings\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nconst (\n\t// CACertEnvVar is the environment variable that specifies the CA certificate to inject\n\tCACertEnvVar = \"COG_CA_CERT\"\n\n\t// CACertFilename is the filename used for the CA cert in the build context and container\n\tCACertFilename = \"cog-ca-cert.crt\"\n\n\t// CACertContainerPath is where the cert is installed in the container\n\tCACertContainerPath = \"/usr/local/share/ca-certificates/\" + CACertFilename\n\n\t// SystemCertBundle is the path to the system certificate bundle after update-ca-certificates\n\tSystemCertBundle = \"/etc/ssl/certs/ca-certificates.crt\"\n)\n\n// ReadCACert reads the CA certificate from the COG_CA_CERT environment variable.\n// It supports multiple input formats:\n//   - File path: /path/to/cert.crt\n//   - Directory: /path/to/certs/ (concatenates all *.crt and *.pem files)\n//   - Inline PEM: -----BEGIN CERTIFICATE-----...\n//   - Base64-encoded PEM: LS0tLS1CRUdJTi...\n//\n// Returns:\n//   - (nil, nil) if COG_CA_CERT is not set (no-op case)\n//   - (certBytes, nil) if a valid certificate was found\n//   - (nil, error) if the input is invalid\nfunc ReadCACert() ([]byte, error) {\n\tvalue := os.Getenv(CACertEnvVar)\n\tif value == \"\" {\n\t\treturn nil, nil\n\t}\n\n\tvalue = strings.TrimSpace(value)\n\n\t// Check if it's a file path\n\tif info, err := os.Stat(value); err == nil { //nolint:gosec // G703: path from trusted COG_CA_CERT env var\n\t\tif info.IsDir() {\n\t\t\treturn readCACertDirectory(value)\n\t\t}\n\t\treturn readCACertFile(value)\n\t}\n\n\t// Check if it's inline PEM\n\tif strings.HasPrefix(value, \"-----BEGIN\") {\n\t\treturn validatePEM([]byte(value))\n\t}\n\n\t// Try base64 decoding\n\tdecoded, err := base64.StdEncoding.DecodeString(value)\n\tif err == nil && strings.HasPrefix(string(decoded), \"-----BEGIN\") {\n\t\treturn validatePEM(decoded)\n\t}\n\n\treturn nil, fmt.Errorf(\"%s: invalid value - must be a file path, directory, PEM certificate, or base64-encoded PEM\", CACertEnvVar)\n}\n\n// readCACertFile reads a single certificate file\nfunc readCACertFile(path string) ([]byte, error) {\n\tdata, err := os.ReadFile(path) //nolint:gosec // G703: path from trusted COG_CA_CERT env var\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"%s: failed to read file %s: %w\", CACertEnvVar, path, err)\n\t}\n\treturn validatePEM(data)\n}\n\n// readCACertDirectory reads all .crt and .pem files from a directory and concatenates them\nfunc readCACertDirectory(dir string) ([]byte, error) {\n\tvar certs []byte\n\n\tentries, err := os.ReadDir(dir)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"%s: failed to read directory %s: %w\", CACertEnvVar, dir, err)\n\t}\n\n\tfor _, entry := range entries {\n\t\tif entry.IsDir() {\n\t\t\tcontinue\n\t\t}\n\t\text := strings.ToLower(filepath.Ext(entry.Name()))\n\t\tif ext != \".crt\" && ext != \".pem\" {\n\t\t\tcontinue\n\t\t}\n\n\t\tpath := filepath.Join(dir, entry.Name())\n\t\tdata, err := os.ReadFile(path) //nolint:gosec // G703: path from trusted COG_CA_CERT env var directory\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"%s: failed to read file %s: %w\", CACertEnvVar, path, err)\n\t\t}\n\n\t\t// Validate each cert\n\t\tif _, err := validatePEM(data); err != nil {\n\t\t\treturn nil, fmt.Errorf(\"%s: invalid certificate in %s: %w\", CACertEnvVar, path, err)\n\t\t}\n\n\t\tif len(certs) > 0 && !strings.HasSuffix(string(certs), \"\\n\") {\n\t\t\tcerts = append(certs, '\\n')\n\t\t}\n\t\tcerts = append(certs, data...)\n\t}\n\n\tif len(certs) == 0 {\n\t\treturn nil, fmt.Errorf(\"%s: no .crt or .pem files found in directory %s\", CACertEnvVar, dir)\n\t}\n\n\treturn certs, nil\n}\n\n// validatePEM checks that the data looks like a valid PEM certificate\nfunc validatePEM(data []byte) ([]byte, error) {\n\tcontent := strings.TrimSpace(string(data))\n\tif !strings.HasPrefix(content, \"-----BEGIN CERTIFICATE-----\") {\n\t\treturn nil, fmt.Errorf(\"invalid PEM: must start with '-----BEGIN CERTIFICATE-----'\")\n\t}\n\tif !strings.Contains(content, \"-----END CERTIFICATE-----\") {\n\t\treturn nil, fmt.Errorf(\"invalid PEM: must contain '-----END CERTIFICATE-----'\")\n\t}\n\treturn []byte(content + \"\\n\"), nil\n}\n\n// GenerateCACertInstall generates the Dockerfile lines to install a CA certificate.\n// It writes the cert to the build context and returns the Dockerfile lines.\n//\n// The returned lines:\n//  1. COPY the cert to /usr/local/share/ca-certificates/\n//  2. RUN update-ca-certificates\n//  3. Set SSL_CERT_FILE and REQUESTS_CA_BUNDLE env vars\n//\n// Parameters:\n//   - certData: The PEM-encoded certificate data\n//   - writeTemp: Function to write a file to the build context (returns COPY lines and container path)\n//\n// Returns the Dockerfile lines to add, or error\nfunc GenerateCACertInstall(certData []byte, writeTemp func(filename string, contents []byte) ([]string, string, error)) (string, error) {\n\tif len(certData) == 0 {\n\t\treturn \"\", nil\n\t}\n\n\tconsole.Infof(\"Injecting CA certificate from %s\", CACertEnvVar)\n\n\t// Write cert to build context\n\tcopyLines, _, err := writeTemp(CACertFilename, certData)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to write CA certificate to build context: %w\", err)\n\t}\n\n\tlines := []string{}\n\tlines = append(lines, copyLines...)\n\n\t// Copy to system CA directory, update the certificate store, and set env vars.\n\t// Also append the cert directly to the bundle file as a fallback for images\n\t// where update-ca-certificates may not work as expected.\n\tlines = append(lines,\n\t\tfmt.Sprintf(\"RUN cp /tmp/%s %s && update-ca-certificates && cat /tmp/%s >> %s\", CACertFilename, CACertContainerPath, CACertFilename, SystemCertBundle),\n\t\tfmt.Sprintf(\"ENV SSL_CERT_FILE=%s\", SystemCertBundle),\n\t\tfmt.Sprintf(\"ENV REQUESTS_CA_BUNDLE=%s\", SystemCertBundle),\n\t)\n\n\treturn strings.Join(lines, \"\\n\"), nil\n}\n"
  },
  {
    "path": "pkg/dockerfile/cacert_test.go",
    "content": "package dockerfile\n\nimport (\n\t\"encoding/base64\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nconst testCertPEM = `-----BEGIN CERTIFICATE-----\nMIIBkTCB+wIJAKHBfpegPjMCMA0GCSqGSIb3DQEBCwUAMBExDzANBgNVBAMMBlRl\nc3RDQTAeFw0yNDAxMDEwMDAwMDBaFw0yNTAxMDEwMDAwMDBaMBExDzANBgNVBAMM\nBlRlc3RDQTBcMA0GCSqGSIb3DQEBAQUAA0sAMEgCQQC7o96WzE5gvnMXvPjNdXjH\nHwjE7F5Q4X5g5W5P5s5Q5Y5V5y5v5p5o5k5f5d5c5b5a5Z5X5W5U5T5S5R5P5N5L\nAgMBAAGjUzBRMB0GA1UdDgQWBBQExample1234567890ABCDEFGHIJKLMN\nMB8GA1UdIwQYMBaAFBQExample1234567890ABCDEFGHIJKLMNMA8GA1Ud\nEwEB/wQFMAMBAf8wDQYJKoZIhvcNAQELBQADQQBExample1234567890\n-----END CERTIFICATE-----`\n\nconst testCertPEM2 = `-----BEGIN CERTIFICATE-----\nMIIBkTCB+wIJAKHBfpegPjMDMA0GCSqGSIb3DQEBCwUAMBExDzANBgNVBAMMBlRl\nc3RDQTAeFw0yNDAxMDEwMDAwMDBaFw0yNTAxMDEwMDAwMDBaMBExDzANBgNVBAMM\nBlRlc3RDQTBcMA0GCSqGSIb3DQEBAQUAA0sAMEgCQQC8p97XzF6hvoPYvQkOeYkI\nHxkF8G6Q5Y6h6X6Q6Z6W6z6w6q6p6l6g6e6d6c6b6a6Y6X6W6V6U6T6S6R6Q6O6M\nAgMBAAGjUzBRMB0GA1UdDgQWBBQExample2222222222ABCDEFGHIJKLMN\nMB8GA1UdIwQYMBaAFBQExample2222222222ABCDEFGHIJKLMNMA8GA1Ud\nEwEB/wQFMAMBAf8wDQYJKoZIhvcNAQELBQADQQBExample2222222222\n-----END CERTIFICATE-----`\n\nfunc TestReadCACert_NotSet(t *testing.T) {\n\tos.Unsetenv(CACertEnvVar)\n\n\tcert, err := ReadCACert()\n\trequire.NoError(t, err)\n\trequire.Nil(t, cert)\n}\n\nfunc TestReadCACert_FilePath(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tcertPath := filepath.Join(tmpDir, \"test.crt\")\n\trequire.NoError(t, os.WriteFile(certPath, []byte(testCertPEM), 0o644))\n\n\tt.Setenv(CACertEnvVar, certPath)\n\n\tcert, err := ReadCACert()\n\trequire.NoError(t, err)\n\trequire.NotNil(t, cert)\n\trequire.Contains(t, string(cert), \"-----BEGIN CERTIFICATE-----\")\n\trequire.Contains(t, string(cert), \"-----END CERTIFICATE-----\")\n}\n\nfunc TestReadCACert_Directory(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\t// Write two cert files\n\trequire.NoError(t, os.WriteFile(filepath.Join(tmpDir, \"cert1.crt\"), []byte(testCertPEM), 0o644))\n\trequire.NoError(t, os.WriteFile(filepath.Join(tmpDir, \"cert2.pem\"), []byte(testCertPEM2), 0o644))\n\t// Also write a non-cert file that should be ignored\n\trequire.NoError(t, os.WriteFile(filepath.Join(tmpDir, \"readme.txt\"), []byte(\"ignore me\"), 0o644))\n\n\tt.Setenv(CACertEnvVar, tmpDir)\n\n\tcert, err := ReadCACert()\n\trequire.NoError(t, err)\n\trequire.NotNil(t, cert)\n\t// Should contain both certificates\n\trequire.Equal(t, 2, strings.Count(string(cert), \"-----BEGIN CERTIFICATE-----\"))\n\trequire.Equal(t, 2, strings.Count(string(cert), \"-----END CERTIFICATE-----\"))\n}\n\nfunc TestReadCACert_InlinePEM(t *testing.T) {\n\tt.Setenv(CACertEnvVar, testCertPEM)\n\n\tcert, err := ReadCACert()\n\trequire.NoError(t, err)\n\trequire.NotNil(t, cert)\n\trequire.Contains(t, string(cert), \"-----BEGIN CERTIFICATE-----\")\n}\n\nfunc TestReadCACert_Base64EncodedPEM(t *testing.T) {\n\tencoded := base64.StdEncoding.EncodeToString([]byte(testCertPEM))\n\tt.Setenv(CACertEnvVar, encoded)\n\n\tcert, err := ReadCACert()\n\trequire.NoError(t, err)\n\trequire.NotNil(t, cert)\n\trequire.Contains(t, string(cert), \"-----BEGIN CERTIFICATE-----\")\n}\n\nfunc TestReadCACert_InvalidPEM(t *testing.T) {\n\tt.Setenv(CACertEnvVar, \"not a valid certificate\")\n\n\tcert, err := ReadCACert()\n\trequire.Error(t, err)\n\trequire.Nil(t, cert)\n\trequire.Contains(t, err.Error(), \"invalid value\")\n}\n\nfunc TestReadCACert_MissingFile(t *testing.T) {\n\tt.Setenv(CACertEnvVar, \"/nonexistent/path/to/cert.crt\")\n\n\tcert, err := ReadCACert()\n\trequire.Error(t, err)\n\trequire.Nil(t, cert)\n\trequire.Contains(t, err.Error(), \"invalid value\")\n}\n\nfunc TestReadCACert_EmptyDirectory(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tt.Setenv(CACertEnvVar, tmpDir)\n\n\tcert, err := ReadCACert()\n\trequire.Error(t, err)\n\trequire.Nil(t, cert)\n\trequire.Contains(t, err.Error(), \"no .crt or .pem files found\")\n}\n\nfunc TestReadCACert_InvalidFileInDirectory(t *testing.T) {\n\ttmpDir := t.TempDir()\n\t// Write an invalid cert file\n\trequire.NoError(t, os.WriteFile(filepath.Join(tmpDir, \"bad.crt\"), []byte(\"not a cert\"), 0o644))\n\n\tt.Setenv(CACertEnvVar, tmpDir)\n\n\tcert, err := ReadCACert()\n\trequire.Error(t, err)\n\trequire.Nil(t, cert)\n\trequire.Contains(t, err.Error(), \"invalid certificate\")\n}\n\nfunc TestReadCACert_TrimsWhitespace(t *testing.T) {\n\t// Test that whitespace is trimmed from the env var value\n\tt.Setenv(CACertEnvVar, \"  \"+testCertPEM+\"  \\n\")\n\n\tcert, err := ReadCACert()\n\trequire.NoError(t, err)\n\trequire.NotNil(t, cert)\n\trequire.True(t, strings.HasPrefix(string(cert), \"-----BEGIN\"))\n}\n\nfunc TestGenerateCACertInstall(t *testing.T) {\n\tvar writtenFilename string\n\tvar writtenContents []byte\n\n\tmockWriteTemp := func(filename string, contents []byte) ([]string, string, error) {\n\t\twrittenFilename = filename\n\t\twrittenContents = contents\n\t\treturn []string{\"COPY .cog/tmp/\" + filename + \" /tmp/\" + filename}, \"/tmp/\" + filename, nil\n\t}\n\n\tresult, err := GenerateCACertInstall([]byte(testCertPEM), mockWriteTemp)\n\trequire.NoError(t, err)\n\n\t// Check that the cert was written\n\trequire.Equal(t, CACertFilename, writtenFilename)\n\trequire.Contains(t, string(writtenContents), \"-----BEGIN CERTIFICATE-----\")\n\n\t// Check the generated Dockerfile lines\n\trequire.Contains(t, result, \"COPY .cog/tmp/\"+CACertFilename)\n\trequire.Contains(t, result, \"update-ca-certificates\")\n\trequire.Contains(t, result, \"ENV SSL_CERT_FILE=\"+SystemCertBundle)\n\trequire.Contains(t, result, \"ENV REQUESTS_CA_BUNDLE=\"+SystemCertBundle)\n}\n\nfunc TestGenerateCACertInstall_EmptyData(t *testing.T) {\n\tmockWriteTemp := func(filename string, contents []byte) ([]string, string, error) {\n\t\tt.Fatal(\"writeTemp should not be called for empty data\")\n\t\treturn nil, \"\", nil\n\t}\n\n\tresult, err := GenerateCACertInstall(nil, mockWriteTemp)\n\trequire.NoError(t, err)\n\trequire.Empty(t, result)\n\n\tresult, err = GenerateCACertInstall([]byte{}, mockWriteTemp)\n\trequire.NoError(t, err)\n\trequire.Empty(t, result)\n}\n"
  },
  {
    "path": "pkg/dockerfile/env.go",
    "content": "package dockerfile\n\nimport (\n\t\"maps\"\n\t\"slices\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc envLineFromConfig(c *config.Config) (string, error) {\n\tvars := c.ParsedEnvironment()\n\tif len(vars) == 0 {\n\t\treturn \"\", nil\n\t}\n\n\tout := \"ENV\"\n\tfor _, name := range slices.Sorted(maps.Keys(vars)) {\n\t\tout = out + \" \" + name + \"=\" + vars[name]\n\t}\n\tout += \"\\n\"\n\n\treturn out, nil\n}\n"
  },
  {
    "path": "pkg/dockerfile/generator.go",
    "content": "package dockerfile\n\nimport (\n\t\"context\"\n\n\t\"github.com/replicate/cog/pkg/weights\"\n)\n\ntype Generator interface {\n\tGenerateInitialSteps(ctx context.Context) (string, error)\n\tSetUseCogBaseImage(bool)\n\tSetUseCogBaseImagePtr(*bool)\n\tGenerateModelBaseWithSeparateWeights(ctx context.Context, imageName string) (string, string, string, error)\n\tCleanup() error\n\tSetStrip(bool)\n\tSetPrecompile(bool)\n\tSetUseCudaBaseImage(string)\n\tIsUsingCogBaseImage() bool\n\tBaseImage(ctx context.Context) (string, error)\n\tGenerateWeightsManifest(ctx context.Context) (*weights.Manifest, error)\n\tGenerateDockerfileWithoutSeparateWeights(ctx context.Context) (string, error)\n\tGenerateModelBase(ctx context.Context) (string, error)\n\tName() string\n\tBuildDir() (string, error)\n\tBuildContexts() (map[string]string, error)\n}\n"
  },
  {
    "path": "pkg/dockerfile/generator_factory.go",
    "content": "package dockerfile\n\nimport (\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\nfunc NewGenerator(config *config.Config, dir string, configFilename string, command command.Command, client registry.Client, requiresCog bool) (Generator, error) {\n\treturn NewStandardGenerator(config, dir, configFilename, command, client, requiresCog)\n}\n"
  },
  {
    "path": "pkg/dockerfile/generator_factory_test.go",
    "content": "package dockerfile\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n\t\"github.com/replicate/cog/pkg/registry/registrytest\"\n)\n\nfunc TestGeneratorFactoryStandardGenerator(t *testing.T) {\n\tdir := t.TempDir()\n\tbuild := config.Build{\n\t\tPythonPackages: []string{\"torch==2.5.1\"},\n\t}\n\tcfg := config.Config{\n\t\tBuild: &build,\n\t}\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgenerator, err := NewGenerator(&cfg, dir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\trequire.Equal(t, generator.Name(), STANDARD_GENERATOR_NAME)\n}\n"
  },
  {
    "path": "pkg/dockerfile/standard_generator.go",
    "content": "package dockerfile\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"os\"\n\t\"path\"\n\t\"path/filepath\"\n\t\"slices\"\n\t\"strings\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/dockercontext\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/requirements\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n\t\"github.com/replicate/cog/pkg/util/version\"\n\t\"github.com/replicate/cog/pkg/weights\"\n\t\"github.com/replicate/cog/pkg/wheels\"\n)\n\nconst DockerignoreHeader = `# generated by replicate/cog\n__pycache__\n*.pyc\n*.pyo\n*.pyd\n.Python\nenv\npip-log.txt\npip-delete-this-directory.txt\n.tox\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.log\n.git\n.mypy_cache\n.pytest_cache\n.hypothesis\n`\nconst LDConfigCacheBuildCommand = \"RUN find / -type f -name \\\"*python*.so\\\" -printf \\\"%h\\\\n\\\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig\"\nconst StripDebugSymbolsCommand = \"find / -type f -name \\\"*python*.so\\\" -not -name \\\"*cpython*.so\\\" -exec strip -S {} \\\\;\"\nconst CFlags = \"ENV CFLAGS=\\\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\\\"\"\nconst UVVersion = \"0.9.26\"\nconst uvCacheMount = \"--mount=type=cache,target=/root/.cache/uv\"\nconst uvPip = \"uv pip\"\nconst PrecompilePythonCommand = \"RUN find / -type f -name \\\"*.py[co]\\\" -delete && find / -type f -name \\\"*.py\\\" -exec touch -t 197001010000 {} \\\\; && find / -type f -name \\\"*.py\\\" -printf \\\"%h\\\\n\\\" | sort -u | /usr/bin/python3 -m compileall --invalidation-mode timestamp -o 2 -j 0\"\nconst STANDARD_GENERATOR_NAME = \"STANDARD_GENERATOR\"\n\ntype StandardGenerator struct {\n\tConfig         *config.Config\n\tDir            string\n\tConfigFilename string // Base filename like \"cog.yaml\" or \"my-config.yaml\"\n\n\t// these are here to make this type testable\n\tGOOS   string\n\tGOARCH string\n\n\tuseCudaBaseImage bool\n\tuseCogBaseImage  *bool\n\tstrip            bool\n\tprecompile       bool\n\n\t// absolute path to tmpDir, a directory that will be cleaned up\n\ttmpDir string\n\t// tmpDir relative to Dir\n\trelativeTmpDir string\n\n\tfileWalker weights.FileWalker\n\n\tmodelDirs  []string\n\tmodelFiles []string\n\n\tpythonRequirementsContents string\n\tcommand                    command.Command\n\tclient                     registry.Client\n\trequiresCog                bool\n\n\t// Optional overrides for wheel configs (used by tests for deterministic output).\n\t// When nil, auto-detection is used (env var → dist/ → PyPI).\n\tcogWheelConfig    *wheels.WheelConfig\n\tcogletWheelConfig *wheels.WheelConfig\n\n\t// Resolved wheel configs — set once by resolveCogWheelConfigs() and shared\n\t// between filterManagedPackages() (for warnings) and installCog() (for install).\n\tresolvedCogConfig    *wheels.WheelConfig\n\tresolvedCogletConfig *wheels.WheelConfig\n}\n\nfunc NewStandardGenerator(config *config.Config, dir string, configFilename string, command command.Command, client registry.Client, requiresCog bool) (*StandardGenerator, error) {\n\ttmpDir, err := dockercontext.BuildTempDir(dir)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\t// tmpDir, but without dir prefix. This is the path used in the Dockerfile.\n\trelativeTmpDir, err := filepath.Rel(dir, tmpDir)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Default to \"cog.yaml\" if not specified\n\tif configFilename == \"\" {\n\t\tconfigFilename = \"cog.yaml\"\n\t}\n\n\treturn &StandardGenerator{\n\t\tConfig:         config,\n\t\tDir:            dir,\n\t\tConfigFilename: configFilename,\n\t\t// Docker build target is always linux/amd64 (see pkg/docker/buildkit.go).\n\t\t// These must match the container platform, not the host.\n\t\tGOOS:             \"linux\",\n\t\tGOARCH:           \"amd64\",\n\t\ttmpDir:           tmpDir,\n\t\trelativeTmpDir:   relativeTmpDir,\n\t\tfileWalker:       filepath.Walk,\n\t\tuseCudaBaseImage: true,\n\t\tuseCogBaseImage:  nil,\n\t\tstrip:            false,\n\t\tprecompile:       false,\n\t\tcommand:          command,\n\t\tclient:           client,\n\t\trequiresCog:      requiresCog,\n\t}, nil\n}\n\nfunc (g *StandardGenerator) SetUseCudaBaseImage(argumentValue string) {\n\t// \"false\" -> false, \"true\" -> true, \"auto\" -> true, \"asdf\" -> true\n\tg.useCudaBaseImage = argumentValue != \"false\"\n}\n\nfunc (g *StandardGenerator) SetUseCogBaseImage(useCogBaseImage bool) {\n\tg.useCogBaseImage = new(bool)\n\t*g.useCogBaseImage = useCogBaseImage\n}\n\nfunc (g *StandardGenerator) SetUseCogBaseImagePtr(useCogBaseImage *bool) {\n\tg.useCogBaseImage = useCogBaseImage\n}\n\nfunc (g *StandardGenerator) IsUsingCogBaseImage() bool {\n\tuseCogBaseImage := g.useCogBaseImage\n\tif useCogBaseImage != nil {\n\t\treturn *useCogBaseImage\n\t}\n\treturn true\n}\n\nfunc (g *StandardGenerator) SetStrip(strip bool) {\n\tg.strip = strip\n}\n\nfunc (g *StandardGenerator) SetPrecompile(precompile bool) {\n\tg.precompile = precompile\n}\n\nfunc (g *StandardGenerator) GenerateInitialSteps(ctx context.Context) (string, error) {\n\tbaseImage, err := g.BaseImage(ctx)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tinstallPython, err := g.installPython()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\taptInstalls, err := g.aptInstalls()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tenvs, err := g.envVars()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\trunCommands, err := g.runCommands()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tpipInstalls, err := g.pipInstalls()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tinstallCog, err := g.installCog()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tinstallCACert, err := g.installCACert()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tif g.IsUsingCogBaseImage() {\n\t\tsteps := []string{\n\t\t\t\"#syntax=docker/dockerfile:1.4\",\n\t\t\t\"FROM \" + baseImage,\n\t\t\tinstallCACert, // First! Before any network requests (apt, pip, etc.)\n\t\t\tenvs,\n\t\t\taptInstalls,\n\t\t\tg.installUV(),\n\t\t}\n\t\tif installCog != \"\" {\n\t\t\tsteps = append(steps, installCog)\n\t\t}\n\t\tsteps = append(steps, pipInstalls)\n\t\tif g.precompile {\n\t\t\tsteps = append(steps, PrecompilePythonCommand)\n\t\t}\n\t\tsteps = append(steps, runCommands)\n\n\t\treturn joinStringsWithoutLineSpace(steps), nil\n\t}\n\n\t// For the CUDA path, uv is installed inside installPython (after the apt step).\n\t// For all other paths (python:X-slim), install uv after apt.\n\tuvInstall := \"\"\n\tif installPython == \"\" {\n\t\tuvInstall = g.installUV()\n\t}\n\tsteps := []string{\n\t\t\"#syntax=docker/dockerfile:1.4\",\n\t\t\"FROM \" + baseImage,\n\t\tg.preamble(),\n\t\tinstallCACert, // Early! Before tini (uses curl), apt, pip, etc.\n\t\tg.installTini(),\n\t\tenvs,\n\t\taptInstalls,\n\t\tuvInstall,\n\t\tinstallPython,\n\t\tpipInstalls,\n\t\tinstallCog,\n\t}\n\tif g.precompile {\n\t\tsteps = append(steps, PrecompilePythonCommand)\n\t}\n\tsteps = append(steps, LDConfigCacheBuildCommand, runCommands)\n\n\treturn joinStringsWithoutLineSpace(steps), nil\n}\n\nfunc (g *StandardGenerator) GenerateModelBase(ctx context.Context) (string, error) {\n\tinitialSteps, err := g.GenerateInitialSteps(ctx)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tsteps := []string{\n\t\tinitialSteps,\n\t\t`WORKDIR /src`,\n\t\t`EXPOSE 5000`,\n\t}\n\tsteps = append(steps, g.cogEnvVars()...)\n\tsteps = append(steps, `CMD [\"python\", \"-m\", \"cog.server.http\"]`)\n\treturn strings.Join(steps, \"\\n\"), nil\n}\n\n// GenerateDockerfileWithoutSeparateWeights generates a Dockerfile that doesn't write model weights to a separate layer.\nfunc (g *StandardGenerator) GenerateDockerfileWithoutSeparateWeights(ctx context.Context) (string, error) {\n\tbase, err := g.GenerateModelBase(ctx)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tbases := []string{\n\t\tbase,\n\t\t`COPY . /src`,\n\t}\n\tif m := g.cpCogYaml(); m != \"\" {\n\t\tbases = append(bases, m)\n\t}\n\treturn joinStringsWithoutLineSpace(bases), nil\n}\n\n// GenerateModelBaseWithSeparateWeights creates the Dockerfile and .dockerignore file contents for model weights\n// It returns four values:\n// - weightsBase: The base image used for Dockerfile generation for model weights.\n// - dockerfile: A string that represents the Dockerfile content generated by the function.\n// - dockerignoreContents: A string that represents the .dockerignore content.\n// - err: An error object if an error occurred during Dockerfile generation; otherwise nil.\nfunc (g *StandardGenerator) GenerateModelBaseWithSeparateWeights(ctx context.Context, imageName string) (weightsBase string, dockerfile string, dockerignoreContents string, err error) {\n\tweightsBase, g.modelDirs, g.modelFiles, err = g.generateForWeights()\n\tif err != nil {\n\t\treturn \"\", \"\", \"\", fmt.Errorf(\"Failed to generate Dockerfile for model weights files: %w\", err)\n\t}\n\tinitialSteps, err := g.GenerateInitialSteps(ctx)\n\tif err != nil {\n\t\treturn \"\", \"\", \"\", err\n\t}\n\n\t// Inject weights base image into initial steps so we can COPY from it\n\tbase := []string{}\n\tinitialStepsLines := strings.Split(initialSteps, \"\\n\")\n\tfor i, line := range initialStepsLines {\n\t\tif strings.HasPrefix(line, \"FROM \") {\n\t\t\tbase = append(base, fmt.Sprintf(\"FROM %s AS %s\", imageName+\"-weights\", \"weights\"))\n\t\t\tbase = append(base, initialStepsLines[i:]...)\n\t\t\tbreak\n\t\t} else {\n\t\t\tbase = append(base, line)\n\t\t}\n\t}\n\n\tfor _, p := range append(g.modelDirs, g.modelFiles...) {\n\t\tbase = append(base, \"COPY --from=weights --link \"+path.Join(\"/src\", p)+\" \"+path.Join(\"/src\", p))\n\t}\n\n\tbase = append(base,\n\t\t`WORKDIR /src`,\n\t\t`EXPOSE 5000`,\n\t)\n\tbase = append(base, g.cogEnvVars()...)\n\tbase = append(base,\n\t\t`CMD [\"python\", \"-m\", \"cog.server.http\"]`,\n\t\t`COPY . /src`,\n\t)\n\tif m := g.cpCogYaml(); m != \"\" {\n\t\tbase = append(base, m)\n\t}\n\n\tdockerignoreContents = makeDockerignoreForWeights(g.modelDirs, g.modelFiles)\n\treturn weightsBase, joinStringsWithoutLineSpace(base), dockerignoreContents, nil\n}\n\n// cogEnvVars returns ENV lines that pass cog.yaml config to the runtime\n// so the container doesn't need to parse cog.yaml at startup.\nfunc (g *StandardGenerator) cogEnvVars() []string {\n\tvar envs []string\n\tif g.Config.Predict != \"\" {\n\t\tenvs = append(envs, fmt.Sprintf(`ENV COG_PREDICT_TYPE_STUB=\"%s\"`, g.Config.Predict))\n\t}\n\tif g.Config.Train != \"\" {\n\t\tenvs = append(envs, fmt.Sprintf(`ENV COG_TRAIN_TYPE_STUB=\"%s\"`, g.Config.Train))\n\t}\n\tif g.Config.Concurrency != nil && g.Config.Concurrency.Max > 0 {\n\t\tenvs = append(envs, fmt.Sprintf(`ENV COG_MAX_CONCURRENCY=%d`, g.Config.Concurrency.Max))\n\t}\n\treturn envs\n}\n\nfunc (g *StandardGenerator) cpCogYaml() string {\n\tif g.ConfigFilename == \"\" || g.ConfigFilename == \"cog.yaml\" {\n\t\treturn \"\"\n\t}\n\t// Absolute filename doesn't work anyway, so it's always relative\n\treturn fmt.Sprintf(\"RUN cp %s /src/cog.yaml\", filepath.Join(\"/src\", g.ConfigFilename))\n}\n\nfunc (g *StandardGenerator) generateForWeights() (string, []string, []string, error) {\n\tmodelDirs, modelFiles, err := weights.FindWeights(g.fileWalker)\n\tif err != nil {\n\t\treturn \"\", nil, nil, err\n\t}\n\t// generate dockerfile to store these model weights files\n\tvar dockerfileContents strings.Builder\n\tdockerfileContents.WriteString(`#syntax=docker/dockerfile:1.4\nFROM scratch\n`)\n\tfor _, p := range append(modelDirs, modelFiles...) {\n\t\tfmt.Fprintf(&dockerfileContents, \"\\nCOPY %s %s\", p, path.Join(\"/src\", p))\n\t}\n\n\treturn dockerfileContents.String(), modelDirs, modelFiles, nil\n}\n\nfunc makeDockerignoreForWeights(dirs, files []string) string {\n\tvar contents strings.Builder\n\tfor _, p := range dirs {\n\t\tfmt.Fprintf(&contents, \"%[1]s\\n%[1]s/**/*\\n\", p)\n\t}\n\tfor _, p := range files {\n\t\tfmt.Fprintf(&contents, \"%[1]s\\n\", p)\n\t}\n\treturn DockerignoreHeader + contents.String()\n}\n\nfunc (g *StandardGenerator) Cleanup() error {\n\tif err := os.RemoveAll(g.tmpDir); err != nil {\n\t\treturn fmt.Errorf(\"Failed to clean up %s: %w\", g.tmpDir, err)\n\t}\n\treturn nil\n}\n\nfunc (g *StandardGenerator) BaseImage(ctx context.Context) (string, error) {\n\tif g.IsUsingCogBaseImage() {\n\t\tbaseImage, err := g.determineBaseImageName(ctx)\n\t\tif err == nil || g.useCogBaseImage != nil {\n\t\t\treturn baseImage, err\n\t\t}\n\t\tconsole.Warnf(\"Could not find a suitable base image, continuing without base image support (%v).\", err)\n\t\tif g.useCogBaseImage == nil {\n\t\t\tg.useCogBaseImage = new(bool)\n\t\t\t*g.useCogBaseImage = false\n\t\t}\n\t}\n\n\tif g.Config.Build.GPU && g.useCudaBaseImage {\n\t\treturn g.Config.CUDABaseImageTag()\n\t}\n\treturn \"python:\" + g.Config.Build.PythonVersion + \"-slim\", nil\n}\n\nfunc (g *StandardGenerator) Name() string {\n\treturn STANDARD_GENERATOR_NAME\n}\n\nfunc (g *StandardGenerator) BuildDir() (string, error) {\n\treturn dockercontext.StandardBuildDirectory, nil\n}\n\nfunc (g *StandardGenerator) BuildContexts() (map[string]string, error) {\n\treturn map[string]string{}, nil\n}\n\nfunc (g *StandardGenerator) preamble() string {\n\treturn `ENV DEBIAN_FRONTEND=noninteractive\nENV PYTHONUNBUFFERED=1\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin\nENV NVIDIA_DRIVER_CAPABILITIES=all`\n}\n\nfunc (g *StandardGenerator) installTini() string {\n\t// Install tini as the image entrypoint to provide signal handling and process\n\t// reaping appropriate for PID 1.\n\t//\n\t// N.B. If you remove/change this, consider removing/changing the `has_init`\n\t// image label applied in image/build.go.\n\tlines := []string{\n\t\t`RUN --mount=type=cache,target=/var/cache/apt,sharing=locked set -eux; \\\napt-get update -qq && \\\napt-get install -qqy --no-install-recommends curl; \\\nrm -rf /var/lib/apt/lists/*; \\\nTINI_VERSION=v0.19.0; \\\nTINI_ARCH=\"$(dpkg --print-architecture)\"; \\\ncurl -sSL -o /sbin/tini \"https://github.com/krallin/tini/releases/download/${TINI_VERSION}/tini-${TINI_ARCH}\"; \\\nchmod +x /sbin/tini`,\n\t\t`ENTRYPOINT [\"/sbin/tini\", \"--\"]`,\n\t}\n\treturn strings.Join(lines, \"\\n\")\n}\n\nfunc (g *StandardGenerator) aptInstalls() (string, error) {\n\tpackages := g.Config.Build.SystemPackages\n\tif len(packages) == 0 {\n\t\treturn \"\", nil\n\t}\n\n\tif g.IsUsingCogBaseImage() {\n\t\t// Filter out packages that are already in the base image\n\t\tpackages = slices.DeleteFunc(slices.Clone(packages), func(pkg string) bool {\n\t\t\treturn slices.Contains(baseImageSystemPackages, pkg)\n\t\t})\n\t}\n\n\treturn \"RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy \" +\n\t\tstrings.Join(packages, \" \") +\n\t\t\" && rm -rf /var/lib/apt/lists/*\", nil\n}\n\nfunc (g *StandardGenerator) installPython() (string, error) {\n\tif g.Config.Build.GPU && g.useCudaBaseImage && !g.IsUsingCogBaseImage() {\n\t\treturn g.installPythonCUDA()\n\t}\n\treturn \"\", nil\n}\n\nfunc (g *StandardGenerator) installUV() string {\n\treturn `COPY --from=ghcr.io/astral-sh/uv:` + UVVersion + ` /uv /uvx /usr/local/bin/\nENV UV_SYSTEM_PYTHON=true`\n}\n\nfunc (g *StandardGenerator) installPythonCUDA() (string, error) {\n\t// TODO: check that python version is valid\n\n\tpy := g.Config.Build.PythonVersion\n\treturn `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy --no-install-recommends \\\n\twget \\\n\tcurl \\\n\txz-utils \\\n\tgit \\\n\tca-certificates \\\n\t&& rm -rf /var/lib/apt/lists/*\n` + g.installUV() + \"\\n\" + fmt.Sprintf(`RUN uv python install %s && \\\n\tln -sf $(uv python find %s) /usr/bin/python3\nENV UV_PYTHON=%s\nENV PATH=\"/usr/local/bin:$PATH\"`, py, py, py), nil\n}\n\n// resolveCogWheelConfigs resolves and caches the cog and coglet wheel configs.\n// It is idempotent — subsequent calls are no-ops. Must be called before\n// filterManagedPackages() and installCog().\n//\n// Precedence for cog SDK:\n//  1. Test override (cogWheelConfig field)\n//  2. COG_SDK_WHEEL env var\n//  3. build.sdk_version in cog.yaml\n//  4. Auto-detect dist/ (dev builds only)\n//  5. Latest PyPI\n//\n// Precedence for coglet:\n//  1. Test override (cogletWheelConfig field)\n//  2. COGLET_WHEEL env var\n//  3. Auto-detect dist/ (dev builds only)\n//  4. Latest PyPI\nfunc (g *StandardGenerator) resolveCogWheelConfigs() error {\n\tif g.resolvedCogConfig != nil {\n\t\treturn nil // already resolved\n\t}\n\n\tvar err error\n\n\t// Resolve cog SDK\n\tif g.cogWheelConfig != nil {\n\t\tg.resolvedCogConfig = g.cogWheelConfig\n\t} else if envVal := os.Getenv(wheels.CogSDKWheelEnvVar); envVal != \"\" {\n\t\tg.resolvedCogConfig, err = wheels.GetCogWheelConfig()\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t} else if g.Config.Build != nil && g.Config.Build.SDKVersion != \"\" {\n\t\tg.resolvedCogConfig = &wheels.WheelConfig{\n\t\t\tSource:  wheels.WheelSourcePyPI,\n\t\t\tVersion: g.Config.Build.SDKVersion,\n\t\t}\n\t} else {\n\t\tg.resolvedCogConfig, err = wheels.GetCogWheelConfig()\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\t// Validate: refuse versions older than the minimum supported SDK\n\tif err := wheels.ValidateSDKVersion(g.resolvedCogConfig, \"cog\"); err != nil {\n\t\treturn err\n\t}\n\n\t// Resolve coglet\n\tif g.cogletWheelConfig != nil {\n\t\tg.resolvedCogletConfig = g.cogletWheelConfig\n\t} else {\n\t\tg.resolvedCogletConfig, err = wheels.GetCogletWheelConfig(g.GOARCH)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\treturn nil\n}\n\n// cogletMinSDKVersion is the minimum SDK version that supports coglet.\n// Older SDKs use the built-in Python HTTP server and are incompatible with coglet.\nconst cogletMinSDKVersion = \"0.17.0\"\n\n// isLegacySDKVersion returns true if the resolved cog SDK version is explicitly\n// pinned below the minimum version that supports coglet. Returns false for\n// unpinned, non-PyPI, or unparseable versions (assume modern).\nfunc (g *StandardGenerator) isLegacySDKVersion() bool {\n\tcfg := g.resolvedCogConfig\n\tif cfg == nil || cfg.Source != wheels.WheelSourcePyPI || cfg.Version == \"\" {\n\t\treturn false\n\t}\n\tbase := cfg.Version\n\tif m := wheels.BaseVersionRe.FindString(base); m != \"\" {\n\t\tbase = m\n\t}\n\tver, err := version.NewVersion(base)\n\tif err != nil {\n\t\treturn false\n\t}\n\treturn !ver.GreaterOrEqual(version.MustVersion(cogletMinSDKVersion))\n}\n\nfunc (g *StandardGenerator) installCog() (string, error) {\n\t// Do not install Cog in base images\n\tif !g.requiresCog {\n\t\treturn \"\", nil\n\t}\n\n\tif err := g.resolveCogWheelConfigs(); err != nil {\n\t\treturn \"\", err\n\t}\n\twheelConfig := g.resolvedCogConfig\n\n\t// Determine if we need --pre flag (pre-release SDK implies pre-release coglet too)\n\tsdkIsPreRelease := wheelConfig.Source == wheels.WheelSourcePyPI && wheels.IsPreRelease(wheelConfig.Version)\n\n\t// Only install coglet explicitly when there's a specific source:\n\t//   - COGLET_WHEEL env var (explicit override)\n\t//   - Local wheel from dist/ (dev/CI auto-detect)\n\t//   - PyPI with pinned version (e.g. COGLET_WHEEL=pypi:0.17.0)\n\t// Otherwise, let the SDK's own dependency handle it — cog >= 0.17.0 declares\n\t// coglet as a hard dependency, older versions don't install it.\n\t//\n\t// Never install coglet when the SDK is explicitly pinned to < 0.17.0 — those\n\t// versions use the built-in Python HTTP server and are incompatible with coglet.\n\tvar installLines string\n\tcogletConfig := g.resolvedCogletConfig\n\texplicitCoglet := cogletConfig != nil &&\n\t\t(cogletConfig.Source == wheels.WheelSourceFile ||\n\t\t\tcogletConfig.Source == wheels.WheelSourceURL ||\n\t\t\t(cogletConfig.Source == wheels.WheelSourcePyPI && cogletConfig.Version != \"\"))\n\tif explicitCoglet && g.isLegacySDKVersion() {\n\t\tconsole.Info(\"Skipping coglet install for legacy SDK\")\n\t\texplicitCoglet = false\n\t}\n\tif explicitCoglet {\n\t\tswitch cogletConfig.Source {\n\t\tcase wheels.WheelSourcePyPI:\n\t\t\tconsole.Infof(\"Using coglet from PyPI: %s\", cogletConfig.PyPIPackageURL(\"coglet\"))\n\t\tcase wheels.WheelSourceURL:\n\t\t\tconsole.Infof(\"Using coglet wheel from URL: %s\", cogletConfig.URL)\n\t\tcase wheels.WheelSourceFile:\n\t\t\tconsole.Debugf(\"Using local coglet wheel: %s\", cogletConfig.Path)\n\t\t}\n\n\t\tcogletIsPreRelease := sdkIsPreRelease ||\n\t\t\t(cogletConfig.Source == wheels.WheelSourcePyPI && wheels.IsPreRelease(cogletConfig.Version))\n\n\t\tcogletInstall, err := g.installCogletWheel(cogletConfig, cogletIsPreRelease)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"failed to install coglet wheel: %w\", err)\n\t\t}\n\t\tif cogletInstall != \"\" {\n\t\t\tinstallLines = cogletInstall\n\t\t}\n\t}\n\n\t// Install cog SDK\n\tvar cogInstall string\n\tvar err error\n\tswitch wheelConfig.Source {\n\tcase wheels.WheelSourcePyPI:\n\t\tcogInstall, err = g.installCogFromPyPI(wheelConfig, sdkIsPreRelease)\n\tcase wheels.WheelSourceURL:\n\t\tconsole.Infof(\"Using cog wheel from URL: %s\", wheelConfig.URL)\n\t\tcogInstall, err = g.installWheelFromURL(wheelConfig.URL)\n\tcase wheels.WheelSourceFile:\n\t\tconsole.Debugf(\"Using local cog wheel: %s\", wheelConfig.Path)\n\t\tcogInstall, err = g.installWheelFromFile(wheelConfig.Path)\n\tdefault:\n\t\treturn \"\", fmt.Errorf(\"unknown wheel source: %v\", wheelConfig.Source)\n\t}\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tif cogInstall != \"\" {\n\t\tif installLines != \"\" {\n\t\t\tinstallLines += \"\\n\"\n\t\t}\n\t\tinstallLines += cogInstall\n\t}\n\n\treturn installLines, nil\n}\n\n// installCogFromPyPI installs the cog SDK from PyPI.\n// preRelease adds --pre to allow pip to resolve pre-release packages.\nfunc (g *StandardGenerator) installCogFromPyPI(config *wheels.WheelConfig, preRelease bool) (string, error) {\n\tpackageSpec := config.PyPIPackageURL(\"cog\")\n\tflags := \"--no-cache\"\n\tif preRelease {\n\t\tflags = \"--pre \" + flags\n\t}\n\tpipInstallLine := \"RUN \" + uvCacheMount + \" \" + uvPip + \" install \" + flags + \" \" + packageSpec\n\tif g.strip {\n\t\tpipInstallLine += \" && \" + StripDebugSymbolsCommand\n\t}\n\tlines := []string{CFlags, pipInstallLine, \"ENV CFLAGS=\"}\n\treturn strings.Join(lines, \"\\n\"), nil\n}\n\n// installWheelFromURL installs a wheel from a URL (when COG_SDK_WHEEL=https://...)\nfunc (g *StandardGenerator) installWheelFromURL(url string) (string, error) {\n\t// Set coglet env vars if this looks like a coglet wheel\n\tvar envLines []string\n\tif strings.Contains(url, \"coglet\") {\n\t\tif !CheckMajorMinorOnly(g.Config.Build.PythonVersion) {\n\t\t\treturn \"\", fmt.Errorf(\"Python version must be . for coglet\")\n\t\t}\n\t\tenvLines = []string{\n\t\t\t\"ENV R8_COG_VERSION=coglet\",\n\t\t\t\"ENV R8_PYTHON_VERSION=\" + g.Config.Build.PythonVersion,\n\t\t}\n\t}\n\n\t// For coglet URLs, uninstall cog first to avoid conflicts with coglet's cog shim package.\n\t// Some base images (e.g. r8.im/cog-base) have cog pre-installed, which conflicts\n\t// with coglet's cog compatibility shim that provides the same module paths.\n\tvar pipPrefix string\n\tif strings.Contains(url, \"coglet\") {\n\t\tpipPrefix = uvPip + \" uninstall cog 2>/dev/null || true && \"\n\t}\n\tpipInstallLine := \"RUN \" + uvCacheMount + \" \" + pipPrefix + uvPip + \" install --no-cache \" + url\n\tif g.strip {\n\t\tpipInstallLine += \" && \" + StripDebugSymbolsCommand\n\t}\n\n\tenvLines = append(envLines, CFlags, pipInstallLine, \"ENV CFLAGS=\")\n\treturn strings.Join(envLines, \"\\n\"), nil\n}\n\n// installWheelFromFile installs a wheel from a local file (when COG_SDK_WHEEL=/path/to/file.whl)\nfunc (g *StandardGenerator) installWheelFromFile(path string) (string, error) {\n\t// Read the local wheel file\n\tdata, err := os.ReadFile(path)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to read wheel file %s: %w\", path, err)\n\t}\n\n\tfilename := filepath.Base(path)\n\tlines, containerPath, err := g.writeTemp(filename, data)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\t// Set coglet env vars if this looks like a coglet wheel\n\tvar pipPrefix string\n\tif strings.Contains(filename, \"coglet\") {\n\t\tif !CheckMajorMinorOnly(g.Config.Build.PythonVersion) {\n\t\t\treturn \"\", fmt.Errorf(\"Python version must be . for coglet\")\n\t\t}\n\t\tlines = append(lines,\n\t\t\t\"ENV R8_COG_VERSION=coglet\",\n\t\t\t\"ENV R8_PYTHON_VERSION=\"+g.Config.Build.PythonVersion,\n\t\t)\n\t\t// Uninstall cog first to avoid conflicts with coglet's cog shim package.\n\t\t// Some base images (e.g. r8.im/cog-base) have cog pre-installed, which conflicts\n\t\t// with coglet's cog compatibility shim that provides the same module paths.\n\t\tpipPrefix = uvPip + \" uninstall cog 2>/dev/null || true && \"\n\t}\n\n\tpipInstallLine := \"RUN \" + uvCacheMount + \" \" + pipPrefix + uvPip + \" install --no-cache \" + containerPath\n\tif g.strip {\n\t\tpipInstallLine += \" && \" + StripDebugSymbolsCommand\n\t}\n\tlines = append(lines, CFlags, pipInstallLine, \"ENV CFLAGS=\")\n\treturn strings.Join(lines, \"\\n\"), nil\n}\n\n// installCogletWheel installs the coglet wheel based on the provided config.\n// preRelease adds --pre to allow pip to resolve pre-release packages.\nfunc (g *StandardGenerator) installCogletWheel(config *wheels.WheelConfig, preRelease bool) (string, error) {\n\tswitch config.Source {\n\tcase wheels.WheelSourcePyPI:\n\t\treturn g.installCogletFromPyPI(config, preRelease)\n\tcase wheels.WheelSourceURL:\n\t\treturn g.installCogletFromURL(config.URL)\n\tcase wheels.WheelSourceFile:\n\t\treturn g.installCogletFromFile(config.Path)\n\tdefault:\n\t\treturn \"\", fmt.Errorf(\"unknown coglet wheel source: %v\", config.Source)\n\t}\n}\n\n// installCogletFromPyPI installs coglet from PyPI.\n// preRelease adds --pre to allow pip to resolve pre-release packages.\nfunc (g *StandardGenerator) installCogletFromPyPI(config *wheels.WheelConfig, preRelease bool) (string, error) {\n\tpackageSpec := config.PyPIPackageURL(\"coglet\")\n\tflags := \"--no-cache\"\n\tif preRelease {\n\t\tflags = \"--pre \" + flags\n\t}\n\tpipInstallLine := \"RUN \" + uvCacheMount + \" \" + uvPip + \" install \" + flags + \" \" + packageSpec\n\tif g.strip {\n\t\tpipInstallLine += \" && \" + StripDebugSymbolsCommand\n\t}\n\tlines := []string{CFlags, pipInstallLine, \"ENV CFLAGS=\"}\n\treturn strings.Join(lines, \"\\n\"), nil\n}\n\n// installCogletFromURL installs coglet from a URL\nfunc (g *StandardGenerator) installCogletFromURL(url string) (string, error) {\n\tpipInstallLine := \"RUN \" + uvCacheMount + \" \" + uvPip + \" install --no-cache \" + url\n\tif g.strip {\n\t\tpipInstallLine += \" && \" + StripDebugSymbolsCommand\n\t}\n\tlines := []string{CFlags, pipInstallLine, \"ENV CFLAGS=\"}\n\treturn strings.Join(lines, \"\\n\"), nil\n}\n\n// installCogletFromFile installs coglet from a local wheel file\nfunc (g *StandardGenerator) installCogletFromFile(path string) (string, error) {\n\tdata, err := os.ReadFile(path)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to read coglet wheel %s: %w\", path, err)\n\t}\n\n\tfilename := filepath.Base(path)\n\tlines, containerPath, err := g.writeTemp(filename, data)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tpipInstallLine := \"RUN \" + uvCacheMount + \" \" + uvPip + \" install --no-cache \" + containerPath\n\tif g.strip {\n\t\tpipInstallLine += \" && \" + StripDebugSymbolsCommand\n\t}\n\n\tlines = append(lines, CFlags, pipInstallLine, \"ENV CFLAGS=\")\n\treturn strings.Join(lines, \"\\n\"), nil\n}\n\n// filterManagedPackages strips cog and coglet from user requirements content,\n// warning loudly for each occurrence. requirements.txt is not the intended\n// mechanism for controlling cog/coglet versions — use build.sdk_version in\n// cog.yaml or the COG_SDK_WHEEL / COGLET_WHEEL environment variables instead.\nfunc (g *StandardGenerator) filterManagedPackages(reqContents string) string {\n\t// Build a human-readable description of what the build system will install\n\t// for each managed package, so the warning is actionable.\n\toverride := func(pkg string) string {\n\t\tvar cfg *wheels.WheelConfig\n\t\tif pkg == \"cog\" {\n\t\t\tcfg = g.resolvedCogConfig\n\t\t} else {\n\t\t\tcfg = g.resolvedCogletConfig\n\t\t}\n\t\tif cfg == nil {\n\t\t\treturn \"latest from PyPI\"\n\t\t}\n\t\tswitch cfg.Source {\n\t\tcase wheels.WheelSourcePyPI:\n\t\t\tif cfg.Version != \"\" {\n\t\t\t\treturn fmt.Sprintf(\"%s==%s from PyPI\", pkg, cfg.Version)\n\t\t\t}\n\t\t\treturn \"latest \" + pkg + \" from PyPI\"\n\t\tcase wheels.WheelSourceURL:\n\t\t\treturn cfg.URL\n\t\tcase wheels.WheelSourceFile:\n\t\t\treturn cfg.Path\n\t\tdefault:\n\t\t\treturn \"unknown source\"\n\t\t}\n\t}\n\n\tmanaged := map[string]bool{\"cog\": true, \"coglet\": true}\n\tvar filtered []string\n\tfor line := range strings.SplitSeq(reqContents, \"\\n\") {\n\t\ttrimmed := strings.TrimSpace(line)\n\t\tif trimmed == \"\" || strings.HasPrefix(trimmed, \"#\") || strings.HasPrefix(trimmed, \"-\") {\n\t\t\tfiltered = append(filtered, line)\n\t\t\tcontinue\n\t\t}\n\t\tpkgName := requirements.PackageName(trimmed)\n\t\tbaseName := strings.ToLower(strings.Split(pkgName, \"[\")[0])\n\t\tif managed[baseName] {\n\t\t\tconsole.Warnf(\n\t\t\t\t\"'%s' found in requirements — overriding with %s. \"+\n\t\t\t\t\t\"Remove it from requirements and use build.sdk_version in cog.yaml or %s to control the version.\",\n\t\t\t\ttrimmed,\n\t\t\t\toverride(baseName),\n\t\t\t\tmap[string]string{\"cog\": \"COG_SDK_WHEEL\", \"coglet\": \"COGLET_WHEEL\"}[baseName],\n\t\t\t)\n\t\t\tcontinue\n\t\t}\n\t\tfiltered = append(filtered, line)\n\t}\n\treturn strings.Join(filtered, \"\\n\")\n}\n\nfunc (g *StandardGenerator) pipInstalls() (string, error) {\n\t// Resolve wheel configs early so filterManagedPackages can emit precise warnings.\n\tif g.requiresCog {\n\t\tif err := g.resolveCogWheelConfigs(); err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t}\n\n\tvar err error\n\tincludePackages := []string{}\n\tif torchVersion, ok := g.Config.TorchVersion(); ok {\n\t\tincludePackages = []string{\"torch==\" + torchVersion}\n\t}\n\tif torchvisionVersion, ok := g.Config.TorchvisionVersion(); ok {\n\t\tincludePackages = append(includePackages, \"torchvision==\"+torchvisionVersion)\n\t}\n\tif torchaudioVersion, ok := g.Config.TorchaudioVersion(); ok {\n\t\tincludePackages = append(includePackages, \"torchaudio==\"+torchaudioVersion)\n\t}\n\tif tensorflowVersion, ok := g.Config.TensorFlowVersion(); ok {\n\t\tincludePackages = append(includePackages, \"tensorflow==\"+tensorflowVersion)\n\t}\n\tg.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, includePackages)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\t// Strip cog/coglet from user requirements — we always install them ourselves\n\t// via installCog(). Leaving them in would cause pip to overwrite our version.\n\tg.pythonRequirementsContents = g.filterManagedPackages(g.pythonRequirementsContents)\n\n\tif strings.Trim(g.pythonRequirementsContents, \"\") == \"\" {\n\t\treturn \"\", nil\n\t}\n\n\tconsole.Debugf(\"Generated requirements.txt:\\n%s\", g.pythonRequirementsContents)\n\tcopyLine, containerPath, err := g.writeTemp(\"requirements.txt\", []byte(g.pythonRequirementsContents))\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tpipInstallLine := \"RUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r \" + containerPath\n\tif g.strip {\n\t\tpipInstallLine += \" && \" + StripDebugSymbolsCommand\n\t}\n\treturn strings.Join([]string{\n\t\tcopyLine[0],\n\t\tCFlags,\n\t\tpipInstallLine,\n\t\t\"ENV CFLAGS=\",\n\t}, \"\\n\"), nil\n}\n\nfunc (g *StandardGenerator) runCommands() (string, error) {\n\trunCommands := g.Config.Build.Run\n\n\t// For backwards compatibility\n\tfor _, command := range g.Config.Build.PreInstall {\n\t\trunCommands = append(runCommands, config.RunItem{Command: command})\n\t}\n\n\tlines := []string{}\n\tfor _, run := range runCommands {\n\t\tcommand := strings.TrimSpace(run.Command)\n\t\tif strings.Contains(command, \"\\n\") {\n\t\t\treturn \"\", fmt.Errorf(`One of the commands in 'run' contains a new line, which won't work. You need to create a new list item in YAML prefixed with '-' for each command.\n\nThis is the offending line: %s`, command)\n\t\t}\n\n\t\tif len(run.Mounts) > 0 {\n\t\t\tmounts := []string{}\n\t\t\tfor _, mount := range run.Mounts {\n\t\t\t\tif mount.Type == \"secret\" {\n\t\t\t\t\tsecretMount := fmt.Sprintf(\"--mount=type=secret,id=%s,target=%s\", mount.ID, mount.Target)\n\t\t\t\t\tmounts = append(mounts, secretMount)\n\t\t\t\t}\n\t\t\t}\n\t\t\tlines = append(lines, fmt.Sprintf(\"RUN %s %s\", strings.Join(mounts, \" \"), command))\n\t\t} else {\n\t\t\tlines = append(lines, \"RUN \"+command)\n\t\t}\n\t}\n\treturn strings.Join(lines, \"\\n\"), nil\n}\n\nfunc (g *StandardGenerator) envVars() (string, error) {\n\treturn envLineFromConfig(g.Config)\n}\n\n// installCACert generates Dockerfile lines to install a custom CA certificate.\n// If COG_CA_CERT is not set, returns empty string (no-op).\nfunc (g *StandardGenerator) installCACert() (string, error) {\n\tcertData, err := ReadCACert()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tif certData == nil {\n\t\treturn \"\", nil\n\t}\n\treturn GenerateCACertInstall(certData, g.writeTemp)\n}\n\n// writeTemp writes a temporary file that can be used as part of the build process\n// It returns the lines to add to Dockerfile to make it available and the filename it ends up as inside the container\nfunc (g *StandardGenerator) writeTemp(filename string, contents []byte) ([]string, string, error) {\n\tpath := filepath.Join(g.tmpDir, filename)\n\tif err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {\n\t\treturn []string{}, \"\", fmt.Errorf(\"Failed to write %s: %w\", filename, err)\n\t}\n\tif err := os.WriteFile(path, contents, 0o644); err != nil {\n\t\treturn []string{}, \"\", fmt.Errorf(\"Failed to write %s: %w\", filename, err)\n\t}\n\treturn []string{fmt.Sprintf(\"COPY %s /tmp/%s\", filepath.Join(g.relativeTmpDir, filename), filename)}, \"/tmp/\" + filename, nil\n}\n\nfunc joinStringsWithoutLineSpace(chunks []string) string {\n\tlines := []string{}\n\tfor _, chunk := range chunks {\n\t\tchunkLines := strings.Split(chunk, \"\\n\")\n\t\tlines = append(lines, chunkLines...)\n\t}\n\treturn strings.Join(filterEmpty(lines), \"\\n\")\n}\n\nfunc filterEmpty(list []string) []string {\n\tfiltered := []string{}\n\tfor _, s := range list {\n\t\tif s != \"\" {\n\t\t\tfiltered = append(filtered, s)\n\t\t}\n\t}\n\treturn filtered\n}\n\nfunc (g *StandardGenerator) GenerateWeightsManifest(ctx context.Context) (*weights.Manifest, error) {\n\tm := weights.NewManifest()\n\n\tfor _, dir := range g.modelDirs {\n\t\terr := g.fileWalker(dir, func(path string, info os.FileInfo, err error) error {\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\tif info.IsDir() {\n\t\t\t\treturn nil\n\t\t\t}\n\n\t\t\treturn m.AddFile(path)\n\t\t})\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\tfor _, path := range g.modelFiles {\n\t\terr := m.AddFile(path)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t}\n\n\treturn m, nil\n}\n\nfunc (g *StandardGenerator) determineBaseImageName(ctx context.Context) (string, error) {\n\tvar changed bool\n\tvar err error\n\n\tcudaVersion := g.Config.Build.CUDA\n\n\tpythonVersion := g.Config.Build.PythonVersion\n\tpythonVersion, changed, err = stripPatchVersion(pythonVersion)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tif changed {\n\t\tconsole.Warnf(\"Stripping patch version from Python version %s to %s\", g.Config.Build.PythonVersion, pythonVersion)\n\t}\n\n\ttorchVersion, _ := g.Config.TorchVersion()\n\n\t// validate that the base image configuration exists\n\timageGenerator, err := NewBaseImageGenerator(ctx, g.client, cudaVersion, pythonVersion, torchVersion, g.command, false)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tbaseImage := BaseImageName(imageGenerator.cudaVersion, imageGenerator.pythonVersion, imageGenerator.torchVersion)\n\treturn baseImage, nil\n}\n\nfunc stripPatchVersion(versionString string) (string, bool, error) {\n\tif versionString == \"\" {\n\t\treturn \"\", false, nil\n\t}\n\n\tv, err := version.NewVersion(versionString)\n\tif err != nil {\n\t\treturn \"\", false, fmt.Errorf(\"Invalid version: %s\", versionString)\n\t}\n\n\tstrippedVersion := fmt.Sprintf(\"%d.%d\", v.Major, v.Minor)\n\tchanged := strippedVersion != versionString\n\n\treturn strippedVersion, changed, nil\n}\n"
  },
  {
    "path": "pkg/dockerfile/standard_generator_test.go",
    "content": "package dockerfile\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\t\"path\"\n\t\"path/filepath\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n\t\"github.com/replicate/cog/pkg/registry/registrytest\"\n\t\"github.com/replicate/cog/pkg/wheels\"\n)\n\nfunc testTini() string {\n\treturn `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked set -eux; \\\napt-get update -qq && \\\napt-get install -qqy --no-install-recommends curl; \\\nrm -rf /var/lib/apt/lists/*; \\\nTINI_VERSION=v0.19.0; \\\nTINI_ARCH=\"$(dpkg --print-architecture)\"; \\\ncurl -sSL -o /sbin/tini \"https://github.com/krallin/tini/releases/download/${TINI_VERSION}/tini-${TINI_ARCH}\"; \\\nchmod +x /sbin/tini\nENTRYPOINT [\"/sbin/tini\", \"--\"]\n`\n}\n\nvar testInstallUVLine = \"COPY --from=ghcr.io/astral-sh/uv:\" + UVVersion + \" /uv /uvx /usr/local/bin/\\nENV UV_SYSTEM_PYTHON=true\"\n\nfunc testInstallCog(stripped bool) string {\n\tstrippedCall := \"\"\n\tif stripped {\n\t\tstrippedCall += \" && find / -type f -name \\\"*python*.so\\\" -not -name \\\"*cpython*.so\\\" -exec strip -S {} \\\\;\"\n\t}\n\t// When coglet has no explicit version pin (empty version via pypiWheels()),\n\t// the SDK's own dependency handles coglet installation — no explicit coglet line.\n\treturn fmt.Sprintf(`ENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/uv uv pip install --no-cache cog%s\nENV CFLAGS=`, strippedCall)\n}\n\n// pypiWheels sets the generator to use unpinned PyPI for both cog and coglet,\n// giving deterministic Dockerfile output regardless of local dist/ contents.\nfunc pypiWheels(gen *StandardGenerator) {\n\tgen.cogWheelConfig = &wheels.WheelConfig{Source: wheels.WheelSourcePyPI}\n\tgen.cogletWheelConfig = &wheels.WheelConfig{Source: wheels.WheelSourcePyPI}\n}\n\nfunc testInstallPython(version string) string {\n\treturn fmt.Sprintf(`RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy --no-install-recommends \\\n\twget \\\n\tcurl \\\n\txz-utils \\\n\tgit \\\n\tca-certificates \\\n\t&& rm -rf /var/lib/apt/lists/*\nCOPY --from=ghcr.io/astral-sh/uv:`+UVVersion+` /uv /uvx /usr/local/bin/\nENV UV_SYSTEM_PYTHON=true\nRUN uv python install %s && \\\n\tln -sf $(uv python find %s) /usr/bin/python3\nENV UV_PYTHON=%s\nENV PATH=\"/usr/local/bin:$PATH\"\n`, version, version, version)\n}\n\nfunc TestGenerateEmptyCPU(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  gpu: false\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM python:3.12-slim\nENV DEBIAN_FRONTEND=noninteractive\nENV PYTHONUNBUFFERED=1\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin\nENV NVIDIA_DRIVER_CAPABILITIES=all\n` + testTini() + testInstallUVLine + \"\\n\" + testInstallCog(gen.strip) + `\nRUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n}\n\nfunc TestGenerateEmptyGPU(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  gpu: true\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04\nENV DEBIAN_FRONTEND=noninteractive\nENV PYTHONUNBUFFERED=1\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin\nENV NVIDIA_DRIVER_CAPABILITIES=all\n` + testTini() + testInstallPython(\"3.12\") + testInstallCog(gen.strip) + `\nRUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n}\n\nfunc TestGenerateFullCPU(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  gpu: false\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==2.3.0\n    - pandas==1.2.0.12\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM python:3.12-slim\nENV DEBIAN_FRONTEND=noninteractive\nENV PYTHONUNBUFFERED=1\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin\nENV NVIDIA_DRIVER_CAPABILITIES=all\n` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*\n` + testInstallUVLine + `\nCOPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt\nENV CFLAGS=\n` + testInstallCog(gen.strip) + `\nRUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\trequire.Equal(t, expected, actual)\n\n\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\trequire.NoError(t, err)\n\n\trequire.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cpu\ntorch==2.3.0\npandas==1.2.0.12`, string(requirements))\n}\n\nfunc TestGenerateFullGPU(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  gpu: true\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==2.0.1\n    - pandas==2.0.3\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04\nENV DEBIAN_FRONTEND=noninteractive\nENV PYTHONUNBUFFERED=1\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin\nENV NVIDIA_DRIVER_CAPABILITIES=all\n` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*\n` + testInstallPython(\"3.12\") + `COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt\nENV CFLAGS=\n` + testInstallCog(gen.strip) + `\nRUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n\n\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu118\ntorch==2.0.1\npandas==2.0.3`, string(requirements))\n}\n\n// pre_install is deprecated but supported for backwards compatibility\nfunc TestPreInstall(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  system_packages:\n    - cowsay\n  pre_install:\n    - \"cowsay moo\"\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM python:3.12-slim\nENV DEBIAN_FRONTEND=noninteractive\nENV PYTHONUNBUFFERED=1\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin\nENV NVIDIA_DRIVER_CAPABILITIES=all\n` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*\n` + testInstallUVLine + `\n` + testInstallCog(gen.strip) + `\nRUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\trequire.Equal(t, expected, actual)\n\n}\n\nfunc TestPythonRequirements(t *testing.T) {\n\ttmpDir := t.TempDir()\n\terr := os.WriteFile(path.Join(tmpDir, \"my-requirements.txt\"), []byte(\"torch==1.0.0\"), 0o644)\n\trequire.NoError(t, err)\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  python_requirements: \"my-requirements.txt\"\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(tmpDir))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\tfmt.Println(actual)\n\trequire.Contains(t, actual, `uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt`)\n}\n\n// mockFileInfo is a test type to mock os.FileInfo\ntype mockFileInfo struct {\n\tsize int64\n}\n\nfunc (mfi mockFileInfo) Size() int64 {\n\treturn mfi.size\n}\nfunc (mfi mockFileInfo) Name() string {\n\treturn \"\"\n}\nfunc (mfi mockFileInfo) Mode() os.FileMode {\n\treturn 0\n}\nfunc (mfi mockFileInfo) ModTime() time.Time {\n\treturn time.Time{}\n}\nfunc (mfi mockFileInfo) IsDir() bool {\n\treturn false\n}\nfunc (mfi mockFileInfo) Sys() any {\n\treturn nil\n}\n\nconst sizeThreshold = 10 * 1024 * 1024\n\nfunc TestGenerateWithLargeModels(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  gpu: true\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==2.0.1\n    - pandas==2.0.3\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\n\tgen.fileWalker = func(root string, walkFn filepath.WalkFunc) error {\n\t\tfor _, path := range []string{\"checkpoints/large-a\", \"models/large-b\", \"root-large\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizeThreshold}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\tpypiWheels(gen)\n\n\tmodelDockerfile, runnerDockerfile, dockerignore, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM scratch\n\nCOPY checkpoints /src/checkpoints\nCOPY models /src/models\nCOPY root-large /src/root-large`\n\n\trequire.Equal(t, expected, modelDockerfile)\n\n\t// model copy should be run before dependency install and code copy\n\texpected = `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04\nENV DEBIAN_FRONTEND=noninteractive\nENV PYTHONUNBUFFERED=1\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin\nENV NVIDIA_DRIVER_CAPABILITIES=all\n` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*` + `\n` + testInstallPython(\"3.12\") + `COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt\nENV CFLAGS=\n` + testInstallCog(gen.strip) + `\nRUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig\nRUN cowsay moo\nCOPY --from=weights --link /src/checkpoints /src/checkpoints\nCOPY --from=weights --link /src/models /src/models\nCOPY --from=weights --link /src/root-large /src/root-large\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, runnerDockerfile)\n\n\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu118\ntorch==2.0.1\npandas==2.0.3`, string(requirements))\n\n\texpected = `# generated by replicate/cog\n__pycache__\n*.pyc\n*.pyo\n*.pyd\n.Python\nenv\npip-log.txt\npip-delete-this-directory.txt\n.tox\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.log\n.git\n.mypy_cache\n.pytest_cache\n.hypothesis\ncheckpoints\ncheckpoints/**/*\nmodels\nmodels/**/*\nroot-large\n`\n\trequire.Equal(t, expected, dockerignore)\n}\n\nfunc TestGenerateDockerfileWithoutSeparateWeights(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  gpu: false\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\tpypiWheels(gen)\n\tactual, err := gen.GenerateDockerfileWithoutSeparateWeights(t.Context())\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM python:3.12-slim\nENV DEBIAN_FRONTEND=noninteractive\nENV PYTHONUNBUFFERED=1\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin\nENV NVIDIA_DRIVER_CAPABILITIES=all\n` + testTini() + testInstallUVLine + \"\\n\" + testInstallCog(gen.strip) + `\nRUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n}\n\nfunc TestGenerateEmptyCPUWithCogBaseImage(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  gpu: false\n  python_version: \"3.12\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(\"\", \"3.12\", \"\"))\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(true)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM r8.im/cog-base:python3.12\n` + testInstallUVLine + `\n` + testInstallCog(gen.strip) + `\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n}\n\nfunc TestGeneratePythonCPUWithCogBaseImage(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  gpu: false\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - pandas==1.2.0.12\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(\"\", \"3.12\", \"\"))\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(true)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM r8.im/cog-base:python3.12\nRUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*\n` + testInstallUVLine + `\n` + testInstallCog(gen.strip) + `\nCOPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt\nENV CFLAGS=\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\trequire.Equal(t, expected, actual)\n\n\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, `pandas==1.2.0.12`, string(requirements))\n}\n\nfunc TestGenerateFullGPUWithCogBaseImage(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tclient := registrytest.NewMockRegistryClient()\n\tcommand := dockertest.NewMockCommand()\n\ttorchVersions := []string{\"2.3\", \"2.3.0\", \"2.3.1\"}\n\tfor _, torchVersion := range torchVersions {\n\t\tyaml := fmt.Sprintf(`\nbuild:\n  gpu: true\n  cuda: \"11.8\"\n  python_version: \"3.11\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==%s\n    - pandas==2.0.3\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`, torchVersion)\n\t\tconf, err := config.FromYAML([]byte(yaml))\n\t\trequire.NoError(t, err)\n\t\trequire.NoError(t, conf.Complete(\"\"))\n\t\tclient.AddMockImage(BaseImageName(\"11.8\", \"3.11\", torchVersion))\n\t\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\t\trequire.NoError(t, err)\n\t\tgen.SetUseCogBaseImage(true)\n\t\tpypiWheels(gen)\n\t\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\t\trequire.NoError(t, err)\n\n\t\t// We add the patch version to the expected torch version\n\t\texpectedTorchVersion := torchVersion\n\t\tif torchVersion == \"2.3\" {\n\t\t\texpectedTorchVersion = \"2.3.1\"\n\t\t}\n\t\texpected := fmt.Sprintf(`#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM r8.im/cog-base:cuda11.8-python3.11-torch%s\nRUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*\n`+testInstallUVLine+`\n`+testInstallCog(gen.strip)+`\nCOPY `+gen.relativeTmpDir+`/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt\nENV CFLAGS=\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`, expectedTorchVersion)\n\n\t\trequire.Equal(t, expected, actual)\n\n\t\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\t\trequire.NoError(t, err)\n\t\texpected = fmt.Sprintf(`--extra-index-url https://download.pytorch.org/whl/cu118\ntorch==%s\npandas==2.0.3`, expectedTorchVersion)\n\t\trequire.Equal(t, expected, string(requirements))\n\t}\n}\n\nfunc TestGenerateTorchWithStrippedModifiedVersion(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: true\n  cuda: \"11.8\"\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==2.3.1+cu118\n    - pandas==2.0.3\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(\"11.8\", \"3.12\", \"2.3.1\"))\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(true)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM r8.im/cog-base:cuda11.8-python3.12-torch2.3.1\nRUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*\n` + testInstallUVLine + `\n` + testInstallCog(gen.strip) + `\nCOPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt\nENV CFLAGS=\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n\n\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu118\ntorch==2.3.1\npandas==2.0.3`, string(requirements))\n}\n\nfunc TestGenerateWithStrip(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: true\n  cuda: \"11.8\"\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==2.3.1\n    - pandas==2.0.3\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(\"11.8\", \"3.12\", \"2.3.1\"))\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(true)\n\tgen.SetStrip(true)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM r8.im/cog-base:cuda11.8-python3.12-torch2.3.1\nRUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*\n` + testInstallUVLine + `\n` + testInstallCog(gen.strip) + `\nCOPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt && find / -type f -name \"*python*.so\" -not -name \"*cpython*.so\" -exec strip -S {} \\;\nENV CFLAGS=\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n\n\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu118\ntorch==2.3.1\npandas==2.0.3`, string(requirements))\n}\n\nfunc TestGenerateDoesNotContainDangerousCFlags(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: true\n  cuda: \"11.8\"\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==2.3.1\n    - pandas==2.0.3\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(\"11.8\", \"3.12\", \"2.3.1\"))\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(true)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\trequire.NotContains(t, actual, \"-march=native\")\n\trequire.NotContains(t, actual, \"-mtune=native\")\n}\n\nfunc TestGenerateWithPrecompile(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: true\n  cuda: \"11.8\"\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==2.3.1\n    - pandas==2.0.3\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(\"11.8\", \"3.12\", \"2.3.1\"))\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(true)\n\tgen.SetStrip(true)\n\tgen.SetPrecompile(true)\n\tpypiWheels(gen)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM r8.im/cog-base:cuda11.8-python3.12-torch2.3.1\nRUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*\n` + testInstallUVLine + `\n` + testInstallCog(gen.strip) + `\nCOPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt && find / -type f -name \"*python*.so\" -not -name \"*cpython*.so\" -exec strip -S {} \\;\nENV CFLAGS=\nRUN find / -type f -name \"*.py[co]\" -delete && find / -type f -name \"*.py\" -exec touch -t 197001010000 {} \\; && find / -type f -name \"*.py\" -printf \"%h\\n\" | sort -u | /usr/bin/python3 -m compileall --invalidation-mode timestamp -o 2 -j 0\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n\n\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu118\ntorch==2.3.1\npandas==2.0.3`, string(requirements))\n}\n\nfunc TestGenerateWithCoglet(t *testing.T) {\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: true\n  cuda: \"11.8\"\n  python_version: \"3.12\"\n  system_packages:\n    - ffmpeg\n    - cowsay\n  python_packages:\n    - torch==2.3.1\n    - pandas==2.0.3\n    - coglet @ https://github.com/replicate/cog-runtime/releases/download/v0.1.0-alpha31/coglet-0.1.0a31-py3-none-any.whl\n  run:\n    - \"cowsay moo\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tclient.AddMockImage(BaseImageName(\"11.8\", \"3.12\", \"2.3.1\"))\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(true)\n\tgen.SetStrip(true)\n\tgen.SetPrecompile(true)\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\t// coglet in python_packages is stripped — the build system always installs coglet\n\t// via installCog(), which runs before pip requirements.\n\texpected := `#syntax=docker/dockerfile:1.4\nFROM r8.im/replicate/cog-test-weights AS weights\nFROM r8.im/cog-base:cuda11.8-python3.12-torch2.3.1\nRUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*\n` + testInstallUVLine + `\n` + testInstallCog(true) + `\nCOPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt\nENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"\nRUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r /tmp/requirements.txt && find / -type f -name \"*python*.so\" -not -name \"*cpython*.so\" -exec strip -S {} \\;\nENV CFLAGS=\nRUN find / -type f -name \"*.py[co]\" -delete && find / -type f -name \"*.py\" -exec touch -t 197001010000 {} \\; && find / -type f -name \"*.py\" -printf \"%h\\n\" | sort -u | /usr/bin/python3 -m compileall --invalidation-mode timestamp -o 2 -j 0\nRUN cowsay moo\nWORKDIR /src\nEXPOSE 5000\nENV COG_PREDICT_TYPE_STUB=\"predict.py:Predictor\"\nCMD [\"python\", \"-m\", \"cog.server.http\"]\nCOPY . /src`\n\n\trequire.Equal(t, expected, actual)\n\n\t// coglet URL is stripped from requirements — build system installs coglet itself\n\trequirements, err := os.ReadFile(path.Join(gen.tmpDir, \"requirements.txt\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu118\ntorch==2.3.1\npandas==2.0.3`, string(requirements))\n}\n\nfunc TestCOGWheelDefault(t *testing.T) {\n\t// Default behavior should install cog from PyPI\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: false\n  python_version: \"3.11\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tpypiWheels(gen)\n\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\t// Should contain uv pip install cog from PyPI.\n\t// Coglet is not explicitly installed when unpinned — the SDK dependency handles it.\n\trequire.Contains(t, actual, \"uv pip install --no-cache cog\")\n}\n\nfunc TestCOGWheelEnvPyPI(t *testing.T) {\n\t// COG_SDK_WHEEL=pypi should install from PyPI\n\tt.Setenv(\"COG_SDK_WHEEL\", \"pypi\")\n\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: false\n  python_version: \"3.11\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\t// Should contain uv pip install cog from PyPI\n\trequire.Contains(t, actual, \"uv pip install --no-cache cog\")\n}\n\nfunc TestCOGWheelEnvPyPIWithVersion(t *testing.T) {\n\t// COG_SDK_WHEEL=pypi:0.17.0 should install specific version from PyPI\n\tt.Setenv(\"COG_SDK_WHEEL\", \"pypi:0.17.0\")\n\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: false\n  python_version: \"3.11\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\t// Should contain uv pip install cog==0.17.0 from PyPI\n\trequire.Contains(t, actual, \"uv pip install --no-cache cog==0.17.0\")\n}\n\nfunc TestCOGWheelEnvURL(t *testing.T) {\n\t// COG_SDK_WHEEL=https://... should install from URL\n\tcustomURL := \"https://example.com/custom-wheel-0.1.0.whl\"\n\tt.Setenv(\"COG_SDK_WHEEL\", customURL)\n\n\ttmpDir := t.TempDir()\n\n\tyaml := `\nbuild:\n  gpu: false\n  python_version: \"3.11\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\t// Should contain uv pip install from custom URL\n\trequire.Contains(t, actual, \"uv pip install --no-cache \"+customURL)\n}\n\nfunc TestCOGWheelEnvFile(t *testing.T) {\n\t// COG_SDK_WHEEL=/path/to/file.whl should install from local file\n\ttmpDir := t.TempDir()\n\n\t// Create a mock wheel file\n\twheelPath := filepath.Join(tmpDir, \"test-cog-0.1.0-py3-none-any.whl\")\n\terr := os.WriteFile(wheelPath, []byte(\"mock wheel content\"), 0o644)\n\trequire.NoError(t, err)\n\n\tt.Setenv(\"COG_SDK_WHEEL\", wheelPath)\n\n\tyaml := `\nbuild:\n  gpu: false\n  python_version: \"3.11\"\npredict: predict.py:Predictor\n`\n\tconf, err := config.FromYAML([]byte(yaml))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\n\t_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights(t.Context(), \"r8.im/replicate/cog-test\")\n\trequire.NoError(t, err)\n\n\t// Should contain uv pip install from temp path (copied into container)\n\trequire.Contains(t, actual, \"uv pip install --no-cache /tmp/test-cog-0.1.0-py3-none-any.whl\")\n}\n\nfunc TestCogletStrippedFromRequirements(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  python_packages:\n    - \"coglet==0.1.0\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.SetUseCogBaseImage(false)\n\tpypiWheels(gen)\n\tdockerfile, err := gen.GenerateInitialSteps(t.Context())\n\trequire.NoError(t, err)\n\t// coglet is NOT explicitly installed — SDK dependency handles it.\n\t// But the user-supplied coglet==0.1.0 must be stripped from requirements\n\t// to avoid conflicting with whatever version the SDK pulls in.\n\trequire.NotContains(t, dockerfile, \"coglet==0.1.0\")\n}\n\nfunc TestInstallCogWithSDKVersion(t *testing.T) {\n\t// build.sdk_version pins the cog SDK version installed from PyPI\n\ttmpDir := t.TempDir()\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  sdk_version: \"0.18.0\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\t// Only pin coglet to PyPI; leave cog to come from config sdk_version\n\tgen.cogletWheelConfig = &wheels.WheelConfig{Source: wheels.WheelSourcePyPI}\n\tgen.SetUseCogBaseImage(false)\n\n\tdockerfile, err := gen.GenerateInitialSteps(t.Context())\n\trequire.NoError(t, err)\n\trequire.Contains(t, dockerfile, \"uv pip install --no-cache cog==0.18.0\")\n\t// No --pre flag for stable release\n\trequire.NotContains(t, dockerfile, \"--pre\")\n}\n\nfunc TestInstallCogWithPreReleaseSDKVersion(t *testing.T) {\n\t// build.sdk_version with a pre-release version adds --pre to cog install\n\ttmpDir := t.TempDir()\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  sdk_version: \"0.18.0a1\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.cogletWheelConfig = &wheels.WheelConfig{Source: wheels.WheelSourcePyPI}\n\tgen.SetUseCogBaseImage(false)\n\n\tdockerfile, err := gen.GenerateInitialSteps(t.Context())\n\trequire.NoError(t, err)\n\t// cog install should have --pre and pinned version\n\trequire.Contains(t, dockerfile, \"uv pip install --pre --no-cache cog==0.18.0a1\")\n\t// coglet is NOT explicitly installed — SDK dependency pulls it in.\n\t// No separate coglet install line expected.\n}\n\nfunc TestInstallCogSDKVersionBelowMinimum(t *testing.T) {\n\t// build.sdk_version below MinimumSDKVersion should return an error\n\ttmpDir := t.TempDir()\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  sdk_version: \"0.15.0\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.cogletWheelConfig = &wheels.WheelConfig{Source: wheels.WheelSourcePyPI}\n\tgen.SetUseCogBaseImage(false)\n\n\t_, err = gen.GenerateInitialSteps(t.Context())\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"0.15.0\")\n\trequire.Contains(t, err.Error(), \"minimum required version\")\n}\n\nfunc TestCOGSDKWheelEnvVarOverridesSDKVersion(t *testing.T) {\n\t// COG_SDK_WHEEL env var overrides build.sdk_version\n\tt.Setenv(\"COG_SDK_WHEEL\", \"pypi:0.17.0\")\n\ttmpDir := t.TempDir()\n\tconf, err := config.FromYAML([]byte(`\nbuild:\n  python_version: \"3.12\"\n  sdk_version: \"0.18.0\"\npredict: predict.py:Predictor\n`))\n\trequire.NoError(t, err)\n\trequire.NoError(t, conf.Complete(\"\"))\n\tcommand := dockertest.NewMockCommand()\n\tclient := registrytest.NewMockRegistryClient()\n\tgen, err := NewStandardGenerator(conf, tmpDir, \"\", command, client, true)\n\trequire.NoError(t, err)\n\tgen.cogletWheelConfig = &wheels.WheelConfig{Source: wheels.WheelSourcePyPI}\n\tgen.SetUseCogBaseImage(false)\n\n\tdockerfile, err := gen.GenerateInitialSteps(t.Context())\n\trequire.NoError(t, err)\n\t// env var wins: should install 0.17.0, not 0.18.0\n\trequire.Contains(t, dockerfile, \"uv pip install --no-cache cog==0.17.0\")\n\trequire.NotContains(t, dockerfile, \"cog==0.18.0\")\n}\n"
  },
  {
    "path": "pkg/dockerfile/version_check.go",
    "content": "package dockerfile\n\nimport (\n\t\"regexp\"\n)\n\n// Version string in the form x.y.z (e.g., Python version, CUDA version)\n// We do not support suffixes like -alpha1 or +cu124\nvar versionRegex = regexp.MustCompile(`^(?P\\d+)(\\.(?P\\d+)(\\.(?P\\d+))?)?$`)\n\nfunc parse(s string) (string, string, string) {\n\tm := versionRegex.FindStringSubmatch(s)\n\tif m == nil {\n\t\treturn \"\", \"\", \"\"\n\t}\n\tmajor := m[versionRegex.SubexpIndex(\"major\")]\n\tminor := m[versionRegex.SubexpIndex(\"minor\")]\n\tpatch := m[versionRegex.SubexpIndex(\"patch\")]\n\treturn major, minor, patch\n\n}\n\nfunc CheckMajorOnly(s string) bool {\n\tmajor, minor, patch := parse(s)\n\treturn major != \"\" && minor == \"\" && patch == \"\"\n}\n\nfunc CheckMajorMinorOnly(s string) bool {\n\tmajor, minor, patch := parse(s)\n\treturn major != \"\" && minor != \"\" && patch == \"\"\n}\n\nfunc CheckMajorMinorPatch(s string) bool {\n\tmajor, minor, patch := parse(s)\n\treturn major != \"\" && minor != \"\" && patch != \"\"\n}\n"
  },
  {
    "path": "pkg/dockerignore/dockerignore.go",
    "content": "package dockerignore\n\nimport (\n\t\"bufio\"\n\t\"os\"\n\t\"path/filepath\"\n\n\tignore \"github.com/sabhiram/go-gitignore\"\n\n\t\"github.com/replicate/cog/pkg/util/files\"\n)\n\nconst DockerIgnoreFilename = \".dockerignore\"\n\nfunc CreateMatcher(dir string) (*ignore.GitIgnore, error) {\n\tdockerIgnorePath := filepath.Join(dir, DockerIgnoreFilename)\n\tdockerIgnoreExists, err := files.Exists(dockerIgnorePath)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif !dockerIgnoreExists {\n\t\treturn nil, nil\n\t}\n\n\tpatterns, err := readDockerIgnore(dockerIgnorePath)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn ignore.CompileIgnoreLines(patterns...), nil\n}\n\nfunc Walk(root string, ignoreMatcher *ignore.GitIgnore, fn filepath.WalkFunc) error {\n\treturn filepath.Walk(root, func(path string, info os.FileInfo, err error) error {\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\t// We ignore files ignored by .dockerignore\n\t\tif ignoreMatcher != nil && ignoreMatcher.MatchesPath(path) {\n\t\t\tif info.IsDir() {\n\t\t\t\treturn filepath.SkipDir\n\t\t\t}\n\t\t\treturn nil\n\t\t}\n\n\t\tif info.IsDir() && info.Name() == \".cog\" {\n\t\t\treturn filepath.SkipDir\n\t\t}\n\n\t\tif info.Name() == DockerIgnoreFilename {\n\t\t\treturn nil\n\t\t}\n\n\t\treturn fn(path, info, err)\n\t})\n}\n\nfunc readDockerIgnore(dockerIgnorePath string) ([]string, error) {\n\tvar patterns []string\n\tfile, err := os.Open(dockerIgnorePath)\n\tif err != nil {\n\t\treturn patterns, err\n\t}\n\tdefer file.Close()\n\n\tscanner := bufio.NewScanner(file)\n\tfor scanner.Scan() {\n\t\tline := scanner.Text()\n\t\tpatterns = append(patterns, line)\n\t}\n\treturn patterns, scanner.Err()\n}\n"
  },
  {
    "path": "pkg/dockerignore/dockerignore_test.go",
    "content": "package dockerignore\n\nimport (\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestWalk(t *testing.T) {\n\tdir := t.TempDir()\n\n\tpredictOtherPyFilename := \"predict_other.py\"\n\tpredictOtherPyFilepath := filepath.Join(dir, predictOtherPyFilename)\n\tpredictOtherPyHandle, err := os.Create(predictOtherPyFilepath)\n\trequire.NoError(t, err)\n\tpredictOtherPyHandle.WriteString(\"import cog\")\n\n\tdockerIgnorePath := filepath.Join(dir, \".dockerignore\")\n\tdockerIgnoreHandle, err := os.Create(dockerIgnorePath)\n\trequire.NoError(t, err)\n\tdockerIgnoreHandle.WriteString(predictOtherPyFilename)\n\n\tpredictPyFilename := \"predict.py\"\n\tpredictPyFilepath := filepath.Join(dir, predictPyFilename)\n\tpredictPyHandle, err := os.Create(predictPyFilepath)\n\trequire.NoError(t, err)\n\tpredictPyHandle.WriteString(\"import cog\")\n\n\tmatcher, err := CreateMatcher(dir)\n\trequire.NoError(t, err)\n\n\tfoundFiles := []string{}\n\terr = Walk(dir, matcher, func(path string, info os.FileInfo, err error) error {\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tif info.IsDir() {\n\t\t\treturn nil\n\t\t}\n\n\t\trelPath, err := filepath.Rel(dir, path)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tfoundFiles = append(foundFiles, relPath)\n\n\t\treturn nil\n\t})\n\trequire.NoError(t, err)\n\n\trequire.Equal(t, []string{predictPyFilename}, foundFiles)\n}\n"
  },
  {
    "path": "pkg/env/env.go",
    "content": "package env\n\nimport \"os\"\n\nconst SchemeEnvVarName = \"R8_SCHEME\"\nconst WebHostEnvVarName = \"R8_WEB_HOST\"\nconst APIHostEnvVarName = \"R8_API_HOST\"\nconst PytorchHostEnvVarName = \"R8_PYTORCH_HOST\"\n\nfunc SchemeFromEnvironment() string {\n\tscheme := os.Getenv(SchemeEnvVarName)\n\tif scheme == \"\" {\n\t\tscheme = \"https\"\n\t}\n\treturn scheme\n}\n\nfunc WebHostFromEnvironment() string {\n\thost := os.Getenv(WebHostEnvVarName)\n\tif host == \"\" {\n\t\thost = \"cog.replicate.com\"\n\t}\n\treturn host\n}\n\nfunc APIHostFromEnvironment() string {\n\thost := os.Getenv(APIHostEnvVarName)\n\tif host == \"\" {\n\t\thost = \"api.replicate.com\"\n\t}\n\treturn host\n}\n\nfunc PytorchHostFromEnvironment() string {\n\thost := os.Getenv(PytorchHostEnvVarName)\n\tif host == \"\" {\n\t\thost = \"download.pytorch.org\"\n\t}\n\treturn host\n}\n"
  },
  {
    "path": "pkg/env/env_test.go",
    "content": "package env\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestSchemeFromEnvironment(t *testing.T) {\n\tconst testScheme = \"myscheme\"\n\tt.Setenv(SchemeEnvVarName, \"myscheme\")\n\trequire.Equal(t, SchemeFromEnvironment(), testScheme)\n}\n\nfunc TestWebHostFromEnvironment(t *testing.T) {\n\tconst testHost = \"web\"\n\tt.Setenv(WebHostEnvVarName, testHost)\n\trequire.Equal(t, WebHostFromEnvironment(), testHost)\n}\n"
  },
  {
    "path": "pkg/errors/common.go",
    "content": "package errors\n\nimport (\n\t\"errors\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n)\n\nvar (\n\tErrorBadRegistryURL = errors.New(\"The image URL must have 3 components in the format of \" + global.ReplicateRegistryHost + \"/your-username/your-model\")\n)\n"
  },
  {
    "path": "pkg/errors/errors.go",
    "content": "package errors\n\nconst (\n\tCodeConfigNotFound = \"CONFIG_NOT_FOUND\"\n)\n\n// Types ////////////////////////////////////////\n\ntype CodedError interface {\n\tCode() string\n}\n\ntype codedError struct {\n\tcode string\n\tmsg  string\n}\n\nfunc (e *codedError) Error() string {\n\treturn e.msg\n}\n\nfunc (e *codedError) Code() string {\n\treturn e.code\n}\n\n// Error Creators ///////////////////////////////\n\n// The Cog config was not found\nfunc ConfigNotFound(msg string) error {\n\treturn &codedError{\n\t\tcode: CodeConfigNotFound,\n\t\tmsg:  msg + ``, // TODO: populate this\n\t}\n}\n\n// Helpers //////////////////////////////////////\n\nfunc IsConfigNotFound(err error) bool {\n\treturn Code(err) == CodeConfigNotFound\n}\n\n// Return the error code, or the empty string\nfunc Code(err error) string {\n\tif cerr, ok := err.(CodedError); ok {\n\t\treturn cerr.Code()\n\t}\n\n\treturn \"\"\n}\n"
  },
  {
    "path": "pkg/global/global.go",
    "content": "package global\n\nimport \"os\"\n\nconst (\n\tDefaultReplicateRegistryHost = \"r8.im\"\n\tReplicateWebsiteHost         = \"replicate.com\"\n)\n\nvar (\n\tVersion               = \"dev\"\n\tCommit                = \"\"\n\tBuildTime             = \"none\"\n\tDebug                 = false\n\tNoColor               = false\n\tProfilingEnabled      = false\n\tReplicateRegistryHost = getDefaultRegistryHost()\n\n\tLabelNamespace          = \"run.cog.\"\n\tCogBuildArtifactsFolder = \".cog\"\n)\n\nfunc getDefaultRegistryHost() string {\n\t// Priority: flag will override at runtime, but env var provides default\n\tif host := os.Getenv(\"COG_REGISTRY_HOST\"); host != \"\" {\n\t\treturn host\n\t}\n\treturn DefaultReplicateRegistryHost\n}\n"
  },
  {
    "path": "pkg/http/client.go",
    "content": "package http\n\nimport (\n\t\"context\"\n\t\"net/http\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/env\"\n\t\"github.com/replicate/cog/pkg/global\"\n)\n\nconst UserAgentHeader = \"User-Agent\"\nconst BearerHeaderPrefix = \"Bearer \"\n\nfunc ProvideHTTPClient(ctx context.Context, dockerCommand command.Command) (*http.Client, error) {\n\tuserInfo, err := dockerCommand.LoadUserInformation(ctx, global.ReplicateRegistryHost)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tclient := http.Client{\n\t\tTransport: &Transport{\n\t\t\theaders: map[string]string{\n\t\t\t\tUserAgentHeader: UserAgent(),\n\t\t\t\t\"Content-Type\":  \"application/json\",\n\t\t\t},\n\t\t\tauthentication: map[string]string{\n\t\t\t\tenv.WebHostFromEnvironment(): BearerHeaderPrefix + userInfo.Token,\n\t\t\t},\n\t\t},\n\t}\n\n\treturn &client, nil\n}\n"
  },
  {
    "path": "pkg/http/client_test.go",
    "content": "package http\n\nimport (\n\t\"net/http\"\n\t\"net/http/httptest\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n)\n\nfunc TestClientDecoratesUserAgent(t *testing.T) {\n\t// Setup mock http server\n\tseenUserAgent := false\n\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\trequire.Equal(t, r.Header.Get(UserAgentHeader), UserAgent())\n\t\tseenUserAgent = true\n\t}))\n\tdefer server.Close()\n\n\tcommand := dockertest.NewMockCommand()\n\tclient, err := ProvideHTTPClient(t.Context(), command)\n\trequire.NoError(t, err)\n\n\t_, err = client.Get(server.URL)\n\trequire.NoError(t, err)\n\n\trequire.True(t, seenUserAgent)\n}\n"
  },
  {
    "path": "pkg/http/transport.go",
    "content": "package http\n\nimport (\n\t\"errors\"\n\t\"net/http\"\n)\n\nconst AuthorizationHeader = \"Authorization\"\n\ntype Transport struct {\n\theaders        map[string]string\n\tauthentication map[string]string\n\tbase           http.RoundTripper\n}\n\nfunc (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {\n\t// Write standard headers\n\tfor k, v := range t.headers {\n\t\tif req.Header.Get(k) == \"\" {\n\t\t\treq.Header.Set(k, v)\n\t\t}\n\t}\n\n\t// Write authentication\n\tif req.Header.Get(AuthorizationHeader) == \"\" {\n\t\tauthorisation, ok := t.authentication[req.URL.Host]\n\t\tif ok {\n\t\t\tif authorisation == BearerHeaderPrefix {\n\t\t\t\treturn nil, errors.New(\"No token supplied for HTTP authorization. Have you run 'cog login'?\")\n\t\t\t}\n\t\t\treq.Header.Set(AuthorizationHeader, authorisation)\n\t\t}\n\t}\n\n\tbase := t.base\n\tif base == nil {\n\t\tbase = http.DefaultTransport\n\t}\n\treturn base.RoundTrip(req)\n}\n"
  },
  {
    "path": "pkg/http/transport_test.go",
    "content": "package http\n\nimport (\n\t\"net/http\"\n\t\"net/http/httptest\"\n\t\"net/url\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestTransportAddsHeaders(t *testing.T) {\n\t// Setup mock http server\n\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\tw.WriteHeader(http.StatusOK)\n\t}))\n\tdefer server.Close()\n\n\tconst testHeader = \"X-Test-Header\"\n\tconst testValue = \"TestValue\"\n\ttransport := Transport{\n\t\theaders: map[string]string{\n\t\t\ttestHeader: testValue,\n\t\t},\n\t}\n\treq, err := http.NewRequest(\"GET\", server.URL, nil)\n\trequire.NoError(t, err)\n\tresp, err := transport.RoundTrip(req)\n\trequire.NoError(t, err)\n\trequire.Equal(t, resp.Request.Header.Get(testHeader), testValue)\n}\n\nfunc TestTransportOnlyAddsHeaderIfMissing(t *testing.T) {\n\t// Setup mock http server\n\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\tw.WriteHeader(http.StatusOK)\n\t}))\n\tdefer server.Close()\n\n\tconst testHeader = \"X-Test-Header\"\n\tconst testValue = \"TestValue\"\n\ttransport := Transport{\n\t\theaders: map[string]string{\n\t\t\ttestHeader: testValue,\n\t\t},\n\t}\n\tconst expectedValue = \"ExpectedValue\"\n\treq, err := http.NewRequest(\"GET\", server.URL, nil)\n\treq.Header.Set(testHeader, expectedValue)\n\trequire.NoError(t, err)\n\tresp, err := transport.RoundTrip(req)\n\trequire.NoError(t, err)\n\trequire.Equal(t, resp.Request.Header.Get(testHeader), expectedValue)\n}\n\nfunc TestTransportSendsErrorWithMissingToken(t *testing.T) {\n\t// Setup mock http server\n\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\tw.WriteHeader(http.StatusOK)\n\t}))\n\tdefer server.Close()\n\tu, err := url.Parse(server.URL)\n\trequire.NoError(t, err)\n\n\ttransport := Transport{\n\t\tauthentication: map[string]string{\n\t\t\tu.Host: BearerHeaderPrefix + \"\",\n\t\t},\n\t}\n\treq, err := http.NewRequest(\"GET\", server.URL, nil)\n\trequire.NoError(t, err)\n\tresp, err := transport.RoundTrip(req)\n\trequire.Error(t, err)\n\trequire.Nil(t, resp)\n}\n"
  },
  {
    "path": "pkg/http/user_agent.go",
    "content": "package http\n\nimport (\n\t\"fmt\"\n\t\"runtime\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n)\n\nfunc UserAgent() string {\n\tvar platform string\n\tswitch runtime.GOOS {\n\tcase \"linux\":\n\t\tplatform = \"Linux\"\n\tcase \"windows\":\n\t\tplatform = \"Windows\"\n\tcase \"darwin\":\n\t\tplatform = \"macOS\"\n\tdefault:\n\t\tplatform = runtime.GOOS\n\t}\n\n\treturn fmt.Sprintf(\"Cog/%s (%s)\", global.Version, platform)\n}\n"
  },
  {
    "path": "pkg/http/user_agent_test.go",
    "content": "package http\n\nimport (\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestUserAgent(t *testing.T) {\n\trequire.True(t, strings.HasPrefix(UserAgent(), \"Cog/\"))\n}\n"
  },
  {
    "path": "pkg/image/build.go",
    "content": "package image\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"crypto/sha256\"\n\t\"encoding/hex\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"maps\"\n\t\"os\"\n\t\"os/exec\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"time\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/dockercontext\"\n\t\"github.com/replicate/cog/pkg/dockerfile\"\n\t\"github.com/replicate/cog/pkg/dockerignore\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/schema\"\n\t\"github.com/replicate/cog/pkg/schema/python\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n\tcogversion \"github.com/replicate/cog/pkg/util/version\"\n\t\"github.com/replicate/cog/pkg/weights\"\n\t\"github.com/replicate/cog/pkg/wheels\"\n)\n\nconst dockerignoreBackupPath = \".dockerignore.cog.bak\"\nconst weightsManifestPath = \".cog/cache/weights_manifest.json\"\nconst bundledSchemaFile = \".cog/openapi_schema.json\"\n\nvar errGit = errors.New(\"git error\")\n\n// Build a Cog model from a config and returns the image ID (sha256:...) on success.\n//\n// This is separated out from docker.Build(), so that can be as close as possible to the behavior of 'docker build'.\nfunc Build(\n\tctx context.Context,\n\tcfg *config.Config,\n\tdir,\n\timageName string,\n\tconfigFilename string,\n\tsecrets []string,\n\tnoCache,\n\tseparateWeights bool,\n\tuseCudaBaseImage string,\n\tprogressOutput string,\n\tschemaFile string,\n\tdockerfileFile string,\n\tuseCogBaseImage *bool,\n\tstrip bool,\n\tprecompile bool,\n\texcludeSource bool,\n\tskipSchemaValidation bool,\n\tskipLabels bool,\n\tannotations map[string]string,\n\tdockerCommand command.Command,\n\tclient registry.Client) (string, error) {\n\t// remove bundled schema files that may be left from previous builds\n\t_ = os.Remove(bundledSchemaFile)\n\n\tif err := checkCompatibleDockerIgnore(dir); err != nil {\n\t\treturn \"\", err\n\t}\n\n\t// Determine whether to use the static schema generator (Go tree-sitter) or\n\t// the legacy runtime path (boot container + python introspection).\n\t//\n\t// Static generation is opt-in via COG_STATIC_SCHEMA=1 for all commands.\n\t// The legacy runtime path (boot container + python -m cog.command.openapi_schema)\n\t// remains the default for `cog build`. For `cog train`, `cog predict`, and\n\t// `cog serve` (skipLabels=true), no schema is generated unless\n\t// COG_STATIC_SCHEMA=1 is set, since these paths return before the post-build\n\t// legacy schema generation step.\n\t//\n\t// The SDK version must be >= 0.17.0 (or unpinned/latest/dev) since older\n\t// SDKs use pydantic-based schemas that cannot be statically analyzed.\n\tneedsSchema := !skipSchemaValidation && schemaFile == \"\"\n\tuseStatic := needsSchema && canUseStaticSchemaGen(cfg)\n\n\t// --- Pre-build static schema generation ---\n\t// When using the static path, generate schema BEFORE the Docker build so we\n\t// fail fast on schema errors and the schema file is in the build context.\n\tvar schemaJSON []byte\n\tswitch {\n\tcase useStatic:\n\t\tconsole.Debug(\"Generating model schema (static)...\")\n\t\tdata, err := generateStaticSchema(cfg, dir)\n\t\tif err == nil {\n\t\t\tschemaJSON = data\n\t\t\tbreak\n\t\t}\n\n\t\t// For `cog build` only: fall back to the post-build legacy runtime\n\t\t// schema generation which can handle types that require Python import\n\t\t// (e.g. package __init__.py modules, pydantic v2 BaseModel subclasses).\n\t\tvar se *schema.SchemaError\n\t\tif !skipLabels && errors.As(err, &se) && se.Kind == schema.ErrUnresolvableType {\n\t\t\tconsole.Warnf(\"Static schema generation failed: %s\", err)\n\t\t\tconsole.Warn(\"Falling back to legacy runtime schema generation...\")\n\t\t\t// leave schemaJSON nil — the post-build legacy path will handle it\n\t\t\tbreak\n\t\t}\n\n\t\treturn \"\", fmt.Errorf(\"image build failed: %w\", err)\n\tcase !skipSchemaValidation && schemaFile != \"\":\n\t\tconsole.Infof(\"Validating model schema from %s...\", schemaFile)\n\t\tdata, err := os.ReadFile(schemaFile)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to read schema file: %w\", err)\n\t\t}\n\t\tschemaJSON = data\n\tcase skipSchemaValidation:\n\t\tconsole.Debug(\"Skipping model schema validation\")\n\t}\n\n\t// Write and validate pre-build schema (static or from file).\n\tif len(schemaJSON) > 0 {\n\t\tif err := writeAndValidateSchema(schemaJSON); err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t}\n\n\t// --- Docker build ---\n\tvar cogBaseImageName string\n\n\ttmpImageId := imageName\n\tisR8imImage := strings.HasPrefix(imageName, \"r8.im\")\n\tif isR8imImage {\n\t\thash := sha256.New()\n\t\t_, err := hash.Write([]byte(imageName))\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\ttmpImageId = fmt.Sprintf(\"cog-tmp:%s\", hex.EncodeToString(hash.Sum(nil)))\n\t}\n\n\tif dockerfileFile != \"\" {\n\t\tdockerfileContents, err := os.ReadFile(dockerfileFile)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to read Dockerfile at %s: %w\", dockerfileFile, err)\n\t\t}\n\n\t\tbuildOpts := command.ImageBuildOptions{\n\t\t\tWorkingDir:         dir,\n\t\t\tDockerfileContents: string(dockerfileContents),\n\t\t\tImageName:          tmpImageId,\n\t\t\tSecrets:            secrets,\n\t\t\tNoCache:            noCache,\n\t\t\tProgressOutput:     progressOutput,\n\t\t\tEpoch:              &config.BuildSourceEpochTimestamp,\n\t\t\tContextDir:         dockercontext.StandardBuildDirectory,\n\t\t}\n\t\tif _, err := dockerCommand.ImageBuild(ctx, buildOpts); err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to build Docker image: %w\", err)\n\t\t}\n\t} else {\n\t\tgenerator, err := dockerfile.NewGenerator(cfg, dir, configFilename, dockerCommand, client, true)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Error creating Dockerfile generator: %w\", err)\n\t\t}\n\t\tcontextDir, err := generator.BuildDir()\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tbuildContexts, err := generator.BuildContexts()\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tdefer func() {\n\t\t\tif err := generator.Cleanup(); err != nil {\n\t\t\t\tconsole.Warnf(\"Error cleaning up Dockerfile generator: %s\", err)\n\t\t\t}\n\t\t}()\n\t\tgenerator.SetStrip(strip)\n\t\tgenerator.SetPrecompile(precompile)\n\t\tgenerator.SetUseCudaBaseImage(useCudaBaseImage)\n\t\tif useCogBaseImage != nil {\n\t\t\tgenerator.SetUseCogBaseImage(*useCogBaseImage)\n\t\t}\n\n\t\tif generator.IsUsingCogBaseImage() {\n\t\t\tcogBaseImageName, err = generator.BaseImage(ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"Failed to get cog base image name: %s\", err)\n\t\t\t}\n\t\t}\n\n\t\tif separateWeights {\n\t\t\tweightsDockerfile, runnerDockerfile, dockerignore, err := generator.GenerateModelBaseWithSeparateWeights(ctx, imageName)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"Failed to generate Dockerfile: %w\", err)\n\t\t\t}\n\n\t\t\tif err := backupDockerignore(); err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"Failed to backup .dockerignore file: %w\", err)\n\t\t\t}\n\n\t\t\tweightsManifest, err := generator.GenerateWeightsManifest(ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"Failed to generate weights manifest: %w\", err)\n\t\t\t}\n\t\t\tcachedManifest, _ := weights.LoadManifest(weightsManifestPath)\n\t\t\tchanged := cachedManifest == nil || !weightsManifest.Equal(cachedManifest)\n\t\t\tif changed {\n\t\t\t\tif err := buildWeightsImage(ctx, dockerCommand, dir, weightsDockerfile, imageName+\"-weights\", secrets, noCache, progressOutput, contextDir, buildContexts); err != nil {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"Failed to build model weights Docker image: %w\", err)\n\t\t\t\t}\n\t\t\t\terr := weightsManifest.Save(weightsManifestPath)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn \"\", fmt.Errorf(\"Failed to save weights hash: %w\", err)\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\tconsole.Info(\"Weights unchanged, skip rebuilding and use cached image...\")\n\t\t\t}\n\n\t\t\tif err := buildRunnerImage(ctx, dockerCommand, dir, runnerDockerfile, dockerignore, imageName, secrets, noCache, progressOutput, contextDir, buildContexts); err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"Failed to build runner Docker image: %w\", err)\n\t\t\t}\n\t\t} else {\n\t\t\tvar dockerfileContents string\n\t\t\tif excludeSource {\n\t\t\t\t// Dev mode (cog serve): same layers as cog build but without\n\t\t\t\t// COPY . /src — source is volume-mounted at runtime instead.\n\t\t\t\t// This shares Docker layer cache with full builds.\n\t\t\t\tdockerfileContents, err = generator.GenerateModelBase(ctx)\n\t\t\t} else {\n\t\t\t\tdockerfileContents, err = generator.GenerateDockerfileWithoutSeparateWeights(ctx)\n\t\t\t}\n\t\t\tif err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"Failed to generate Dockerfile: %w\", err)\n\t\t\t}\n\n\t\t\tbuildOpts := command.ImageBuildOptions{\n\t\t\t\tWorkingDir:         dir,\n\t\t\t\tDockerfileContents: dockerfileContents,\n\t\t\t\tImageName:          tmpImageId,\n\t\t\t\tSecrets:            secrets,\n\t\t\t\tNoCache:            noCache,\n\t\t\t\tProgressOutput:     progressOutput,\n\t\t\t\tEpoch:              &config.BuildSourceEpochTimestamp,\n\t\t\t\tContextDir:         contextDir,\n\t\t\t\tBuildContexts:      buildContexts,\n\t\t\t}\n\n\t\t\tif _, err := dockerCommand.ImageBuild(ctx, buildOpts); err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"Failed to build Docker image: %w\", err)\n\t\t\t}\n\t\t}\n\t}\n\n\t// --- Post-build legacy schema generation ---\n\t// For SDK < 0.17.0 (or when static gen was not used), generate the schema\n\t// by running the built image with python -m cog.command.openapi_schema.\n\t// This must run before the skipLabels early return so that cog train/predict/serve\n\t// have a schema available for input validation and -i flag parsing.\n\tif len(schemaJSON) == 0 && !skipSchemaValidation {\n\t\tconsole.Info(\"Validating model schema...\")\n\t\tenableGPU := cfg.Build != nil && cfg.Build.GPU\n\t\t// When excludeSource is true (cog serve/predict/train), /src was not\n\t\t// COPYed into the image, so volume-mount the project directory.\n\t\tsourceDir := \"\"\n\t\tif excludeSource {\n\t\t\tsourceDir = dir\n\t\t}\n\t\tlegacySchema, err := GenerateOpenAPISchema(ctx, dockerCommand, tmpImageId, enableGPU, sourceDir)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to get type signature: %w\", err)\n\t\t}\n\t\tdata, err := json.Marshal(legacySchema)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to convert type signature to JSON: %w\", err)\n\t\t}\n\t\tschemaJSON = data\n\n\t\tif err := writeAndValidateSchema(schemaJSON); err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t}\n\n\t// When skipLabels is true (cog run/predict/serve/train), skip the expensive\n\t// label-adding phase. This image is for local use only and won't be distributed,\n\t// so we don't need metadata labels, pip freeze, or git info.\n\t// We still need the schema bundled, so do a minimal second build to add it.\n\tif skipLabels {\n\t\tif len(schemaJSON) > 0 {\n\t\t\t// Use trailing \"/\" on the destination so Docker creates the .cog/\n\t\t\t// directory even in ExcludeSource images where COPY . /src was\n\t\t\t// skipped and .cog/ does not yet exist.\n\t\t\tschemaDockerfile := fmt.Sprintf(\"FROM %s\\nCOPY %s .cog/\\n\", tmpImageId, bundledSchemaFile)\n\t\t\tbuildOpts := command.ImageBuildOptions{\n\t\t\t\tDockerfileContents: schemaDockerfile,\n\t\t\t\tImageName:          tmpImageId,\n\t\t\t\tProgressOutput:     progressOutput,\n\t\t\t}\n\t\t\tif _, err := dockerCommand.ImageBuild(ctx, buildOpts); err != nil {\n\t\t\t\treturn \"\", fmt.Errorf(\"Failed to bundle schema into image: %w\", err)\n\t\t\t}\n\t\t}\n\t\treturn tmpImageId, nil\n\t}\n\n\tconsole.Info(\"Adding labels to image...\")\n\tconsole.Info(\"\")\n\n\t// We used to set the cog_version and config labels in Dockerfile, because we didn't require running the\n\t// built image to get those. But, the escaping of JSON inside a label inside a Dockerfile was gnarly, and\n\t// doesn't seem to be a problem here, so do it here instead.\n\tconfigJSON, err := json.Marshal(cfg)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"Failed to convert config to JSON: %w\", err)\n\t}\n\n\tpipFreeze, err := GeneratePipFreeze(ctx, dockerCommand, tmpImageId)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"Failed to generate pip freeze from image: %w\", err)\n\t}\n\n\tlabels := map[string]string{\n\t\tcommand.CogVersionLabelKey:           global.Version,\n\t\tcommand.CogConfigLabelKey:            string(bytes.TrimSpace(configJSON)),\n\t\tcommand.CogOpenAPISchemaLabelKey:     string(schemaJSON),\n\t\tglobal.LabelNamespace + \"pip_freeze\": pipFreeze,\n\t\t// Mark the image as having an appropriate init entrypoint. We can use this\n\t\t// to decide how/if to shim the image.\n\t\tglobal.LabelNamespace + \"has_init\": \"true\",\n\t}\n\n\tif cogBaseImageName != \"\" {\n\t\tlabels[global.LabelNamespace+\"cog-base-image-name\"] = cogBaseImageName\n\n\t\t// name.Insecure allows HTTP fallback for local/test registries,\n\t\t// consistent with ParseReference calls in pkg/registry/.\n\t\tref, err := name.ParseReference(cogBaseImageName, name.Insecure)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to parse cog base image reference: %w\", err)\n\t\t}\n\n\t\timg, err := remote.Image(ref)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to fetch cog base image: %w\", err)\n\t\t}\n\n\t\tlayers, err := img.Layers()\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to get layers for cog base image: %w\", err)\n\t\t}\n\n\t\tif len(layers) == 0 {\n\t\t\treturn \"\", fmt.Errorf(\"Cog base image has no layers: %s\", cogBaseImageName)\n\t\t}\n\n\t\tlastLayerIndex := len(layers) - 1\n\t\tlayerLayerDigest, err := layers[lastLayerIndex].DiffID()\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to get last layer digest for cog base image: %w\", err)\n\t\t}\n\n\t\tlastLayer := layerLayerDigest.String()\n\t\tconsole.Debugf(\"Last layer of the cog base image: %s\", lastLayer)\n\n\t\tlabels[global.LabelNamespace+\"cog-base-image-last-layer-sha\"] = lastLayer\n\t\tlabels[global.LabelNamespace+\"cog-base-image-last-layer-idx\"] = fmt.Sprintf(\"%d\", lastLayerIndex)\n\t}\n\n\tif commit, err := gitHead(ctx, dir); commit != \"\" && err == nil {\n\t\tlabels[\"org.opencontainers.image.revision\"] = commit\n\t} else {\n\t\tconsole.Debug(\"Unable to determine Git commit\")\n\t}\n\n\tif tag, err := gitTag(ctx, dir); tag != \"\" && err == nil {\n\t\tlabels[\"org.opencontainers.image.version\"] = tag\n\t} else {\n\t\tconsole.Debug(\"Unable to determine Git tag\")\n\t}\n\n\tmaps.Copy(labels, annotations)\n\n\t// The final image ID comes from the label-adding step.\n\t// When schema validation is skipped (cog run), there is no schema file to bundle.\n\tschemaFileToBundle := bundledSchemaFile\n\tif skipSchemaValidation {\n\t\tschemaFileToBundle = \"\"\n\t}\n\timageID, err := BuildAddLabelsAndSchemaToImage(ctx, dockerCommand, tmpImageId, imageName, labels, schemaFileToBundle, progressOutput)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"Failed to add labels to image: %w\", err)\n\t}\n\n\t// We created a temp image, so delete it. Don't \"-f\" so it doesn't blow anything up\n\tif isR8imImage {\n\t\tif err = dockerCommand.RemoveImage(ctx, tmpImageId); err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t}\n\n\treturn imageID, nil\n}\n\n// BuildAddLabelsAndSchemaToImage builds a cog model with labels and schema.\n// Returns the image ID (sha256:...) of the final image.\n//\n// The new image is based on the provided image with the labels and schema file appended to it.\n// tmpName is the source image to build from, image is the final image name/tag.\nfunc BuildAddLabelsAndSchemaToImage(ctx context.Context, dockerClient command.Command, tmpName, image string, labels map[string]string, bundledSchemaFile string, progressOutput string) (string, error) {\n\tvar dockerfile string\n\tif bundledSchemaFile != \"\" {\n\t\tdockerfile = fmt.Sprintf(\"FROM %s\\nCOPY %s .cog\\n\", tmpName, bundledSchemaFile)\n\t} else {\n\t\tdockerfile = fmt.Sprintf(\"FROM %s\\n\", tmpName)\n\t}\n\n\tbuildOpts := command.ImageBuildOptions{\n\t\tDockerfileContents: dockerfile,\n\t\tImageName:          image,\n\t\tLabels:             labels,\n\t\tProgressOutput:     progressOutput,\n\t}\n\n\timageID, err := dockerClient.ImageBuild(ctx, buildOpts)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"Failed to add labels and schema to image: %w\", err)\n\t}\n\treturn imageID, nil\n}\n\n// staticSchemaGenMinSDKVersion is the minimum SDK version that supports\n// static schema generation. Older SDK versions use pydantic-based runtime\n// introspection and must fall back to the legacy Docker-based path.\nconst staticSchemaGenMinSDKVersion = \"0.17.0\"\n\n// canUseStaticSchemaGen returns true if the user has opted in to static schema\n// generation via COG_STATIC_SCHEMA=1 (or \"true\").\n//\n// Even when opted in, returns false when the SDK version is explicitly\n// pinned < 0.17.0, since older SDKs use pydantic-based schemas that the\n// static parser cannot analyze.\nfunc canUseStaticSchemaGen(cfg *config.Config) bool {\n\tenv := strings.ToLower(os.Getenv(\"COG_STATIC_SCHEMA\"))\n\tif env != \"1\" && env != \"true\" {\n\t\treturn false\n\t}\n\n\tsdkVersion := resolveSDKVersion(cfg)\n\tif sdkVersion != \"\" {\n\t\tbase := sdkVersion\n\t\tif m := wheels.BaseVersionRe.FindString(base); m != \"\" {\n\t\t\tbase = m\n\t\t}\n\t\tif ver, err := cogversion.NewVersion(base); err == nil {\n\t\t\tminVer := cogversion.MustVersion(staticSchemaGenMinSDKVersion)\n\t\t\tif !ver.GreaterOrEqual(minVer) {\n\t\t\t\tconsole.Infof(\"SDK version %s < %s, using legacy runtime schema generation\", sdkVersion, staticSchemaGenMinSDKVersion)\n\t\t\t\treturn false\n\t\t\t}\n\t\t}\n\t}\n\treturn true\n}\n\n// resolveSDKVersion determines the SDK version that will be installed in the\n// container, using the same precedence as the Dockerfile generator:\n//  1. COG_SDK_WHEEL env var (parse version from \"pypi:X.Y.Z\")\n//  2. build.sdk_version in cog.yaml\n//  3. Auto-detect from dist/ wheel filename\n//  4. Empty string (latest/unpinned)\nfunc resolveSDKVersion(cfg *config.Config) string {\n\tif envVal := os.Getenv(wheels.CogSDKWheelEnvVar); envVal != \"\" {\n\t\twc := wheels.ParseWheelValue(envVal)\n\t\tif wc != nil && wc.Source == wheels.WheelSourcePyPI && wc.Version != \"\" {\n\t\t\treturn wc.Version\n\t\t}\n\t\treturn \"\"\n\t}\n\tif cfg.Build != nil && cfg.Build.SDKVersion != \"\" {\n\t\treturn cfg.Build.SDKVersion\n\t}\n\tif v := wheels.DetectLocalSDKVersion(); v != \"\" {\n\t\treturn v\n\t}\n\treturn \"\"\n}\n\n// generateStaticSchema runs the Go tree-sitter parser to produce the OpenAPI schema.\n// When both predict and train are configured, it generates both and merges them.\nfunc generateStaticSchema(cfg *config.Config, dir string) ([]byte, error) {\n\tif cfg.Predict == \"\" && cfg.Train == \"\" {\n\t\treturn nil, fmt.Errorf(\"no predict or train reference found in cog.yaml\")\n\t}\n\treturn schema.GenerateCombined(dir, cfg.Predict, cfg.Train, python.ParsePredictor)\n\n}\n\n// writeAndValidateSchema writes the schema JSON to the bundled schema file and\n// validates it as a well-formed OpenAPI 3.0 specification.\nfunc writeAndValidateSchema(schemaJSON []byte) error {\n\tif err := os.MkdirAll(filepath.Dir(bundledSchemaFile), 0o755); err != nil {\n\t\treturn fmt.Errorf(\"failed to create directory for %s: %w\", bundledSchemaFile, err)\n\t}\n\tif err := os.WriteFile(bundledSchemaFile, schemaJSON, 0o644); err != nil {\n\t\treturn fmt.Errorf(\"failed to store bundled schema file %s: %w\", bundledSchemaFile, err)\n\t}\n\n\tloader := openapi3.NewLoader()\n\tloader.IsExternalRefsAllowed = true\n\tdoc, err := loader.LoadFromData(schemaJSON)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Failed to load model schema JSON: %w\", err)\n\t}\n\tif err := doc.Validate(loader.Context); err != nil {\n\t\treturn fmt.Errorf(\"Model schema is invalid: %w\\n\\n%s\", err, string(schemaJSON))\n\t}\n\treturn nil\n}\n\nfunc isGitWorkTree(ctx context.Context, dir string) bool {\n\tctx, cancel := context.WithTimeout(ctx, 3*time.Second)\n\tdefer cancel()\n\n\tout, err := exec.CommandContext(ctx, \"git\", \"-C\", dir, \"rev-parse\", \"--is-inside-work-tree\").Output()\n\tif err != nil {\n\t\treturn false\n\t}\n\n\treturn strings.TrimSpace(string(out)) == \"true\"\n}\n\nfunc gitHead(ctx context.Context, dir string) (string, error) {\n\tif v, ok := os.LookupEnv(\"GITHUB_SHA\"); ok && v != \"\" {\n\t\treturn v, nil\n\t}\n\n\tif isGitWorkTree(ctx, dir) {\n\t\tctx, cancel := context.WithTimeout(ctx, 3*time.Second)\n\t\tdefer cancel()\n\n\t\tout, err := exec.CommandContext(ctx, \"git\", \"-C\", dir, \"rev-parse\", \"HEAD\").Output()\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn string(bytes.TrimSpace(out)), nil\n\t}\n\n\treturn \"\", fmt.Errorf(\"Failed to find HEAD commit: %w\", errGit)\n}\n\nfunc gitTag(ctx context.Context, dir string) (string, error) {\n\tif v, ok := os.LookupEnv(\"GITHUB_REF_NAME\"); ok && v != \"\" {\n\t\treturn v, nil\n\t}\n\n\tif isGitWorkTree(ctx, dir) {\n\t\tctx, cancel := context.WithTimeout(ctx, 3*time.Second)\n\t\tdefer cancel()\n\n\t\tout, err := exec.CommandContext(ctx, \"git\", \"-C\", dir, \"describe\", \"--tags\", \"--dirty\").Output()\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\treturn string(bytes.TrimSpace(out)), nil\n\t}\n\n\treturn \"\", fmt.Errorf(\"Failed to find ref name: %w\", errGit)\n}\n\nfunc buildWeightsImage(ctx context.Context, dockerClient command.Command, dir, dockerfileContents, imageName string, secrets []string, noCache bool, progressOutput string, contextDir string, buildContexts map[string]string) error {\n\tif err := makeDockerignoreForWeightsImage(); err != nil {\n\t\treturn fmt.Errorf(\"Failed to create .dockerignore file: %w\", err)\n\t}\n\tbuildOpts := command.ImageBuildOptions{\n\t\tWorkingDir:         dir,\n\t\tDockerfileContents: dockerfileContents,\n\t\tImageName:          imageName,\n\t\tSecrets:            secrets,\n\t\tNoCache:            noCache,\n\t\tProgressOutput:     progressOutput,\n\t\tEpoch:              &config.BuildSourceEpochTimestamp,\n\t\tContextDir:         contextDir,\n\t\tBuildContexts:      buildContexts,\n\t}\n\tif _, err := dockerClient.ImageBuild(ctx, buildOpts); err != nil {\n\t\treturn fmt.Errorf(\"Failed to build Docker image for model weights: %w\", err)\n\t}\n\treturn nil\n}\n\nfunc buildRunnerImage(ctx context.Context, dockerClient command.Command, dir, dockerfileContents, dockerignoreContents, imageName string, secrets []string, noCache bool, progressOutput string, contextDir string, buildContexts map[string]string) error {\n\tif err := writeDockerignore(dockerignoreContents); err != nil {\n\t\treturn fmt.Errorf(\"Failed to write .dockerignore file with weights included: %w\", err)\n\t}\n\tbuildOpts := command.ImageBuildOptions{\n\t\tWorkingDir:         dir,\n\t\tDockerfileContents: dockerfileContents,\n\t\tImageName:          imageName,\n\t\tSecrets:            secrets,\n\t\tNoCache:            noCache,\n\t\tProgressOutput:     progressOutput,\n\t\tEpoch:              &config.BuildSourceEpochTimestamp,\n\t\tContextDir:         contextDir,\n\t\tBuildContexts:      buildContexts,\n\t}\n\tif _, err := dockerClient.ImageBuild(ctx, buildOpts); err != nil {\n\t\treturn fmt.Errorf(\"Failed to build Docker image: %w\", err)\n\t}\n\tif err := restoreDockerignore(); err != nil {\n\t\treturn fmt.Errorf(\"Failed to restore backup .dockerignore file: %w\", err)\n\t}\n\treturn nil\n}\n\nfunc makeDockerignoreForWeightsImage() error {\n\tif err := backupDockerignore(); err != nil {\n\t\treturn fmt.Errorf(\"Failed to backup .dockerignore file: %w\", err)\n\t}\n\n\tif err := writeDockerignore(dockerfile.DockerignoreHeader); err != nil {\n\t\treturn fmt.Errorf(\"Failed to write .dockerignore file: %w\", err)\n\t}\n\treturn nil\n}\n\nfunc writeDockerignore(contents string) error {\n\t// read existing file contents from .dockerignore.cog.bak if it exists, and append to the new contents\n\tif _, err := os.Stat(dockerignoreBackupPath); err == nil {\n\t\texistingContents, err := os.ReadFile(dockerignoreBackupPath)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tcontents = string(existingContents) + \"\\n\" + contents\n\t}\n\n\treturn os.WriteFile(\".dockerignore\", []byte(contents), 0o644)\n}\n\nfunc backupDockerignore() error {\n\tif _, err := os.Stat(\".dockerignore\"); err != nil {\n\t\tif os.IsNotExist(err) {\n\t\t\t// .dockerignore file does not exist, nothing to backup\n\t\t\treturn nil\n\t\t}\n\t\treturn err\n\t}\n\n\t// rename the .dockerignore file to a new name\n\treturn os.Rename(\".dockerignore\", dockerignoreBackupPath)\n}\n\nfunc restoreDockerignore() error {\n\tif err := os.Remove(\".dockerignore\"); err != nil {\n\t\treturn err\n\t}\n\n\tif _, err := os.Stat(dockerignoreBackupPath); err != nil {\n\t\tif os.IsNotExist(err) {\n\t\t\t// .dockerignore backup file does not exist, nothing to restore\n\t\t\treturn nil\n\t\t}\n\t\treturn err\n\t}\n\n\treturn os.Rename(dockerignoreBackupPath, \".dockerignore\")\n}\n\nfunc checkCompatibleDockerIgnore(dir string) error {\n\tmatcher, err := dockerignore.CreateMatcher(dir)\n\tif err != nil {\n\t\treturn err\n\t}\n\t// If the matcher is nil and we don't have an error, we don't have a .dockerignore to scan.\n\tif matcher == nil {\n\t\treturn nil\n\t}\n\tif matcher.MatchesPath(\".cog\") {\n\t\treturn errors.New(\"The .cog tmp path cannot be ignored by docker in .dockerignore\")\n\t}\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/image/build_test.go",
    "content": "package image\n\nimport (\n\t\"context\"\n\t\"os\"\n\t\"os/exec\"\n\t\"path/filepath\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nvar hasGit = (func() bool {\n\t_, err := exec.LookPath(\"git\")\n\treturn err == nil\n})()\n\nfunc gitRun(ctx context.Context, argv []string, t *testing.T) {\n\tctx, cancel := context.WithTimeout(ctx, 2*time.Second)\n\tt.Cleanup(cancel)\n\n\tout, err := exec.CommandContext(ctx, \"git\", argv...).CombinedOutput()\n\tt.Logf(\"git output:\\n%s\", string(out))\n\n\trequire.NoError(t, err)\n}\n\nfunc setupGitWorkTree(t *testing.T) string {\n\tctx := t.Context()\n\tif !hasGit {\n\t\tt.Skip(\"no git executable available\")\n\t\treturn \"\"\n\t}\n\n\tr := require.New(t)\n\n\ttmp := filepath.Join(t.TempDir(), \"wd\")\n\tr.NoError(os.MkdirAll(tmp, 0o755))\n\n\tgitRun(ctx, []string{\"init\", tmp}, t)\n\tgitRun(ctx, []string{\"-C\", tmp, \"config\", \"user.email\", \"cog@localhost\"}, t)\n\tgitRun(ctx, []string{\"-C\", tmp, \"config\", \"user.name\", \"Cog Tests\"}, t)\n\tgitRun(ctx, []string{\"-C\", tmp, \"commit\", \"--allow-empty\", \"-m\", \"walrus\"}, t)\n\tgitRun(ctx, []string{\"-C\", tmp, \"tag\", \"-a\", \"v0.0.1+walrus\", \"-m\", \"walrus time\"}, t)\n\n\treturn tmp\n}\n\nfunc TestIsGitWorkTree(t *testing.T) {\n\tctx := t.Context()\n\tr := require.New(t)\n\n\tr.False(isGitWorkTree(ctx, \"/dev/null\"))\n\tr.True(isGitWorkTree(ctx, setupGitWorkTree(t)))\n}\n\nfunc TestGitHead(t *testing.T) {\n\tt.Run(\"via github env\", func(t *testing.T) {\n\t\tt.Setenv(\"GITHUB_SHA\", \"fafafaf\")\n\n\t\thead, err := gitHead(t.Context(), \"/dev/null\")\n\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"fafafaf\", head)\n\t})\n\n\tt.Run(\"via git\", func(t *testing.T) {\n\t\ttmp := setupGitWorkTree(t)\n\t\tif tmp == \"\" {\n\t\t\treturn\n\t\t}\n\n\t\tt.Setenv(\"GITHUB_SHA\", \"\")\n\n\t\thead, err := gitHead(t.Context(), tmp)\n\t\trequire.NoError(t, err)\n\t\trequire.NotEqual(t, \"\", head)\n\t})\n\n\tt.Run(\"unavailable\", func(t *testing.T) {\n\t\tt.Setenv(\"GITHUB_SHA\", \"\")\n\n\t\thead, err := gitHead(t.Context(), \"/dev/null\")\n\t\trequire.Error(t, err)\n\t\trequire.Equal(t, \"\", head)\n\t})\n}\n\nfunc TestGitTag(t *testing.T) {\n\tt.Run(\"via github env\", func(t *testing.T) {\n\t\tt.Setenv(\"GITHUB_REF_NAME\", \"v0.0.1+manatee\")\n\n\t\ttag, err := gitTag(t.Context(), \"/dev/null\")\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"v0.0.1+manatee\", tag)\n\t})\n\n\tt.Run(\"via git\", func(t *testing.T) {\n\t\ttmp := setupGitWorkTree(t)\n\t\tif tmp == \"\" {\n\t\t\treturn\n\t\t}\n\n\t\tt.Setenv(\"GITHUB_REF_NAME\", \"\")\n\n\t\ttag, err := gitTag(t.Context(), tmp)\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"v0.0.1+walrus\", tag)\n\t})\n\n\tt.Run(\"unavailable\", func(t *testing.T) {\n\t\tt.Setenv(\"GITHUB_REF_NAME\", \"\")\n\n\t\ttag, err := gitTag(t.Context(), \"/dev/null\")\n\t\trequire.Error(t, err)\n\t\trequire.Equal(t, \"\", tag)\n\t})\n}\n\nfunc TestCanUseStaticSchemaGen(t *testing.T) {\n\t// Helper to build a config with a specific SDK version.\n\tcfgWithSDK := func(version string) *config.Config {\n\t\treturn &config.Config{\n\t\t\tBuild: &config.Build{SDKVersion: version},\n\t\t}\n\t}\n\tnoBuild := &config.Config{}\n\n\ttests := []struct {\n\t\tname     string\n\t\tcfg      *config.Config\n\t\tenvVar   string // COG_STATIC_SCHEMA value\n\t\tsdkWheel string // COG_SDK_WHEEL value\n\t\twant     bool\n\t}{\n\t\t{\n\t\t\tname: \"disabled by default (env not set)\",\n\t\t\tcfg:  cfgWithSDK(\"0.18.0\"),\n\t\t\twant: false,\n\t\t},\n\t\t{\n\t\t\tname:   \"disabled when env is empty string\",\n\t\t\tcfg:    cfgWithSDK(\"0.18.0\"),\n\t\t\tenvVar: \"\",\n\t\t\twant:   false,\n\t\t},\n\t\t{\n\t\t\tname:   \"disabled when env is 0\",\n\t\t\tcfg:    cfgWithSDK(\"0.18.0\"),\n\t\t\tenvVar: \"0\",\n\t\t\twant:   false,\n\t\t},\n\t\t{\n\t\t\tname:   \"enabled when env is 1\",\n\t\t\tcfg:    cfgWithSDK(\"0.18.0\"),\n\t\t\tenvVar: \"1\",\n\t\t\twant:   true,\n\t\t},\n\t\t{\n\t\t\tname:   \"enabled when env is true\",\n\t\t\tcfg:    cfgWithSDK(\"0.18.0\"),\n\t\t\tenvVar: \"true\",\n\t\t\twant:   true,\n\t\t},\n\t\t{\n\t\t\tname:   \"enabled when env is True (mixed case)\",\n\t\t\tcfg:    cfgWithSDK(\"0.18.0\"),\n\t\t\tenvVar: \"True\",\n\t\t\twant:   true,\n\t\t},\n\t\t{\n\t\t\tname:   \"enabled when env is TRUE (upper case)\",\n\t\t\tcfg:    cfgWithSDK(\"0.18.0\"),\n\t\t\tenvVar: \"TRUE\",\n\t\t\twant:   true,\n\t\t},\n\t\t{\n\t\t\tname:   \"disabled for old SDK even when opted in\",\n\t\t\tcfg:    cfgWithSDK(\"0.16.12\"),\n\t\t\tenvVar: \"1\",\n\t\t\twant:   false,\n\t\t},\n\t\t{\n\t\t\tname:   \"disabled for pre-release old SDK\",\n\t\t\tcfg:    cfgWithSDK(\"0.16.0a1\"),\n\t\t\tenvVar: \"1\",\n\t\t\twant:   false,\n\t\t},\n\t\t{\n\t\t\tname:   \"enabled for SDK 0.17.0 when opted in\",\n\t\t\tcfg:    cfgWithSDK(\"0.17.0\"),\n\t\t\tenvVar: \"1\",\n\t\t\twant:   true,\n\t\t},\n\t\t{\n\t\t\tname:   \"enabled for SDK 0.18.0 when opted in\",\n\t\t\tcfg:    cfgWithSDK(\"0.18.0\"),\n\t\t\tenvVar: \"1\",\n\t\t\twant:   true,\n\t\t},\n\t\t{\n\t\t\tname:   \"enabled for unpinned SDK when opted in\",\n\t\t\tcfg:    noBuild,\n\t\t\tenvVar: \"1\",\n\t\t\twant:   true,\n\t\t},\n\t\t{\n\t\t\tname:     \"disabled for old SDK via COG_SDK_WHEEL even when opted in\",\n\t\t\tcfg:      noBuild,\n\t\t\tenvVar:   \"1\",\n\t\t\tsdkWheel: \"pypi:0.16.12\",\n\t\t\twant:     false,\n\t\t},\n\t\t{\n\t\t\tname:     \"enabled for new SDK via COG_SDK_WHEEL when opted in\",\n\t\t\tcfg:      noBuild,\n\t\t\tenvVar:   \"1\",\n\t\t\tsdkWheel: \"pypi:0.18.0\",\n\t\t\twant:     true,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tt.Setenv(\"COG_STATIC_SCHEMA\", tt.envVar)\n\t\t\tt.Setenv(\"COG_SDK_WHEEL\", tt.sdkWheel)\n\n\t\t\tgot := canUseStaticSchemaGen(tt.cfg)\n\t\t\trequire.Equal(t, tt.want, got)\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "pkg/image/config.go",
    "content": "package image\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\n\t\"github.com/docker/docker/api/types/image\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n)\n\nfunc CogConfigFromManifest(ctx context.Context, manifest *image.InspectResponse) (*config.Config, error) {\n\tconfigString := manifest.Config.Labels[command.CogConfigLabelKey]\n\tif configString == \"\" {\n\t\t// Deprecated. Remove for 1.0.\n\t\tconfigString = manifest.Config.Labels[\"org.cogmodel.config\"]\n\t}\n\tif configString == \"\" {\n\t\t// TODO[md]: find the tag/ref and return that in the error instead of the ID\n\t\treturn nil, fmt.Errorf(\"Image %s does not appear to be a Cog model\", friendlyName(manifest))\n\t}\n\tconf := new(config.Config)\n\tif err := json.Unmarshal([]byte(configString), conf); err != nil {\n\t\t// TODO[md]: find the tag/ref and return that in the error instead of the ID\n\t\treturn nil, fmt.Errorf(\"Failed to parse config from %s: %w\", friendlyName(manifest), err)\n\t}\n\treturn conf, nil\n}\n\nfunc friendlyName(manifest *image.InspectResponse) string {\n\t// this appears to get the base image name, which we don't really want\n\t// name := manifest.Config.Labels[\"org.opencontainers.image.title\"]\n\t// if name != \"\" {\n\t// \treturn name\n\t// }\n\n\tif len(manifest.RepoTags) > 0 {\n\t\treturn manifest.RepoTags[0]\n\t}\n\n\treturn manifest.ID\n}\n"
  },
  {
    "path": "pkg/image/openapi_schema.go",
    "content": "package image\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/json\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// GenerateOpenAPISchema generates the OpenAPI schema by running the built Docker\n// image with `python -m cog.command.openapi_schema`. This is the legacy path used\n// for SDK versions < 0.17.0 where the schema must be generated at runtime via\n// pydantic introspection rather than static analysis.\n//\n// sourceDir, when non-empty, is volume-mounted as /src. This is needed for\n// ExcludeSource builds (cog serve/predict/train) where COPY . /src was skipped.\nfunc GenerateOpenAPISchema(ctx context.Context, dockerClient command.Command, imageName string, enableGPU bool, sourceDir string) (map[string]any, error) {\n\tconsole.Debugf(\"=== image.GenerateOpenAPISchema %s\", imageName)\n\tvar stdout bytes.Buffer\n\tvar stderr bytes.Buffer\n\n\tgpus := \"\"\n\tif enableGPU {\n\t\tgpus = \"all\"\n\t}\n\n\trunOpts := command.RunOptions{\n\t\tImage: imageName,\n\t\tArgs: []string{\n\t\t\t\"python\", \"-m\", \"cog.command.openapi_schema\",\n\t\t},\n\t\tGPUs: gpus,\n\t}\n\tif sourceDir != \"\" {\n\t\trunOpts.Volumes = []command.Volume{{Source: sourceDir, Destination: \"/src\"}}\n\t}\n\n\terr := docker.RunWithIO(ctx, dockerClient, runOpts, nil, &stdout, &stderr)\n\n\tif enableGPU && err == docker.ErrMissingDeviceDriver {\n\t\tconsole.Debug(stdout.String())\n\t\tconsole.Debug(stderr.String())\n\t\tconsole.Debug(\"Missing device driver, re-trying without GPU\")\n\t\treturn GenerateOpenAPISchema(ctx, dockerClient, imageName, false, sourceDir)\n\t}\n\n\tif err != nil {\n\t\tconsole.Info(stdout.String())\n\t\tconsole.Info(stderr.String())\n\t\treturn nil, err\n\t}\n\n\tvar schema map[string]any\n\tif err := json.Unmarshal(stdout.Bytes(), &schema); err != nil {\n\t\tconsole.Info(stdout.String())\n\t\tconsole.Info(stderr.String())\n\t\treturn nil, err\n\t}\n\n\treturn schema, nil\n}\n"
  },
  {
    "path": "pkg/image/pip_freeze.go",
    "content": "package image\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// GeneratePipFreeze by running a pip freeze on the image.\n// This will be run as part of the build process then added as a label to the image.\nfunc GeneratePipFreeze(ctx context.Context, dockerClient command.Command, imageName string) (string, error) {\n\tvar stdout bytes.Buffer\n\tvar stderr bytes.Buffer\n\n\targs := []string{\"python\", \"-m\", \"pip\", \"freeze\"}\n\terr := docker.RunWithIO(ctx, dockerClient, command.RunOptions{\n\t\tImage: imageName,\n\t\tArgs:  args,\n\t}, nil, &stdout, &stderr)\n\n\tif err != nil {\n\t\tconsole.Info(stdout.String())\n\t\tconsole.Info(stderr.String())\n\t\treturn \"\", err\n\t}\n\n\treturn stdout.String(), nil\n}\n"
  },
  {
    "path": "pkg/model/artifact.go",
    "content": "package model\n\nimport v1 \"github.com/google/go-containerregistry/pkg/v1\"\n\n// ArtifactType identifies the kind of artifact.\ntype ArtifactType int\n\nconst (\n\t// ArtifactTypeImage is a container image artifact.\n\tArtifactTypeImage ArtifactType = iota + 1\n\t// ArtifactTypeWeight is a model weight artifact.\n\tArtifactTypeWeight\n)\n\n// String returns the human-readable name of the artifact type.\nfunc (t ArtifactType) String() string {\n\tswitch t {\n\tcase ArtifactTypeImage:\n\t\treturn \"image\"\n\tcase ArtifactTypeWeight:\n\t\treturn \"weight\"\n\tdefault:\n\t\treturn \"unknown\"\n\t}\n}\n\n// ArtifactSpec declares what artifact will be produced.\n// It contains all inputs needed to build that artifact.\n// Specs are derived from analyzing the Source (cog.yaml + project directory).\ntype ArtifactSpec interface {\n\tType() ArtifactType\n\tName() string\n}\n\n// Artifact is the immutable result of building a spec.\n// It contains the OCI descriptor and enough information for a pusher to upload it.\ntype Artifact interface {\n\tType() ArtifactType\n\tName() string\n\tDescriptor() v1.Descriptor\n}\n"
  },
  {
    "path": "pkg/model/artifact_image.go",
    "content": "package model\n\nimport (\n\t\"encoding/json\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/global\"\n)\n\n// ImageSource indicates where an image was loaded from.\ntype ImageSource string\n\nconst (\n\tImageSourceLocal  ImageSource = \"local\"  // Docker daemon\n\tImageSourceRemote ImageSource = \"remote\" // Registry\n\tImageSourceBuild  ImageSource = \"build\"  // Just built\n)\n\n// Platform describes the OS and architecture of an image.\ntype Platform struct {\n\tOS           string\n\tArchitecture string\n\tVariant      string\n}\n\n// Label keys for Cog-specific metadata stored in image labels.\nvar (\n\tLabelConfig          = global.LabelNamespace + \"config\"\n\tLabelVersion         = global.LabelNamespace + \"version\"\n\tLabelOpenAPISchema   = global.LabelNamespace + \"openapi_schema\"\n\tLabelWeightsManifest = global.LabelNamespace + \"r8_weights_manifest\"\n)\n\n// =============================================================================\n// ImageSpec\n// =============================================================================\n\n// ImageSpecOption configures optional fields on ImageSpec.\ntype ImageSpecOption func(*ImageSpec)\n\n// WithImageSecrets sets build-time secrets for the image build.\nfunc WithImageSecrets(secrets []string) ImageSpecOption {\n\treturn func(s *ImageSpec) {\n\t\ts.Secrets = secrets\n\t}\n}\n\n// WithImageNoCache disables build cache for the image build.\nfunc WithImageNoCache(noCache bool) ImageSpecOption {\n\treturn func(s *ImageSpec) {\n\t\ts.NoCache = noCache\n\t}\n}\n\n// ImageSpec declares an image to be built.\n// It implements ArtifactSpec.\n//\n// TODO: ImageBuilder currently reads build options from BuildOptions (passed at\n// construction) rather than from ImageSpec fields. When the build pipeline fully\n// migrates to specs, ImageName/Secrets/NoCache should be the source of truth.\ntype ImageSpec struct {\n\tname      string\n\tImageName string\n\tSecrets   []string\n\tNoCache   bool\n}\n\n// NewImageSpec creates an ImageSpec with the given name and image name.\n// Optional configuration can be provided via ImageSpecOption functions.\nfunc NewImageSpec(name, imageName string, opts ...ImageSpecOption) *ImageSpec {\n\ts := &ImageSpec{\n\t\tname:      name,\n\t\tImageName: imageName,\n\t}\n\tfor _, opt := range opts {\n\t\topt(s)\n\t}\n\treturn s\n}\n\n// Type returns ArtifactTypeImage.\nfunc (s *ImageSpec) Type() ArtifactType { return ArtifactTypeImage }\n\n// Name returns the spec's logical name.\nfunc (s *ImageSpec) Name() string { return s.name }\n\n// =============================================================================\n// ImageArtifact\n// =============================================================================\n\n// ImageArtifact represents an OCI container image.\n// It serves as both the build artifact (in Model.Artifacts) and the general-purpose\n// image metadata type throughout the codebase.\n// It implements the Artifact interface.\ntype ImageArtifact struct {\n\t// Artifact fields (set when used as a build artifact)\n\tname       string\n\tdescriptor v1.Descriptor\n\n\t// Image metadata\n\tReference string            // Full image reference (e.g., \"r8.im/user/model:latest\")\n\tDigest    string            // Content-addressable digest (sha256:...)\n\tLabels    map[string]string // Docker/OCI image labels\n\tPlatform  *Platform         // OS/architecture\n\tSource    ImageSource       // Where loaded from (local/remote/build)\n}\n\n// NewImageArtifact creates an ImageArtifact from a build result.\nfunc NewImageArtifact(name string, desc v1.Descriptor, reference string) *ImageArtifact {\n\treturn &ImageArtifact{\n\t\tname:       name,\n\t\tdescriptor: desc,\n\t\tReference:  reference,\n\t}\n}\n\n// Type returns ArtifactTypeImage.\nfunc (a *ImageArtifact) Type() ArtifactType { return ArtifactTypeImage }\n\n// Name returns the artifact's logical name.\nfunc (a *ImageArtifact) Name() string { return a.name }\n\n// Descriptor returns the OCI descriptor for this image.\nfunc (a *ImageArtifact) Descriptor() v1.Descriptor { return a.descriptor }\n\n// =============================================================================\n// Image metadata methods (formerly on *Image)\n// =============================================================================\n\n// IsCogModel returns true if this image has Cog labels indicating it's a Cog model.\nfunc (a *ImageArtifact) IsCogModel() bool {\n\tif a.Labels == nil {\n\t\treturn false\n\t}\n\t_, ok := a.Labels[LabelConfig]\n\treturn ok\n}\n\n// CogVersion returns the Cog version that built this image, or empty string if not set.\nfunc (a *ImageArtifact) CogVersion() string {\n\tif a.Labels == nil {\n\t\treturn \"\"\n\t}\n\treturn a.Labels[LabelVersion]\n}\n\n// Config returns the raw cog.yaml config stored in image labels, or empty string if not set.\nfunc (a *ImageArtifact) Config() string {\n\tif a.Labels == nil {\n\t\treturn \"\"\n\t}\n\treturn a.Labels[LabelConfig]\n}\n\n// OpenAPISchema returns the OpenAPI schema stored in image labels, or empty string if not set.\nfunc (a *ImageArtifact) OpenAPISchema() string {\n\tif a.Labels == nil {\n\t\treturn \"\"\n\t}\n\treturn a.Labels[LabelOpenAPISchema]\n}\n\n// ParsedConfig returns the parsed cog.yaml config from image labels.\n// Returns nil without error if no config label is present.\n// Returns error if the label contains invalid JSON.\nfunc (a *ImageArtifact) ParsedConfig() (*config.Config, error) {\n\traw := a.Config()\n\tif raw == \"\" {\n\t\treturn nil, nil\n\t}\n\n\tcfg := new(config.Config)\n\tif err := json.Unmarshal([]byte(raw), cfg); err != nil {\n\t\treturn nil, err\n\t}\n\treturn cfg, nil\n}\n\n// ParsedOpenAPISchema returns the parsed OpenAPI schema from image labels.\n// Returns nil without error if no schema label is present.\n// Returns error if the label contains invalid JSON.\nfunc (a *ImageArtifact) ParsedOpenAPISchema() (*openapi3.T, error) {\n\traw := a.OpenAPISchema()\n\tif raw == \"\" {\n\t\treturn nil, nil\n\t}\n\n\tloader := openapi3.NewLoader()\n\tschema, err := loader.LoadFromData([]byte(raw))\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn schema, nil\n}\n\n// ToModel converts the ImageArtifact to a Model by parsing its labels.\n// Returns error if the image is not a valid Cog model or if labels contain invalid JSON.\nfunc (a *ImageArtifact) ToModel() (*Model, error) {\n\tif !a.IsCogModel() {\n\t\treturn nil, ErrNotCogModel\n\t}\n\n\tcfg, err := a.ParsedConfig()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tschema, err := a.ParsedOpenAPISchema()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &Model{\n\t\tImage:      a,\n\t\tConfig:     cfg,\n\t\tSchema:     schema,\n\t\tCogVersion: a.CogVersion(),\n\t}, nil\n}\n"
  },
  {
    "path": "pkg/model/artifact_image_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestImageSpec_ImplementsArtifactSpec(t *testing.T) {\n\tspec := NewImageSpec(\"model\", \"r8.im/user/model:latest\")\n\n\tvar _ ArtifactSpec = spec // compile-time interface check\n\n\trequire.Equal(t, ArtifactTypeImage, spec.Type())\n\trequire.Equal(t, \"model\", spec.Name())\n}\n\nfunc TestImageSpec_Fields(t *testing.T) {\n\tspec := NewImageSpec(\"model\", \"r8.im/user/model:latest\",\n\t\tWithImageSecrets([]string{\"secret1\", \"secret2\"}),\n\t\tWithImageNoCache(true),\n\t)\n\n\trequire.Equal(t, \"r8.im/user/model:latest\", spec.ImageName)\n\trequire.Equal(t, []string{\"secret1\", \"secret2\"}, spec.Secrets)\n\trequire.True(t, spec.NoCache)\n}\n\nfunc TestImageSpec_DefaultFields(t *testing.T) {\n\tspec := NewImageSpec(\"model\", \"myimage:latest\")\n\n\trequire.Equal(t, \"myimage:latest\", spec.ImageName)\n\trequire.Nil(t, spec.Secrets)\n\trequire.False(t, spec.NoCache)\n}\n\nfunc TestImageArtifact_ImplementsArtifact(t *testing.T) {\n\tdesc := v1.Descriptor{\n\t\tDigest: v1.Hash{Algorithm: \"sha256\", Hex: \"abc123\"},\n\t\tSize:   1024,\n\t}\n\tartifact := NewImageArtifact(\"model\", desc, \"r8.im/user/model@sha256:abc123\")\n\n\tvar _ Artifact = artifact // compile-time interface check\n\n\trequire.Equal(t, ArtifactTypeImage, artifact.Type())\n\trequire.Equal(t, \"model\", artifact.Name())\n\trequire.Equal(t, desc, artifact.Descriptor())\n}\n\nfunc TestImageArtifact_Reference(t *testing.T) {\n\tdesc := v1.Descriptor{\n\t\tDigest: v1.Hash{Algorithm: \"sha256\", Hex: \"abc123\"},\n\t\tSize:   1024,\n\t}\n\tartifact := NewImageArtifact(\"model\", desc, \"r8.im/user/model@sha256:abc123\")\n\n\trequire.Equal(t, \"r8.im/user/model@sha256:abc123\", artifact.Reference)\n}\n"
  },
  {
    "path": "pkg/model/artifact_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestArtifactType_String(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\tat     ArtifactType\n\t\texpect string\n\t}{\n\t\t{name: \"image type\", at: ArtifactTypeImage, expect: \"image\"},\n\t\t{name: \"weight type\", at: ArtifactTypeWeight, expect: \"weight\"},\n\t\t{name: \"zero value\", at: ArtifactType(0), expect: \"unknown\"},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\trequire.Equal(t, tt.expect, tt.at.String())\n\t\t})\n\t}\n}\n\nfunc TestArtifactType_Values(t *testing.T) {\n\t// Ensure types are distinct\n\trequire.NotEqual(t, ArtifactTypeImage, ArtifactTypeWeight)\n}\n"
  },
  {
    "path": "pkg/model/artifact_weight.go",
    "content": "package model\n\nimport (\n\t\"time\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n)\n\n// Media types for weight artifacts (OCI 1.1 conventions).\nconst (\n\t// MediaTypeWeightArtifact is the artifactType for weight manifests.\n\tMediaTypeWeightArtifact = \"application/vnd.cog.weight.v1\"\n\t// MediaTypeWeightConfig is the media type for weight config blobs.\n\tMediaTypeWeightConfig = \"application/vnd.cog.weight.config.v1+json\"\n\t// MediaTypeWeightLayer is the media type for uncompressed weight layers.\n\tMediaTypeWeightLayer = \"application/vnd.cog.weight.layer.v1\"\n\t// MediaTypeWeightLayerGzip is the media type for gzip-compressed weight layers.\n\tMediaTypeWeightLayerGzip = \"application/vnd.cog.weight.layer.v1+gzip\"\n\t// MediaTypeWeightLayerZstd is the media type for zstd-compressed weight layers (future).\n\tMediaTypeWeightLayerZstd = \"application/vnd.cog.weight.layer.v1+zstd\"\n)\n\n// Annotation keys for weight file layers in OCI manifests.\nconst (\n\tAnnotationWeightName             = \"vnd.cog.weight.name\"\n\tAnnotationWeightDest             = \"vnd.cog.weight.dest\"\n\tAnnotationWeightDigestOriginal   = \"vnd.cog.weight.digest.original\"\n\tAnnotationWeightSizeUncompressed = \"vnd.cog.weight.size.uncompressed\"\n)\n\n// WeightSpec declares a weight artifact to be built.\n// It implements ArtifactSpec.\ntype WeightSpec struct {\n\tname string\n\t// Source is the local file path to the weight file.\n\tSource string\n\t// Target is the container mount path for this weight.\n\tTarget string\n}\n\n// NewWeightSpec creates a WeightSpec with the given name, source path, and target mount path.\nfunc NewWeightSpec(name, source, target string) *WeightSpec {\n\treturn &WeightSpec{\n\t\tname:   name,\n\t\tSource: source,\n\t\tTarget: target,\n\t}\n}\n\n// Type returns ArtifactTypeWeight.\nfunc (s *WeightSpec) Type() ArtifactType { return ArtifactTypeWeight }\n\n// Name returns the spec's logical name.\nfunc (s *WeightSpec) Name() string { return s.name }\n\n// WeightArtifact is a built weight artifact ready to push as an OCI artifact.\n// It implements Artifact.\ntype WeightArtifact struct {\n\tname       string\n\tdescriptor v1.Descriptor\n\n\t// FilePath is the local file path to the weight data (for pushing layers).\n\tFilePath string\n\t// Target is the container mount path for this weight.\n\tTarget string\n\t// Config is the weight metadata for the config blob.\n\tConfig WeightConfig\n}\n\n// NewWeightArtifact creates a WeightArtifact from a build result.\nfunc NewWeightArtifact(name string, desc v1.Descriptor, filePath, target string, cfg WeightConfig) *WeightArtifact {\n\treturn &WeightArtifact{\n\t\tname:       name,\n\t\tdescriptor: desc,\n\t\tFilePath:   filePath,\n\t\tTarget:     target,\n\t\tConfig:     cfg,\n\t}\n}\n\n// Type returns ArtifactTypeWeight.\nfunc (a *WeightArtifact) Type() ArtifactType { return ArtifactTypeWeight }\n\n// Name returns the artifact's logical name.\nfunc (a *WeightArtifact) Name() string { return a.name }\n\n// Descriptor returns the OCI descriptor for this weight artifact.\nfunc (a *WeightArtifact) Descriptor() v1.Descriptor { return a.descriptor }\n\n// WeightConfig contains metadata about a weight artifact.\n// This is serialized as the config blob in the OCI manifest.\n// The schema is versioned via SchemaVersion to allow evolution.\ntype WeightConfig struct {\n\tSchemaVersion string    `json:\"schemaVersion\"`\n\tCogVersion    string    `json:\"cogVersion\"`\n\tName          string    `json:\"name\"`\n\tTarget        string    `json:\"target\"`\n\tCreated       time.Time `json:\"created\"` // RFC 3339 format when serialized to JSON\n}\n"
  },
  {
    "path": "pkg/model/artifact_weight_test.go",
    "content": "package model\n\nimport (\n\t\"encoding/json\"\n\t\"testing\"\n\t\"time\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestWeightSpec_ImplementsArtifactSpec(t *testing.T) {\n\tspec := NewWeightSpec(\"my-model-weights\", \"/data/weights.bin\", \"/weights/model.bin\")\n\n\tvar _ ArtifactSpec = spec // compile-time interface check\n\n\trequire.Equal(t, ArtifactTypeWeight, spec.Type())\n\trequire.Equal(t, \"my-model-weights\", spec.Name())\n}\n\nfunc TestWeightSpec_Fields(t *testing.T) {\n\tspec := NewWeightSpec(\"llama-7b\", \"/data/llama-7b.safetensors\", \"/weights/llama-7b.safetensors\")\n\n\trequire.Equal(t, \"/data/llama-7b.safetensors\", spec.Source)\n\trequire.Equal(t, \"/weights/llama-7b.safetensors\", spec.Target)\n}\n\nfunc TestWeightArtifact_ImplementsArtifact(t *testing.T) {\n\tdesc := v1.Descriptor{\n\t\tDigest: v1.Hash{Algorithm: \"sha256\", Hex: \"def456\"},\n\t\tSize:   4096,\n\t}\n\tcfg := WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"my-weights\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Date(2026, 2, 5, 12, 0, 0, 0, time.UTC),\n\t}\n\tartifact := NewWeightArtifact(\"my-weights\", desc, \"/data/weights.bin\", \"/weights/model.bin\", cfg)\n\n\tvar _ Artifact = artifact // compile-time interface check\n\n\trequire.Equal(t, ArtifactTypeWeight, artifact.Type())\n\trequire.Equal(t, \"my-weights\", artifact.Name())\n\trequire.Equal(t, desc, artifact.Descriptor())\n}\n\nfunc TestWeightArtifact_Fields(t *testing.T) {\n\tdesc := v1.Descriptor{\n\t\tDigest: v1.Hash{Algorithm: \"sha256\", Hex: \"def456\"},\n\t\tSize:   4096,\n\t}\n\tcfg := WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"my-weights\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Date(2026, 2, 5, 12, 0, 0, 0, time.UTC),\n\t}\n\tartifact := NewWeightArtifact(\"my-weights\", desc, \"/data/weights.bin\", \"/weights/model.bin\", cfg)\n\n\trequire.Equal(t, \"/data/weights.bin\", artifact.FilePath)\n\trequire.Equal(t, \"/weights/model.bin\", artifact.Target)\n\trequire.Equal(t, cfg, artifact.Config)\n}\n\nfunc TestWeightConfig_JSONRoundTrip(t *testing.T) {\n\toriginal := WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"llama-7b\",\n\t\tTarget:        \"/weights/llama-7b\",\n\t\tCreated:       time.Date(2026, 2, 5, 12, 0, 0, 0, time.UTC),\n\t}\n\n\tdata, err := json.Marshal(original)\n\trequire.NoError(t, err)\n\n\t// Verify JSON structure\n\tvar raw map[string]any\n\terr = json.Unmarshal(data, &raw)\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"1.0\", raw[\"schemaVersion\"])\n\trequire.Equal(t, \"0.15.0\", raw[\"cogVersion\"])\n\trequire.Equal(t, \"llama-7b\", raw[\"name\"])\n\trequire.Equal(t, \"/weights/llama-7b\", raw[\"target\"])\n\n\t// Round-trip\n\tvar decoded WeightConfig\n\terr = json.Unmarshal(data, &decoded)\n\trequire.NoError(t, err)\n\trequire.Equal(t, original.SchemaVersion, decoded.SchemaVersion)\n\trequire.Equal(t, original.CogVersion, decoded.CogVersion)\n\trequire.Equal(t, original.Name, decoded.Name)\n\trequire.Equal(t, original.Target, decoded.Target)\n\trequire.True(t, original.Created.Equal(decoded.Created))\n}\n\nfunc TestWeightMediaTypeConstants(t *testing.T) {\n\t// Verify media type constants have expected values\n\trequire.Equal(t, \"application/vnd.cog.weight.v1\", MediaTypeWeightArtifact)\n\trequire.Equal(t, \"application/vnd.cog.weight.config.v1+json\", MediaTypeWeightConfig)\n\trequire.Equal(t, \"application/vnd.cog.weight.layer.v1\", MediaTypeWeightLayer)\n\trequire.Equal(t, \"application/vnd.cog.weight.layer.v1+gzip\", MediaTypeWeightLayerGzip)\n}\n"
  },
  {
    "path": "pkg/model/builder.go",
    "content": "package model\n\nimport \"context\"\n\n// Builder builds an artifact from a spec.\n// Each builder handles one artifact type (image, weight, etc.).\ntype Builder interface {\n\tBuild(ctx context.Context, spec ArtifactSpec) (Artifact, error)\n}\n"
  },
  {
    "path": "pkg/model/builder_test.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/stretchr/testify/require\"\n)\n\n// mockBuilder is a test double that implements the Builder interface.\ntype mockBuilder struct {\n\tbuildFn func(ctx context.Context, spec ArtifactSpec) (Artifact, error)\n}\n\nfunc (m *mockBuilder) Build(ctx context.Context, spec ArtifactSpec) (Artifact, error) {\n\treturn m.buildFn(ctx, spec)\n}\n\nfunc TestBuilderInterface_Satisfiable(t *testing.T) {\n\t// Compile-time check: mockBuilder satisfies Builder.\n\tvar _ Builder = &mockBuilder{}\n\n\t// Runtime check: a mock builder can be called and returns an artifact.\n\tmb := &mockBuilder{\n\t\tbuildFn: func(_ context.Context, spec ArtifactSpec) (Artifact, error) {\n\t\t\treturn NewImageArtifact(spec.Name(), v1.Descriptor{}, \"test-ref\"), nil\n\t\t},\n\t}\n\n\tartifact, err := mb.Build(context.Background(), NewImageSpec(\"test\", \"test-image\"))\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"test\", artifact.Name())\n}\n"
  },
  {
    "path": "pkg/model/errors.go",
    "content": "package model\n\nimport \"errors\"\n\n// Sentinel errors for Resolver operations.\nvar (\n\t// ErrNotCogModel indicates the image exists but is not a valid Cog model.\n\t// This occurs when the image lacks the required run.cog.config label.\n\tErrNotCogModel = errors.New(\"image is not a Cog model\")\n\n\t// ErrNotFound indicates the image was not found in the requested location(s).\n\tErrNotFound = errors.New(\"image not found\")\n)\n"
  },
  {
    "path": "pkg/model/errors_test.go",
    "content": "package model\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestSentinelErrors(t *testing.T) {\n\t// Test that sentinel errors can be wrapped and unwrapped\n\tt.Run(\"ErrNotCogModel can be wrapped and detected\", func(t *testing.T) {\n\t\twrapped := fmt.Errorf(\"failed to inspect image: %w\", ErrNotCogModel)\n\t\trequire.True(t, errors.Is(wrapped, ErrNotCogModel))\n\t})\n\n\tt.Run(\"ErrNotFound can be wrapped and detected\", func(t *testing.T) {\n\t\twrapped := fmt.Errorf(\"image my-image:latest: %w\", ErrNotFound)\n\t\trequire.True(t, errors.Is(wrapped, ErrNotFound))\n\t})\n\n\tt.Run(\"errors are distinct\", func(t *testing.T) {\n\t\trequire.False(t, errors.Is(ErrNotCogModel, ErrNotFound))\n\t\trequire.False(t, errors.Is(ErrNotFound, ErrNotCogModel))\n\t})\n}\n"
  },
  {
    "path": "pkg/model/factory.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/image\"\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// Factory is the build backend interface.\n// Different implementations handle different build strategies.\ntype Factory interface {\n\t// Build creates a Docker image from source and returns ImageArtifact metadata.\n\t// For dev mode (cog serve), set ExcludeSource=true in BuildOptions to skip\n\t// COPY . /src — the source directory is volume-mounted at runtime instead.\n\tBuild(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error)\n\n\t// Name returns the factory name for logging/debugging.\n\tName() string\n}\n\n// DockerfileFactory wraps existing Dockerfile-based build.\ntype DockerfileFactory struct {\n\tdocker   command.Command\n\tregistry registry.Client\n}\n\n// NewDockerfileFactory creates a Factory that uses the existing Dockerfile-based build.\nfunc NewDockerfileFactory(docker command.Command, registry registry.Client) Factory {\n\treturn &DockerfileFactory{docker: docker, registry: registry}\n}\n\n// Name returns the factory name.\nfunc (f *DockerfileFactory) Name() string {\n\treturn \"dockerfile\"\n}\n\n// Build delegates to the existing image.Build() function.\nfunc (f *DockerfileFactory) Build(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error) {\n\timageID, err := image.Build(\n\t\tctx,\n\t\tsrc.Config,\n\t\tsrc.ProjectDir,\n\t\topts.ImageName,\n\t\tsrc.ConfigFilename,\n\t\topts.Secrets,\n\t\topts.NoCache,\n\t\topts.SeparateWeights,\n\t\topts.UseCudaBaseImage,\n\t\topts.ProgressOutput,\n\t\topts.SchemaFile,\n\t\topts.DockerfileFile,\n\t\topts.UseCogBaseImage,\n\t\topts.Strip,\n\t\topts.Precompile,\n\t\topts.ExcludeSource,\n\t\topts.SkipSchemaValidation,\n\t\topts.SkipLabels,\n\t\topts.Annotations,\n\t\tf.docker,\n\t\tf.registry,\n\t)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &ImageArtifact{\n\t\tReference: opts.ImageName,\n\t\tDigest:    imageID,\n\t\tSource:    ImageSourceBuild,\n\t}, nil\n}\n\n// defaultFactory returns a Factory based on environment variables.\n// It checks COG_BUILDER and COGPACK to select the appropriate backend.\n//\n// TODO: When FrontendFactory is implemented, check COG_BUILDER env var.\n// TODO: When CogpacksFactory is implemented, check COGPACK env var.\nfunc defaultFactory(docker command.Command, registry registry.Client) Factory {\n\treturn NewDockerfileFactory(docker, registry)\n}\n"
  },
  {
    "path": "pkg/model/factory_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n\t\"github.com/replicate/cog/pkg/registry/registrytest\"\n)\n\nfunc TestDockerfileFactory_Name(t *testing.T) {\n\tdocker := dockertest.NewMockCommand()\n\tregistry := registrytest.NewMockRegistryClient()\n\n\tfactory := NewDockerfileFactory(docker, registry)\n\n\trequire.Equal(t, \"dockerfile\", factory.Name())\n}\n\nfunc TestDockerfileFactory_ImplementsInterface(t *testing.T) {\n\tdocker := dockertest.NewMockCommand()\n\tregistry := registrytest.NewMockRegistryClient()\n\n\t// Verify that DockerfileFactory implements the Factory interface\n\tvar _ = NewDockerfileFactory(docker, registry)\n}\n\nfunc TestDefaultFactory_ReturnsDockerfileFactory(t *testing.T) {\n\tdocker := dockertest.NewMockCommand()\n\tregistry := registrytest.NewMockRegistryClient()\n\n\tfactory := defaultFactory(docker, registry)\n\n\trequire.Equal(t, \"dockerfile\", factory.Name())\n}\n"
  },
  {
    "path": "pkg/model/format.go",
    "content": "package model\n\nimport \"os\"\n\n// TODO(md): OCIIndexEnabled is a temporary gate for the OCI Image Index push path.\n// When COG_OCI_INDEX=1, builds produce weight artifacts and pushes create an OCI\n// Image Index instead of a single image manifest. Remove this gate (and always use\n// the index path) once we've validated index compatibility with all registries.\nfunc OCIIndexEnabled() bool {\n\treturn os.Getenv(\"COG_OCI_INDEX\") == \"1\"\n}\n"
  },
  {
    "path": "pkg/model/format_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestOCIIndexEnabled_Default(t *testing.T) {\n\tt.Setenv(\"COG_OCI_INDEX\", \"\")\n\trequire.False(t, OCIIndexEnabled())\n}\n\nfunc TestOCIIndexEnabled_Enabled(t *testing.T) {\n\tt.Setenv(\"COG_OCI_INDEX\", \"1\")\n\trequire.True(t, OCIIndexEnabled())\n}\n\nfunc TestOCIIndexEnabled_OtherValue(t *testing.T) {\n\tt.Setenv(\"COG_OCI_INDEX\", \"0\")\n\trequire.False(t, OCIIndexEnabled())\n}\n"
  },
  {
    "path": "pkg/model/hash.go",
    "content": "package model\n\nimport (\n\t\"crypto/sha256\"\n\t\"encoding/hex\"\n\t\"io\"\n\t\"os\"\n)\n\n// hashFile computes SHA256 digest and size of a file by streaming.\nfunc hashFile(path string) (digest string, size int64, err error) {\n\tf, err := os.Open(path)\n\tif err != nil {\n\t\treturn \"\", 0, err\n\t}\n\tdefer f.Close()\n\n\th := sha256.New()\n\tsize, err = io.Copy(h, f)\n\tif err != nil {\n\t\treturn \"\", 0, err\n\t}\n\n\tdigest = \"sha256:\" + hex.EncodeToString(h.Sum(nil))\n\treturn digest, size, nil\n}\n"
  },
  {
    "path": "pkg/model/image_builder.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n)\n\n// ImageBuilder builds an ImageArtifact from an ImageSpec.\n// It delegates to a Factory for the docker build, inspects the result\n// to populate labels and the canonical digest, and returns a fully\n// populated ImageArtifact.\ntype ImageBuilder struct {\n\tfactory Factory\n\tdocker  command.Command\n\tsource  *Source\n\topts    BuildOptions\n}\n\n// NewImageBuilder creates an ImageBuilder.\nfunc NewImageBuilder(factory Factory, docker command.Command, source *Source, opts BuildOptions) *ImageBuilder {\n\treturn &ImageBuilder{\n\t\tfactory: factory,\n\t\tdocker:  docker,\n\t\tsource:  source,\n\t\topts:    opts,\n\t}\n}\n\n// Build builds an ImageArtifact from an ImageSpec.\n// It delegates to the Factory for the docker build, inspects the result\n// to populate labels and the canonical digest, and returns a fully\n// populated ImageArtifact.\nfunc (b *ImageBuilder) Build(ctx context.Context, spec ArtifactSpec) (Artifact, error) {\n\tis, ok := spec.(*ImageSpec)\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"image builder: expected *ImageSpec, got %T\", spec)\n\t}\n\n\t// Build the image via the factory (returns partially populated ImageArtifact)\n\timg, err := b.factory.Build(ctx, b.source, b.opts)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"image build failed: %w\", err)\n\t}\n\n\t// Inspect the built image to get labels and canonical digest.\n\t// Prefer digest (ID) for stable lookups, fall back to reference.\n\tinspectRef := img.Digest\n\tif inspectRef == \"\" {\n\t\tinspectRef = img.Reference\n\t}\n\n\tresp, err := b.docker.Inspect(ctx, inspectRef)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"inspect built image: %w\", err)\n\t}\n\n\t// Populate the artifact with inspect results\n\timg.name = is.Name()\n\timg.Labels = resp.Config.Labels\n\timg.Digest = resp.ID\n\timg.Source = ImageSourceBuild\n\n\tdigest, err := v1.NewHash(resp.ID)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"parse image digest %q: %w\", resp.ID, err)\n\t}\n\timg.descriptor = v1.Descriptor{Digest: digest}\n\n\treturn img, nil\n}\n"
  },
  {
    "path": "pkg/model/image_builder_test.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"testing\"\n\n\t\"github.com/docker/docker/api/types/image\"\n\tdockerspec \"github.com/moby/docker-image-spec/specs-go/v1\"\n\tocispec \"github.com/opencontainers/image-spec/specs-go/v1\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc TestImageBuilder_HappyPath(t *testing.T) {\n\t// Setup mock factory that returns a built image\n\tfactory := &mockFactory{\n\t\tbuildFunc: func(_ context.Context, _ *Source, opts BuildOptions) (*ImageArtifact, error) {\n\t\t\treturn &ImageArtifact{\n\t\t\t\tReference: opts.ImageName,\n\t\t\t\tDigest:    \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\",\n\t\t\t\tSource:    ImageSourceBuild,\n\t\t\t}, nil\n\t\t},\n\t}\n\n\t// Setup mock docker that returns inspect results with labels\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(_ context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\t\"org.cogmodel.cog_version\": \"0.15.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tsrc := NewSourceFromConfig(&config.Config{\n\t\tImage: \"my-model:latest\",\n\t}, \"/project\")\n\n\tib := NewImageBuilder(factory, docker, src, BuildOptions{\n\t\tImageName: \"my-model:latest\",\n\t})\n\n\tspec := NewImageSpec(\"model\", \"my-model:latest\")\n\tartifact, err := ib.Build(context.Background(), spec)\n\trequire.NoError(t, err)\n\trequire.NotNil(t, artifact)\n\n\t// Type assertion\n\tia, ok := artifact.(*ImageArtifact)\n\trequire.True(t, ok, \"expected *ImageArtifact, got %T\", artifact)\n\n\t// Check artifact interface\n\trequire.Equal(t, ArtifactTypeImage, ia.Type())\n\trequire.Equal(t, \"model\", ia.Name())\n\n\t// Check descriptor has the digest\n\tdesc := ia.Descriptor()\n\trequire.Equal(t, \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\", desc.Digest.String())\n\n\t// Check image-specific fields\n\trequire.Equal(t, \"my-model:latest\", ia.Reference)\n}\n\nfunc TestImageBuilder_ErrorWrongSpecType(t *testing.T) {\n\tsrc := NewSourceFromConfig(&config.Config{}, \"/project\")\n\tib := NewImageBuilder(&mockFactory{}, &mockDocker{}, src, BuildOptions{})\n\n\t// Pass a WeightSpec instead of ImageSpec\n\tweightSpec := NewWeightSpec(\"model\", \"model.bin\", \"/weights/model.bin\")\n\t_, err := ib.Build(context.Background(), weightSpec)\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"expected *ImageSpec\")\n}\n\nfunc TestImageBuilder_ErrorFactoryBuildFails(t *testing.T) {\n\tfactory := &mockFactory{\n\t\tbuildFunc: func(_ context.Context, _ *Source, _ BuildOptions) (*ImageArtifact, error) {\n\t\t\treturn nil, errors.New(\"docker build failed: out of disk\")\n\t\t},\n\t}\n\n\tsrc := NewSourceFromConfig(&config.Config{}, \"/project\")\n\tib := NewImageBuilder(factory, &mockDocker{}, src, BuildOptions{})\n\n\tspec := NewImageSpec(\"model\", \"test-image\")\n\t_, err := ib.Build(context.Background(), spec)\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"image build failed\")\n\trequire.Contains(t, err.Error(), \"out of disk\")\n}\n\nfunc TestImageBuilder_ErrorInspectFails(t *testing.T) {\n\tfactory := &mockFactory{\n\t\tbuildFunc: func(_ context.Context, _ *Source, opts BuildOptions) (*ImageArtifact, error) {\n\t\t\treturn &ImageArtifact{\n\t\t\t\tReference: opts.ImageName,\n\t\t\t\tDigest:    \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\",\n\t\t\t\tSource:    ImageSourceBuild,\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(_ context.Context, _ string) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"image not found\")\n\t\t},\n\t}\n\n\tsrc := NewSourceFromConfig(&config.Config{}, \"/project\")\n\tib := NewImageBuilder(factory, docker, src, BuildOptions{})\n\n\tspec := NewImageSpec(\"model\", \"test-image\")\n\t_, err := ib.Build(context.Background(), spec)\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"inspect built image\")\n}\n\nfunc TestImageBuilder_ImplementsBuilderInterface(t *testing.T) {\n\tsrc := NewSourceFromConfig(&config.Config{}, \"/project\")\n\t// Compile-time check\n\tvar _ Builder = NewImageBuilder(&mockFactory{}, &mockDocker{}, src, BuildOptions{})\n}\n"
  },
  {
    "path": "pkg/model/image_pusher.go",
    "content": "package model\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"net/http\"\n\t\"os\"\n\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote/transport\"\n\t\"github.com/google/go-containerregistry/pkg/v1/tarball\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n\t\"golang.org/x/sync/errgroup\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/registry\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// ImagePusher pushes container images to a registry.\n//\n// It first attempts an OCI chunked push (export from Docker -> tarball ->\n// push layers via registry client), then falls back to Docker's native push\n// on any non-fatal error. This bypasses size limits on Docker's monolithic\n// push path while maintaining backwards compatibility.\ntype ImagePusher struct {\n\tdocker   command.Command\n\tregistry registry.Client\n}\n\n// newImagePusher creates a new ImagePusher.\nfunc newImagePusher(docker command.Command, reg registry.Client) *ImagePusher {\n\treturn &ImagePusher{\n\t\tdocker:   docker,\n\t\tregistry: reg,\n\t}\n}\n\n// imagePushOptions holds the resolved configuration for an image push.\ntype imagePushOptions struct {\n\tprogressFn func(PushProgress)\n\tonFallback func()\n}\n\n// ImagePushOption is a functional option for configuring ImagePusher.Push.\ntype ImagePushOption func(*imagePushOptions)\n\n// WithProgressFn sets a callback for reporting per-layer upload progress.\nfunc WithProgressFn(fn func(PushProgress)) ImagePushOption {\n\treturn func(o *imagePushOptions) {\n\t\to.progressFn = fn\n\t}\n}\n\n// WithOnFallback sets a callback invoked when OCI push fails and the push is\n// about to fall back to Docker push. This allows the caller to clean up any\n// OCI-specific progress display before Docker push starts its own output.\nfunc WithOnFallback(fn func()) ImagePushOption {\n\treturn func(o *imagePushOptions) {\n\t\to.onFallback = fn\n\t}\n}\n\n// Push pushes a container image to the registry.\n//\n// Tries the OCI chunked push path first (if enabled and registry client is\n// available), then falls back to Docker push on any non-fatal error.\n// The artifact must have a valid Reference.\nfunc (p *ImagePusher) Push(ctx context.Context, artifact *ImageArtifact, opts ...ImagePushOption) error {\n\tif artifact == nil {\n\t\treturn fmt.Errorf(\"image artifact is nil\")\n\t}\n\tif artifact.Reference == \"\" {\n\t\treturn fmt.Errorf(\"image artifact has no reference\")\n\t}\n\n\tvar opt imagePushOptions\n\tfor _, apply := range opts {\n\t\tapply(&opt)\n\t}\n\n\timageRef := artifact.Reference\n\n\tif p.canOCIPush() {\n\t\terr := p.ociPush(ctx, imageRef, opt)\n\t\tif err == nil {\n\t\t\treturn nil\n\t\t}\n\t\tif !shouldFallbackToDocker(err) {\n\t\t\treturn fmt.Errorf(\"OCI chunked push: %w\", err)\n\t\t}\n\t\tif opt.onFallback != nil {\n\t\t\topt.onFallback()\n\t\t}\n\t\tconsole.Warnf(\"OCI chunked push failed, falling back to Docker push: %v\", sanitizeError(err))\n\t}\n\n\treturn p.docker.Push(ctx, imageRef)\n}\n\n// canOCIPush returns true if OCI chunked push is enabled.\nfunc (p *ImagePusher) canOCIPush() bool {\n\treturn os.Getenv(\"COG_PUSH_OCI\") == \"1\"\n}\n\n// ociPush exports the image from Docker daemon as a tar, then pushes all layers,\n// config, and manifest to the registry using chunked uploads.\nfunc (p *ImagePusher) ociPush(ctx context.Context, imageRef string, opt imagePushOptions) error {\n\tconsole.Debugf(\"Exporting image %s from Docker daemon...\", imageRef)\n\n\tref, err := name.ParseReference(imageRef, name.Insecure)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"parse image reference %q: %w\", imageRef, err)\n\t}\n\n\t// Get the Docker tar stream directly from the docker command\n\trc, err := p.docker.ImageSave(ctx, imageRef)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"export image from daemon: %w\", err)\n\t}\n\tdefer rc.Close() //nolint:errcheck\n\n\t// Write the tar to a temp file so we can seek on it\n\ttmpTar, err := os.CreateTemp(\"\", \"cog-image-*.tar\")\n\tif err != nil {\n\t\treturn fmt.Errorf(\"create temp tar file: %w\", err)\n\t}\n\tdefer func() { _ = os.Remove(tmpTar.Name()) }() //nolint:gosec // G703: path from os.CreateTemp, not user input\n\tdefer tmpTar.Close()                            //nolint:errcheck\n\n\tif _, err := io.Copy(tmpTar, rc); err != nil {\n\t\treturn fmt.Errorf(\"write image tar: %w\", err)\n\t}\n\t_ = rc.Close()\n\n\t// Load image from Docker tar using go-containerregistry.\n\t// tarball.ImageFromPath returns a lazy image that reads layers on-demand\n\t// from the tar file rather than loading them all into memory at once.\n\ttag, ok := ref.(name.Tag)\n\tif !ok {\n\t\t// If reference is a digest, use tag \"latest\" as a fallback\n\t\ttag = ref.Context().Tag(\"latest\")\n\t}\n\n\timg, err := tarball.ImageFromPath(tmpTar.Name(), &tag)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"load image from tar: %w\", err)\n\t}\n\n\treturn p.pushImage(ctx, imageRef, img, opt)\n}\n\n// pushImage pushes a v1.Image (layers, config, manifest) to the registry.\nfunc (p *ImagePusher) pushImage(ctx context.Context, imageRef string, img v1.Image, opt imagePushOptions) error {\n\trepo := repoFromReference(imageRef)\n\n\tif err := p.pushLayers(ctx, repo, img, opt); err != nil {\n\t\treturn fmt.Errorf(\"push layers: %w\", err)\n\t}\n\n\tif err := p.pushConfig(ctx, repo, img); err != nil {\n\t\treturn fmt.Errorf(\"push config: %w\", err)\n\t}\n\n\tconsole.Debugf(\"Pushing image manifest for %s\", imageRef)\n\tif err := p.registry.PushImage(ctx, imageRef, img); err != nil {\n\t\treturn fmt.Errorf(\"push manifest: %w\", err)\n\t}\n\n\treturn nil\n}\n\n// pushLayers pushes all image layers concurrently using the registry client's\n// WriteLayer method, which handles chunked uploads, retry, and progress reporting.\nfunc (p *ImagePusher) pushLayers(ctx context.Context, repo string, img v1.Image, opt imagePushOptions) error {\n\tlayers, err := img.Layers()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"get image layers: %w\", err)\n\t}\n\n\tif len(layers) == 0 {\n\t\treturn nil\n\t}\n\n\tconcurrency := GetPushConcurrency()\n\tconsole.Debugf(\"Pushing %d layers with concurrency %d\", len(layers), concurrency)\n\n\tg, ctx := errgroup.WithContext(ctx)\n\tg.SetLimit(concurrency)\n\n\tfor _, layer := range layers {\n\t\tg.Go(func() error {\n\t\t\treturn p.pushLayer(ctx, repo, layer, opt)\n\t\t})\n\t}\n\n\treturn g.Wait()\n}\n\n// pushLayer pushes a single layer with progress reporting.\nfunc (p *ImagePusher) pushLayer(ctx context.Context, repo string, layer v1.Layer, opt imagePushOptions) error {\n\tdigest, err := layer.Digest()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"get layer digest: %w\", err)\n\t}\n\n\tsize, err := layer.Size()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"get layer size: %w\", err)\n\t}\n\n\tconsole.Debugf(\"Pushing layer %s (%d bytes)\", digest, size)\n\n\tvar onProgress func(v1.Update)\n\tif opt.progressFn != nil {\n\t\tdigestStr := digest.String()\n\t\tonProgress = func(update v1.Update) {\n\t\t\topt.progressFn(PushProgress{\n\t\t\t\tLayerDigest: digestStr,\n\t\t\t\tComplete:    update.Complete,\n\t\t\t\tTotal:       update.Total,\n\t\t\t})\n\t\t}\n\t}\n\n\twriteErr := writeLayerWithProgress(ctx, p.registry, registry.WriteLayerOptions{\n\t\tRepo:  repo,\n\t\tLayer: layer,\n\t}, onProgress)\n\n\tif writeErr != nil {\n\t\treturn fmt.Errorf(\"push layer %s: %w\", digest, writeErr)\n\t}\n\n\treturn nil\n}\n\n// pushConfig pushes the image config blob to the registry.\nfunc (p *ImagePusher) pushConfig(ctx context.Context, repo string, img v1.Image) error {\n\tcfgBlob, err := img.RawConfigFile()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"get config: %w\", err)\n\t}\n\n\tcfgName, err := img.ConfigName()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"get config digest: %w\", err)\n\t}\n\n\tconsole.Debugf(\"Pushing config blob %s (%d bytes)\", cfgName, len(cfgBlob))\n\n\tconfigLayer := &configBlobLayer{\n\t\tdata:   cfgBlob,\n\t\tdigest: cfgName,\n\t}\n\n\treturn p.registry.WriteLayer(ctx, registry.WriteLayerOptions{\n\t\tRepo:  repo,\n\t\tLayer: configLayer,\n\t})\n}\n\n// shouldFallbackToDocker returns true if the error is safe to fall back from.\n// We do NOT fall back on context errors (cancellation/timeout) or authentication\n// errors (401/403), since Docker push would fail with the same credentials.\nfunc shouldFallbackToDocker(err error) bool {\n\tif err == nil {\n\t\treturn false\n\t}\n\tif errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {\n\t\treturn false\n\t}\n\tvar transportErr *transport.Error\n\tif errors.As(err, &transportErr) {\n\t\tswitch transportErr.StatusCode {\n\t\tcase http.StatusUnauthorized, http.StatusForbidden:\n\t\t\treturn false\n\t\t}\n\t}\n\treturn true\n}\n\n// sanitizeError returns a clean, user-friendly error message.\n//\n// Registry errors from go-containerregistry's transport.Error can contain the\n// entire HTTP response body which produces unreadable terminal output. This function extracts\n// just the HTTP status code and status text for those cases.\nfunc sanitizeError(err error) error {\n\tvar transportErr *transport.Error\n\tif errors.As(err, &transportErr) {\n\t\treturn fmt.Errorf(\"HTTP %d %s\", transportErr.StatusCode, http.StatusText(transportErr.StatusCode))\n\t}\n\treturn err\n}\n\n// configBlobLayer wraps a config blob to satisfy the v1.Layer interface\n// required by WriteLayerOptions.\ntype configBlobLayer struct {\n\tdata   []byte\n\tdigest v1.Hash\n}\n\nfunc (c *configBlobLayer) Digest() (v1.Hash, error) {\n\treturn c.digest, nil\n}\n\n// DiffID returns the same hash as Digest. For uncompressed config blobs,\n// the compressed and uncompressed representations are identical, so DiffID\n// (hash of uncompressed content) equals Digest (hash of compressed content).\nfunc (c *configBlobLayer) DiffID() (v1.Hash, error) {\n\treturn c.digest, nil\n}\n\nfunc (c *configBlobLayer) Compressed() (io.ReadCloser, error) {\n\treturn io.NopCloser(bytes.NewReader(c.data)), nil\n}\n\nfunc (c *configBlobLayer) Uncompressed() (io.ReadCloser, error) {\n\treturn io.NopCloser(bytes.NewReader(c.data)), nil\n}\n\nfunc (c *configBlobLayer) Size() (int64, error) {\n\treturn int64(len(c.data)), nil\n}\n\nfunc (c *configBlobLayer) MediaType() (types.MediaType, error) {\n\treturn types.OCIConfigJSON, nil\n}\n"
  },
  {
    "path": "pkg/model/image_pusher_test.go",
    "content": "package model\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"sync\"\n\t\"testing\"\n\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/empty\"\n\t\"github.com/google/go-containerregistry/pkg/v1/mutate\"\n\t\"github.com/google/go-containerregistry/pkg/v1/random\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote/transport\"\n\t\"github.com/google/go-containerregistry/pkg/v1/tarball\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// ociMockClient implements registry.Client for testing ImagePusher.\ntype ociMockClient struct {\n\tmu              sync.Mutex\n\twrittenLayers   []v1.Hash\n\tpushedImages    []string\n\twriteLayerErr   error\n\tpushImageErr    error\n\twriteLayerCount int\n}\n\nfunc (m *ociMockClient) WriteLayer(_ context.Context, opts registry.WriteLayerOptions) error {\n\tm.mu.Lock()\n\tdefer m.mu.Unlock()\n\tm.writeLayerCount++\n\tif m.writeLayerErr != nil {\n\t\treturn m.writeLayerErr\n\t}\n\tdigest, err := opts.Layer.Digest()\n\tif err != nil {\n\t\treturn err\n\t}\n\tm.writtenLayers = append(m.writtenLayers, digest)\n\n\t// Send progress if channel is provided\n\tif opts.ProgressCh != nil {\n\t\tsize, _ := opts.Layer.Size()\n\t\topts.ProgressCh <- v1.Update{Complete: size, Total: size}\n\t}\n\treturn nil\n}\n\nfunc (m *ociMockClient) PushImage(_ context.Context, ref string, _ v1.Image) error {\n\tm.mu.Lock()\n\tdefer m.mu.Unlock()\n\tif m.pushImageErr != nil {\n\t\treturn m.pushImageErr\n\t}\n\tm.pushedImages = append(m.pushedImages, ref)\n\treturn nil\n}\n\nfunc (m *ociMockClient) Inspect(context.Context, string, *registry.Platform) (*registry.ManifestResult, error) {\n\treturn nil, nil\n}\nfunc (m *ociMockClient) GetImage(context.Context, string, *registry.Platform) (v1.Image, error) {\n\treturn nil, nil\n}\nfunc (m *ociMockClient) Exists(context.Context, string) (bool, error) { return false, nil }\nfunc (m *ociMockClient) GetDescriptor(context.Context, string) (v1.Descriptor, error) {\n\treturn v1.Descriptor{}, nil\n}\nfunc (m *ociMockClient) PushIndex(context.Context, string, v1.ImageIndex) error { return nil }\n\n// testArtifact creates an *ImageArtifact for testing with the given reference string.\nfunc testArtifact(ref string) *ImageArtifact {\n\treturn &ImageArtifact{Reference: ref}\n}\n\n// fakeImageSaveFunc creates a fake ImageSave function that produces a Docker-format tar\n// from the given v1.Image. This simulates Docker's ImageSave API.\nfunc fakeImageSaveFunc(img v1.Image, tagStr string) func(context.Context, string) (io.ReadCloser, error) {\n\treturn func(_ context.Context, _ string) (io.ReadCloser, error) {\n\t\ttag, err := name.NewTag(tagStr, name.Insecure)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"parse tag: %w\", err)\n\t\t}\n\t\tvar buf bytes.Buffer\n\t\trefToImage := map[name.Tag]v1.Image{tag: img}\n\t\tif err := tarball.MultiWrite(refToImage, &buf); err != nil {\n\t\t\treturn nil, fmt.Errorf(\"create test tar: %w\", err)\n\t\t}\n\t\treturn io.NopCloser(bytes.NewReader(buf.Bytes())), nil\n\t}\n}\n\n// =============================================================================\n// ImagePusher.Push — OCI chunked push tests\n// =============================================================================\n\nfunc TestImagePusher_Push(t *testing.T) {\n\tt.Setenv(\"COG_PUSH_OCI\", \"1\")\n\n\tt.Run(\"pushes layers config and manifest via OCI path\", func(t *testing.T) {\n\t\timg, err := random.Image(1024, 2) // 2 layers of 1KB\n\t\trequire.NoError(t, err)\n\n\t\tmock := &ociMockClient{}\n\t\ttag := \"example.com/test/repo:v1\"\n\t\tdocker := &mockDocker{imageSaveFunc: fakeImageSaveFunc(img, tag)}\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag))\n\t\trequire.NoError(t, err)\n\n\t\t// Should have pushed 2 layers + 1 config blob = 3 WriteLayer calls\n\t\tassert.Equal(t, 3, mock.writeLayerCount)\n\n\t\t// Should have pushed the manifest\n\t\trequire.Len(t, mock.pushedImages, 1)\n\t\tassert.Equal(t, tag, mock.pushedImages[0])\n\t})\n\n\tt.Run(\"reports progress via callback\", func(t *testing.T) {\n\t\timg, err := random.Image(1024, 1)\n\t\trequire.NoError(t, err)\n\n\t\tmock := &ociMockClient{}\n\t\ttag := \"example.com/test/repo:v1\"\n\t\tdocker := &mockDocker{imageSaveFunc: fakeImageSaveFunc(img, tag)}\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\tvar mu sync.Mutex\n\t\tvar progressUpdates []PushProgress\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag), WithProgressFn(func(p PushProgress) {\n\t\t\tmu.Lock()\n\t\t\tdefer mu.Unlock()\n\t\t\tprogressUpdates = append(progressUpdates, p)\n\t\t}))\n\t\trequire.NoError(t, err)\n\n\t\tmu.Lock()\n\t\tdefer mu.Unlock()\n\t\tassert.NotEmpty(t, progressUpdates)\n\t\tfor _, p := range progressUpdates {\n\t\t\tassert.NotEmpty(t, p.LayerDigest)\n\t\t\tassert.True(t, p.Complete > 0)\n\t\t\tassert.True(t, p.Total > 0)\n\t\t}\n\t})\n\n\tt.Run(\"falls back to docker when WriteLayer fails\", func(t *testing.T) {\n\t\timg, err := random.Image(1024, 1)\n\t\trequire.NoError(t, err)\n\n\t\tvar dockerPushed bool\n\t\tmock := &ociMockClient{writeLayerErr: errors.New(\"upload failed\")}\n\t\ttag := \"example.com/test/repo:v1\"\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: fakeImageSaveFunc(img, tag),\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tdockerPushed = true\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag))\n\t\trequire.NoError(t, err)\n\t\tassert.True(t, dockerPushed)\n\t})\n\n\tt.Run(\"falls back to docker when PushImage fails\", func(t *testing.T) {\n\t\timg, err := random.Image(1024, 1)\n\t\trequire.NoError(t, err)\n\n\t\tvar dockerPushed bool\n\t\tmock := &ociMockClient{pushImageErr: errors.New(\"manifest push failed\")}\n\t\ttag := \"example.com/test/repo:v1\"\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: fakeImageSaveFunc(img, tag),\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tdockerPushed = true\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag))\n\t\trequire.NoError(t, err)\n\t\tassert.True(t, dockerPushed)\n\t})\n\n\tt.Run(\"falls back to docker when ImageSave fails\", func(t *testing.T) {\n\t\tmock := &ociMockClient{}\n\n\t\tvar dockerPushed bool\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: func(_ context.Context, _ string) (io.ReadCloser, error) {\n\t\t\t\treturn nil, errors.New(\"docker daemon unavailable\")\n\t\t\t},\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tdockerPushed = true\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr := pusher.Push(context.Background(), testArtifact(\"example.com/test/repo:v1\"))\n\t\trequire.NoError(t, err)\n\t\tassert.True(t, dockerPushed)\n\t})\n\n\tt.Run(\"handles empty image with no layers\", func(t *testing.T) {\n\t\timg := empty.Image\n\t\timg, err := mutate.Config(img, v1.Config{})\n\t\trequire.NoError(t, err)\n\n\t\tmock := &ociMockClient{}\n\t\ttag := \"example.com/test/repo:empty\"\n\t\tdocker := &mockDocker{imageSaveFunc: fakeImageSaveFunc(img, tag)}\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag))\n\t\trequire.NoError(t, err)\n\n\t\t// Only config blob should be written (no layers)\n\t\tassert.Equal(t, 1, mock.writeLayerCount)\n\t\trequire.Len(t, mock.pushedImages, 1)\n\t})\n}\n\n// =============================================================================\n// ImagePusher.Push with artifact tests\n// =============================================================================\n\nfunc TestImagePusher_PushArtifact(t *testing.T) {\n\tt.Run(\"pushes artifact by reference\", func(t *testing.T) {\n\t\tvar dockerPushed string\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(_ context.Context, ref string) error {\n\t\t\t\tdockerPushed = ref\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\t// No registry — will use Docker push directly\n\t\tpusher := newImagePusher(docker, nil)\n\t\tartifact := &ImageArtifact{Reference: \"r8.im/user/model:latest\"}\n\n\t\terr := pusher.Push(context.Background(), artifact)\n\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"r8.im/user/model:latest\", dockerPushed)\n\t})\n\n\tt.Run(\"returns error for nil artifact\", func(t *testing.T) {\n\t\tpusher := newImagePusher(&mockDocker{}, nil)\n\n\t\terr := pusher.Push(context.Background(), nil)\n\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"nil\")\n\t})\n\n\tt.Run(\"returns error for empty reference\", func(t *testing.T) {\n\t\tpusher := newImagePusher(&mockDocker{}, nil)\n\n\t\terr := pusher.Push(context.Background(), &ImageArtifact{Reference: \"\"})\n\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"no reference\")\n\t})\n\n\tt.Run(\"propagates docker push error\", func(t *testing.T) {\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\treturn errors.New(\"unauthorized: authentication required\")\n\t\t\t},\n\t\t}\n\n\t\tpusher := newImagePusher(docker, nil)\n\t\tartifact := &ImageArtifact{Reference: \"r8.im/user/model:latest\"}\n\n\t\terr := pusher.Push(context.Background(), artifact)\n\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"unauthorized\")\n\t})\n}\n\n// =============================================================================\n// Docker fallback behavior tests\n// =============================================================================\n\nfunc TestImagePusher_Fallback(t *testing.T) {\n\tt.Setenv(\"COG_PUSH_OCI\", \"1\")\n\n\tt.Run(\"uses OCI push when it succeeds\", func(t *testing.T) {\n\t\timg, err := random.Image(512, 1)\n\t\trequire.NoError(t, err)\n\n\t\tmock := &ociMockClient{}\n\t\ttag := \"example.com/test/repo:v1\"\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: fakeImageSaveFunc(img, tag),\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tt.Fatal(\"docker push should not be called when OCI succeeds\")\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag))\n\t\trequire.NoError(t, err)\n\t})\n\n\tt.Run(\"falls back to docker on OCI error\", func(t *testing.T) {\n\t\tvar dockerPushed bool\n\t\tmock := &ociMockClient{writeLayerErr: errors.New(\"connection reset\")}\n\t\ttag := \"example.com/test/repo:v1\"\n\n\t\timg, err := random.Image(512, 1)\n\t\trequire.NoError(t, err)\n\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: fakeImageSaveFunc(img, tag),\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tdockerPushed = true\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag))\n\t\trequire.NoError(t, err)\n\t\tassert.True(t, dockerPushed)\n\t})\n\n\tt.Run(\"does not fall back on context cancellation\", func(t *testing.T) {\n\t\tctx, cancel := context.WithCancel(context.Background())\n\t\tcancel() // cancel immediately\n\n\t\tmock := &ociMockClient{}\n\t\ttag := \"example.com/test/repo:v1\"\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: func(ctx context.Context, _ string) (io.ReadCloser, error) {\n\t\t\t\treturn nil, ctx.Err()\n\t\t\t},\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tt.Fatal(\"docker push should not be called on context cancellation\")\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr := pusher.Push(ctx, testArtifact(tag))\n\t\trequire.Error(t, err)\n\t})\n\n\tt.Run(\"does not fall back on 401 unauthorized\", func(t *testing.T) {\n\t\tmock := &ociMockClient{writeLayerErr: &transport.Error{StatusCode: 401}}\n\t\ttag := \"example.com/test/repo:v1\"\n\n\t\timg, err := random.Image(512, 1)\n\t\trequire.NoError(t, err)\n\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: fakeImageSaveFunc(img, tag),\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tt.Fatal(\"docker push should not be called on auth error\")\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag))\n\t\trequire.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"OCI chunked push\")\n\t})\n\n\tt.Run(\"does not fall back on 403 forbidden\", func(t *testing.T) {\n\t\tmock := &ociMockClient{writeLayerErr: &transport.Error{StatusCode: 403}}\n\t\ttag := \"example.com/test/repo:v1\"\n\n\t\timg, err := random.Image(512, 1)\n\t\trequire.NoError(t, err)\n\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: fakeImageSaveFunc(img, tag),\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tt.Fatal(\"docker push should not be called on auth error\")\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag))\n\t\trequire.Error(t, err)\n\t\tassert.Contains(t, err.Error(), \"OCI chunked push\")\n\t})\n\n\tt.Run(\"uses docker directly when registry is nil\", func(t *testing.T) {\n\t\tvar dockerPushed bool\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tdockerPushed = true\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := newImagePusher(docker, nil)\n\n\t\terr := pusher.Push(context.Background(), testArtifact(\"example.com/test/repo:v1\"))\n\t\trequire.NoError(t, err)\n\t\tassert.True(t, dockerPushed)\n\t})\n}\n\n// =============================================================================\n// shouldFallbackToDocker tests\n// =============================================================================\n\nfunc TestShouldFallbackToDocker(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\terr      error\n\t\texpected bool\n\t}{\n\t\t{\"nil error\", nil, false},\n\t\t{\"context canceled\", context.Canceled, false},\n\t\t{\"context deadline\", context.DeadlineExceeded, false},\n\t\t{\"401 unauthorized\", &transport.Error{StatusCode: 401}, false},\n\t\t{\"403 forbidden\", &transport.Error{StatusCode: 403}, false},\n\t\t{\"wrapped 401\", fmt.Errorf(\"push failed: %w\", &transport.Error{StatusCode: 401}), false},\n\t\t{\"500 server error\", &transport.Error{StatusCode: 500}, true},\n\t\t{\"network error\", errors.New(\"connection refused\"), true},\n\t\t{\"unknown error\", errors.New(\"something unexpected\"), true},\n\t\t{\"export error\", errors.New(\"export OCI layout: daemon error\"), true},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tassert.Equal(t, tt.expected, shouldFallbackToDocker(tt.err))\n\t\t})\n\t}\n}\n\n// =============================================================================\n// sanitizeError tests\n// =============================================================================\n\nfunc TestSanitizeError(t *testing.T) {\n\tt.Run(\"strips HTML body from transport error\", func(t *testing.T) {\n\t\thtmlBody := `413 Request Entity Too Large
413 Request Entity Too Large
cloudflare
`\n\t\ttransportErr := &transport.Error{\n\t\t\tStatusCode: 413,\n\t\t\tErrors:     nil,\n\t\t}\n\t\t// The rawBody field is unexported, so we test via the wrapped error behavior.\n\t\t// A transport.Error with no Errors slice and status 413 produces a message\n\t\t// that includes the raw body — sanitizeError should replace it.\n\t\t_ = htmlBody // illustrates the problem scenario\n\n\t\tresult := sanitizeError(transportErr)\n\t\tassert.Equal(t, \"HTTP 413 Request Entity Too Large\", result.Error())\n\t})\n\n\tt.Run(\"strips body from 502 transport error\", func(t *testing.T) {\n\t\ttransportErr := &transport.Error{\n\t\t\tStatusCode: 502,\n\t\t}\n\n\t\tresult := sanitizeError(transportErr)\n\t\tassert.Equal(t, \"HTTP 502 Bad Gateway\", result.Error())\n\t})\n\n\tt.Run(\"passes through non-transport errors unchanged\", func(t *testing.T) {\n\t\terr := errors.New(\"connection refused\")\n\t\tresult := sanitizeError(err)\n\t\tassert.Equal(t, \"connection refused\", result.Error())\n\t})\n\n\tt.Run(\"passes through wrapped transport errors\", func(t *testing.T) {\n\t\ttransportErr := &transport.Error{\n\t\t\tStatusCode: 413,\n\t\t}\n\t\twrapped := fmt.Errorf(\"pushing layer: %w\", transportErr)\n\n\t\tresult := sanitizeError(wrapped)\n\t\tassert.Equal(t, \"HTTP 413 Request Entity Too Large\", result.Error())\n\t})\n}\n\n// =============================================================================\n// OnFallback callback tests\n// =============================================================================\n\nfunc TestImagePusher_OnFallback(t *testing.T) {\n\tt.Setenv(\"COG_PUSH_OCI\", \"1\")\n\n\tt.Run(\"calls OnFallback before docker push on OCI failure\", func(t *testing.T) {\n\t\tvar callOrder []string\n\n\t\tmock := &ociMockClient{writeLayerErr: errors.New(\"connection reset\")}\n\t\ttag := \"example.com/test/repo:v1\"\n\n\t\timg, err := random.Image(512, 1)\n\t\trequire.NoError(t, err)\n\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: fakeImageSaveFunc(img, tag),\n\t\t\tpushFunc: func(_ context.Context, _ string) error {\n\t\t\t\tcallOrder = append(callOrder, \"docker-push\")\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\terr = pusher.Push(context.Background(), testArtifact(tag), WithOnFallback(func() {\n\t\t\tcallOrder = append(callOrder, \"on-fallback\")\n\t\t}))\n\t\trequire.NoError(t, err)\n\t\tassert.Equal(t, []string{\"on-fallback\", \"docker-push\"}, callOrder)\n\t})\n\n\tt.Run(\"does not call OnFallback when OCI push succeeds\", func(t *testing.T) {\n\t\tmock := &ociMockClient{}\n\t\ttag := \"example.com/test/repo:v1\"\n\n\t\timg, err := random.Image(512, 1)\n\t\trequire.NoError(t, err)\n\n\t\tdocker := &mockDocker{\n\t\t\timageSaveFunc: fakeImageSaveFunc(img, tag),\n\t\t}\n\n\t\tpusher := newImagePusher(docker, mock)\n\n\t\tvar fallbackCalled bool\n\t\terr = pusher.Push(context.Background(), testArtifact(tag), WithOnFallback(func() {\n\t\t\tfallbackCalled = true\n\t\t}))\n\t\trequire.NoError(t, err)\n\t\tassert.False(t, fallbackCalled)\n\t})\n}\n\n// =============================================================================\n// configBlobLayer tests\n// =============================================================================\n\nfunc TestConfigBlobLayer(t *testing.T) {\n\tdata := []byte(`{\"architecture\":\"amd64\",\"os\":\"linux\"}`)\n\tdigest := v1.Hash{Algorithm: \"sha256\", Hex: \"abc123\"}\n\n\tlayer := &configBlobLayer{data: data, digest: digest}\n\n\tt.Run(\"Digest\", func(t *testing.T) {\n\t\td, err := layer.Digest()\n\t\trequire.NoError(t, err)\n\t\tassert.Equal(t, digest, d)\n\t})\n\n\tt.Run(\"DiffID equals Digest for uncompressed config\", func(t *testing.T) {\n\t\td, err := layer.DiffID()\n\t\trequire.NoError(t, err)\n\t\tassert.Equal(t, digest, d)\n\t})\n\n\tt.Run(\"Size\", func(t *testing.T) {\n\t\tsize, err := layer.Size()\n\t\trequire.NoError(t, err)\n\t\tassert.Equal(t, int64(len(data)), size)\n\t})\n\n\tt.Run(\"MediaType\", func(t *testing.T) {\n\t\tmt, err := layer.MediaType()\n\t\trequire.NoError(t, err)\n\t\tassert.Equal(t, types.OCIConfigJSON, mt)\n\t})\n\n\tt.Run(\"Compressed returns data\", func(t *testing.T) {\n\t\trc, err := layer.Compressed()\n\t\trequire.NoError(t, err)\n\t\tdefer rc.Close()\n\t\tgot, err := io.ReadAll(rc)\n\t\trequire.NoError(t, err)\n\t\tassert.Equal(t, data, got)\n\t})\n\n\tt.Run(\"Uncompressed returns data\", func(t *testing.T) {\n\t\trc, err := layer.Uncompressed()\n\t\trequire.NoError(t, err)\n\t\tdefer rc.Close()\n\t\tgot, err := io.ReadAll(rc)\n\t\trequire.NoError(t, err)\n\t\tassert.Equal(t, data, got)\n\t})\n}\n\n// =============================================================================\n// GetPushConcurrency tests\n// =============================================================================\n\nfunc TestGetPushConcurrency(t *testing.T) {\n\tt.Run(\"returns default when env not set\", func(t *testing.T) {\n\t\tt.Setenv(\"COG_PUSH_CONCURRENCY\", \"\")\n\t\tassert.Equal(t, DefaultPushConcurrency, GetPushConcurrency())\n\t})\n\n\tt.Run(\"returns env var value\", func(t *testing.T) {\n\t\tt.Setenv(\"COG_PUSH_CONCURRENCY\", \"8\")\n\t\tassert.Equal(t, 8, GetPushConcurrency())\n\t})\n\n\tt.Run(\"returns default for invalid value\", func(t *testing.T) {\n\t\tt.Setenv(\"COG_PUSH_CONCURRENCY\", \"not-a-number\")\n\t\tassert.Equal(t, DefaultPushConcurrency, GetPushConcurrency())\n\t})\n\n\tt.Run(\"returns default for zero\", func(t *testing.T) {\n\t\tt.Setenv(\"COG_PUSH_CONCURRENCY\", \"0\")\n\t\tassert.Equal(t, DefaultPushConcurrency, GetPushConcurrency())\n\t})\n\n\tt.Run(\"returns default for negative\", func(t *testing.T) {\n\t\tt.Setenv(\"COG_PUSH_CONCURRENCY\", \"-1\")\n\t\tassert.Equal(t, DefaultPushConcurrency, GetPushConcurrency())\n\t})\n}\n"
  },
  {
    "path": "pkg/model/image_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc TestImage_IsCogModel(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\timage  *ImageArtifact\n\t\texpect bool\n\t}{\n\t\t{\n\t\t\tname:   \"nil labels\",\n\t\t\timage:  &ImageArtifact{Labels: nil},\n\t\t\texpect: false,\n\t\t},\n\t\t{\n\t\t\tname:   \"empty labels\",\n\t\t\timage:  &ImageArtifact{Labels: map[string]string{}},\n\t\t\texpect: false,\n\t\t},\n\t\t{\n\t\t\tname: \"has cog config label\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig: `{\"build\": {}}`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpect: true,\n\t\t},\n\t\t{\n\t\t\tname: \"has other labels but not cog config\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\t\"some.other.label\": \"value\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpect: false,\n\t\t},\n\t\t{\n\t\t\tname: \"has cog version but not config\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpect: false,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.image.IsCogModel()\n\t\t\trequire.Equal(t, tt.expect, result)\n\t\t})\n\t}\n}\n\nfunc TestImage_CogVersion(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\timage  *ImageArtifact\n\t\texpect string\n\t}{\n\t\t{\n\t\t\tname:   \"nil labels\",\n\t\t\timage:  &ImageArtifact{Labels: nil},\n\t\t\texpect: \"\",\n\t\t},\n\t\t{\n\t\t\tname:   \"empty labels\",\n\t\t\timage:  &ImageArtifact{Labels: map[string]string{}},\n\t\t\texpect: \"\",\n\t\t},\n\t\t{\n\t\t\tname: \"has version label\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpect: \"0.10.0\",\n\t\t},\n\t\t{\n\t\t\tname: \"has other labels but not version\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig: `{\"build\": {}}`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpect: \"\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.image.CogVersion()\n\t\t\trequire.Equal(t, tt.expect, result)\n\t\t})\n\t}\n}\n\nfunc TestImage_Config(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\timage  *ImageArtifact\n\t\texpect string\n\t}{\n\t\t{\n\t\t\tname:   \"nil labels\",\n\t\t\timage:  &ImageArtifact{Labels: nil},\n\t\t\texpect: \"\",\n\t\t},\n\t\t{\n\t\t\tname: \"has config label\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig: `{\"build\": {\"python_version\": \"3.11\"}}`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpect: `{\"build\": {\"python_version\": \"3.11\"}}`,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.image.Config()\n\t\t\trequire.Equal(t, tt.expect, result)\n\t\t})\n\t}\n}\n\nfunc TestImage_OpenAPISchema(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\timage  *ImageArtifact\n\t\texpect string\n\t}{\n\t\t{\n\t\t\tname:   \"nil labels\",\n\t\t\timage:  &ImageArtifact{Labels: nil},\n\t\t\texpect: \"\",\n\t\t},\n\t\t{\n\t\t\tname: \"has openapi schema label\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelOpenAPISchema: `{\"openapi\": \"3.0.0\"}`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpect: `{\"openapi\": \"3.0.0\"}`,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.image.OpenAPISchema()\n\t\t\trequire.Equal(t, tt.expect, result)\n\t\t})\n\t}\n}\n\nfunc TestImageSource_Values(t *testing.T) {\n\t// Verify the constants have expected string values\n\trequire.Equal(t, ImageSource(\"local\"), ImageSourceLocal)\n\trequire.Equal(t, ImageSource(\"remote\"), ImageSourceRemote)\n\trequire.Equal(t, ImageSource(\"build\"), ImageSourceBuild)\n}\n\nfunc TestLabelKeys(t *testing.T) {\n\t// Verify label keys have expected prefixes\n\trequire.Equal(t, \"run.cog.config\", LabelConfig)\n\trequire.Equal(t, \"run.cog.version\", LabelVersion)\n\trequire.Equal(t, \"run.cog.openapi_schema\", LabelOpenAPISchema)\n\trequire.Equal(t, \"run.cog.r8_weights_manifest\", LabelWeightsManifest)\n}\n\n// =============================================================================\n// Parsed accessor tests\n// =============================================================================\n\nfunc TestImage_ParsedConfig(t *testing.T) {\n\ttests := []struct {\n\t\tname        string\n\t\timage       *ImageArtifact\n\t\texpectNil   bool\n\t\texpectErr   bool\n\t\tcheckConfig func(t *testing.T, cfg *config.Config)\n\t}{\n\t\t{\n\t\t\tname:      \"nil labels returns nil without error\",\n\t\t\timage:     &ImageArtifact{Labels: nil},\n\t\t\texpectNil: true,\n\t\t\texpectErr: false,\n\t\t},\n\t\t{\n\t\t\tname:      \"empty labels returns nil without error\",\n\t\t\timage:     &ImageArtifact{Labels: map[string]string{}},\n\t\t\texpectNil: true,\n\t\t\texpectErr: false,\n\t\t},\n\t\t{\n\t\t\tname: \"missing config label returns nil without error\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectNil: true,\n\t\t\texpectErr: false,\n\t\t},\n\t\t{\n\t\t\tname: \"valid config JSON parses correctly\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig: `{\"build\":{\"python_version\":\"3.12\",\"gpu\":true},\"predict\":\"predict.py:Predictor\"}`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectNil: false,\n\t\t\texpectErr: false,\n\t\t\tcheckConfig: func(t *testing.T, cfg *config.Config) {\n\t\t\t\trequire.Equal(t, \"3.12\", cfg.Build.PythonVersion)\n\t\t\t\trequire.True(t, cfg.Build.GPU)\n\t\t\t\trequire.Equal(t, \"predict.py:Predictor\", cfg.Predict)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"invalid JSON returns error\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig: `{invalid json`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectNil: true,\n\t\t\texpectErr: true,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tcfg, err := tt.image.ParsedConfig()\n\n\t\t\tif tt.expectErr {\n\t\t\t\trequire.Error(t, err)\n\t\t\t} else {\n\t\t\t\trequire.NoError(t, err)\n\t\t\t}\n\n\t\t\tif tt.expectNil {\n\t\t\t\trequire.Nil(t, cfg)\n\t\t\t} else {\n\t\t\t\trequire.NotNil(t, cfg)\n\t\t\t\tif tt.checkConfig != nil {\n\t\t\t\t\ttt.checkConfig(t, cfg)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestImage_ToModel(t *testing.T) {\n\ttests := []struct {\n\t\tname       string\n\t\timage      *ImageArtifact\n\t\texpectErr  error\n\t\tcheckModel func(t *testing.T, m *Model)\n\t}{\n\t\t{\n\t\t\tname:      \"not a cog model returns ErrNotCogModel\",\n\t\t\timage:     &ImageArtifact{Labels: map[string]string{}},\n\t\t\texpectErr: ErrNotCogModel,\n\t\t},\n\t\t{\n\t\t\tname: \"valid cog model with config and schema\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tReference: \"my-image:latest\",\n\t\t\t\tDigest:    \"sha256:abc123\",\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig:        `{\"build\":{\"python_version\":\"3.12\"},\"predict\":\"predict.py:Predictor\"}`,\n\t\t\t\t\tLabelVersion:       \"0.10.0\",\n\t\t\t\t\tLabelOpenAPISchema: `{\"openapi\":\"3.0.2\",\"info\":{\"title\":\"Cog\",\"version\":\"0.1.0\"},\"paths\":{}}`,\n\t\t\t\t},\n\t\t\t\tSource: ImageSourceLocal,\n\t\t\t},\n\t\t\tcheckModel: func(t *testing.T, m *Model) {\n\t\t\t\trequire.NotNil(t, m.Image)\n\t\t\t\trequire.Equal(t, \"my-image:latest\", m.Image.Reference)\n\t\t\t\trequire.Equal(t, \"0.10.0\", m.CogVersion)\n\t\t\t\trequire.NotNil(t, m.Config)\n\t\t\t\trequire.Equal(t, \"3.12\", m.Config.Build.PythonVersion)\n\t\t\t\trequire.NotNil(t, m.Schema)\n\t\t\t\trequire.Equal(t, \"Cog\", m.Schema.Info.Title)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"valid cog model without schema\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig:  `{\"build\":{}}`,\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t},\n\t\t\tcheckModel: func(t *testing.T, m *Model) {\n\t\t\t\trequire.NotNil(t, m.Config)\n\t\t\t\trequire.Nil(t, m.Schema)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"invalid config JSON returns error\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig: `{invalid`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectErr: nil, // Will have an error, just not ErrNotCogModel\n\t\t},\n\t\t{\n\t\t\tname: \"invalid schema JSON returns error\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig:        `{\"build\":{}}`,\n\t\t\t\t\tLabelOpenAPISchema: `{invalid schema`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectErr: nil, // Will have an error, just not ErrNotCogModel\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tm, err := tt.image.ToModel()\n\n\t\t\tif tt.expectErr != nil {\n\t\t\t\trequire.ErrorIs(t, err, tt.expectErr)\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\tif tt.name == \"invalid config JSON returns error\" || tt.name == \"invalid schema JSON returns error\" {\n\t\t\t\trequire.Error(t, err)\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\trequire.NoError(t, err)\n\t\t\tif tt.checkModel != nil {\n\t\t\t\ttt.checkModel(t, m)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestImage_ParsedOpenAPISchema(t *testing.T) {\n\ttests := []struct {\n\t\tname        string\n\t\timage       *ImageArtifact\n\t\texpectNil   bool\n\t\texpectErr   bool\n\t\tcheckSchema func(t *testing.T, schema *openapi3.T)\n\t}{\n\t\t{\n\t\t\tname:      \"nil labels returns nil without error\",\n\t\t\timage:     &ImageArtifact{Labels: nil},\n\t\t\texpectNil: true,\n\t\t\texpectErr: false,\n\t\t},\n\t\t{\n\t\t\tname:      \"empty labels returns nil without error\",\n\t\t\timage:     &ImageArtifact{Labels: map[string]string{}},\n\t\t\texpectNil: true,\n\t\t\texpectErr: false,\n\t\t},\n\t\t{\n\t\t\tname: \"missing schema label returns nil without error\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig: `{\"build\":{}}`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectNil: true,\n\t\t\texpectErr: false,\n\t\t},\n\t\t{\n\t\t\tname: \"valid OpenAPI JSON parses correctly\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelOpenAPISchema: `{\"openapi\":\"3.0.2\",\"info\":{\"title\":\"Cog\",\"version\":\"0.1.0\"},\"paths\":{}}`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectNil: false,\n\t\t\texpectErr: false,\n\t\t\tcheckSchema: func(t *testing.T, schema *openapi3.T) {\n\t\t\t\trequire.Equal(t, \"3.0.2\", schema.OpenAPI)\n\t\t\t\trequire.Equal(t, \"Cog\", schema.Info.Title)\n\t\t\t\trequire.Equal(t, \"0.1.0\", schema.Info.Version)\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname: \"invalid JSON returns error\",\n\t\t\timage: &ImageArtifact{\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelOpenAPISchema: `{invalid json`,\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectNil: true,\n\t\t\texpectErr: true,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tschema, err := tt.image.ParsedOpenAPISchema()\n\n\t\t\tif tt.expectErr {\n\t\t\t\trequire.Error(t, err)\n\t\t\t} else {\n\t\t\t\trequire.NoError(t, err)\n\t\t\t}\n\n\t\t\tif tt.expectNil {\n\t\t\t\trequire.Nil(t, schema)\n\t\t\t} else {\n\t\t\t\trequire.NotNil(t, schema)\n\t\t\t\tif tt.checkSchema != nil {\n\t\t\t\t\ttt.checkSchema(t, schema)\n\t\t\t\t}\n\t\t\t}\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "pkg/model/index.go",
    "content": "// pkg/model/index.go\npackage model\n\n// Index represents an OCI Image Index containing multiple manifests.\ntype Index struct {\n\t// Digest is the index digest (sha256:...).\n\tDigest string\n\t// Reference is the full image reference.\n\tReference string\n\t// MediaType is typically application/vnd.oci.image.index.v1+json.\n\tMediaType string\n\t// Manifests are the child manifests in this index.\n\tManifests []IndexManifest\n}\n\n// IndexManifest represents a single manifest within an index.\ntype IndexManifest struct {\n\t// Digest is the manifest digest.\n\tDigest string\n\t// MediaType is the manifest media type.\n\tMediaType string\n\t// Size is the manifest size in bytes.\n\tSize int64\n\t// Platform is the target platform (nil for artifacts).\n\tPlatform *Platform\n\t// Annotations are OCI annotations on this manifest.\n\tAnnotations map[string]string\n\t// Type is derived from platform/annotations (image or weights).\n\tType ManifestType\n}\n\n// ManifestType identifies the type of manifest in an index.\ntype ManifestType string\n\nconst (\n\t// ManifestTypeImage is a runnable container image.\n\tManifestTypeImage ManifestType = \"image\"\n\t// ManifestTypeWeights is a weights artifact.\n\tManifestTypeWeights ManifestType = \"weights\"\n)\n\n// Annotation keys for weights manifests.\nconst (\n\tAnnotationReferenceType   = \"vnd.cog.reference.type\"\n\tAnnotationReferenceDigest = \"vnd.cog.reference.digest\"\n)\n\n// Annotation values.\nconst (\n\t// AnnotationValueWeights is the value for AnnotationReferenceType indicating a weights manifest.\n\tAnnotationValueWeights = \"weights\"\n)\n\n// Platform values for non-platform-specific artifacts.\nconst (\n\t// PlatformUnknown is used for artifacts that are not platform-specific (e.g., weights).\n\tPlatformUnknown = \"unknown\"\n)\n"
  },
  {
    "path": "pkg/model/index_factory.go",
    "content": "package model\n\nimport (\n\t\"fmt\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/empty\"\n\t\"github.com/google/go-containerregistry/pkg/v1/mutate\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n)\n\n// IndexBuilder builds an OCI Image Index from pre-pushed manifest descriptors.\ntype IndexBuilder struct {\n\timageDescriptor   *v1.Descriptor\n\timagePlatform     *v1.Platform\n\tweightDescriptors []weightDescEntry\n}\n\n// weightDescEntry pairs a weight descriptor with the image digest it references.\ntype weightDescEntry struct {\n\tdescriptor  v1.Descriptor\n\timageDigest string\n\tname        string\n\ttarget      string\n}\n\n// NewIndexBuilder creates a new index builder.\nfunc NewIndexBuilder() *IndexBuilder {\n\treturn &IndexBuilder{}\n}\n\n// SetImageDescriptor sets the image manifest descriptor.\nfunc (b *IndexBuilder) SetImageDescriptor(desc v1.Descriptor, platform *v1.Platform) {\n\tb.imageDescriptor = &desc\n\tb.imagePlatform = platform\n}\n\n// AddWeightDescriptor adds a weight manifest descriptor.\n// imageDigest is the digest of the model image, used in the reference annotation.\n// name and target are optional weight metadata for index annotations.\nfunc (b *IndexBuilder) AddWeightDescriptor(desc v1.Descriptor, imageDigest, name, target string) {\n\tb.weightDescriptors = append(b.weightDescriptors, weightDescEntry{\n\t\tdescriptor:  desc,\n\t\timageDigest: imageDigest,\n\t\tname:        name,\n\t\ttarget:      target,\n\t})\n}\n\n// BuildFromDescriptors creates an OCI Image Index from pre-pushed manifest descriptors.\n// This works with bare descriptors returned from push operations, avoiding the need\n// to fetch images back from the registry.\nfunc (b *IndexBuilder) BuildFromDescriptors() (v1.ImageIndex, error) {\n\tif b.imageDescriptor == nil {\n\t\treturn nil, fmt.Errorf(\"image descriptor not set\")\n\t}\n\n\tidx := mutate.IndexMediaType(empty.Index, types.OCIImageIndex)\n\n\t// Add image manifest\n\tidx = mutate.AppendManifests(idx, mutate.IndexAddendum{\n\t\tAdd: &descriptorAppendable{desc: *b.imageDescriptor},\n\t\tDescriptor: v1.Descriptor{\n\t\t\tMediaType: b.imageDescriptor.MediaType,\n\t\t\tSize:      b.imageDescriptor.Size,\n\t\t\tDigest:    b.imageDescriptor.Digest,\n\t\t\tPlatform:  b.imagePlatform,\n\t\t},\n\t})\n\n\t// Add weight manifest(s)\n\tfor _, entry := range b.weightDescriptors {\n\t\tannotations := map[string]string{\n\t\t\tAnnotationReferenceType: AnnotationValueWeights,\n\t\t}\n\t\tif entry.imageDigest != \"\" {\n\t\t\tannotations[AnnotationReferenceDigest] = entry.imageDigest\n\t\t}\n\t\tif entry.name != \"\" {\n\t\t\tannotations[AnnotationWeightName] = entry.name\n\t\t}\n\t\tif entry.target != \"\" {\n\t\t\tannotations[AnnotationWeightDest] = entry.target\n\t\t}\n\n\t\tidx = mutate.AppendManifests(idx, mutate.IndexAddendum{\n\t\t\tAdd: &descriptorAppendable{desc: entry.descriptor},\n\t\t\tDescriptor: v1.Descriptor{\n\t\t\t\tMediaType: entry.descriptor.MediaType,\n\t\t\t\tSize:      entry.descriptor.Size,\n\t\t\t\tDigest:    entry.descriptor.Digest,\n\t\t\t\tPlatform: &v1.Platform{\n\t\t\t\t\tOS:           PlatformUnknown,\n\t\t\t\t\tArchitecture: PlatformUnknown,\n\t\t\t\t},\n\t\t\t\tAnnotations: annotations,\n\t\t\t},\n\t\t})\n\t}\n\n\treturn idx, nil\n}\n\n// descriptorAppendable wraps a v1.Descriptor to implement mutate.Appendable.\n// This allows building an OCI index from descriptors without needing full v1.Image objects.\ntype descriptorAppendable struct {\n\tdesc v1.Descriptor\n}\n\nfunc (d *descriptorAppendable) MediaType() (types.MediaType, error) {\n\treturn d.desc.MediaType, nil\n}\n\nfunc (d *descriptorAppendable) Digest() (v1.Hash, error) {\n\treturn d.desc.Digest, nil\n}\n\nfunc (d *descriptorAppendable) Size() (int64, error) {\n\treturn d.desc.Size, nil\n}\n"
  },
  {
    "path": "pkg/model/index_factory_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestIndexBuilder_BuildFromDescriptors(t *testing.T) {\n\tt.Run(\"builds index from image and weight descriptors\", func(t *testing.T) {\n\t\timgDesc := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      1234,\n\t\t\tDigest: v1.Hash{\n\t\t\t\tAlgorithm: \"sha256\",\n\t\t\t\tHex:       \"aaaa\",\n\t\t\t},\n\t\t}\n\t\tweightDesc := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      5678,\n\t\t\tDigest: v1.Hash{\n\t\t\t\tAlgorithm: \"sha256\",\n\t\t\t\tHex:       \"bbbb\",\n\t\t\t},\n\t\t}\n\n\t\tbuilder := NewIndexBuilder()\n\t\tbuilder.SetImageDescriptor(imgDesc, &v1.Platform{OS: \"linux\", Architecture: \"amd64\"})\n\t\tbuilder.AddWeightDescriptor(weightDesc, imgDesc.Digest.String(), \"model-v1\", \"/cache/model.safetensors\")\n\n\t\tidx, err := builder.BuildFromDescriptors()\n\t\trequire.NoError(t, err)\n\n\t\tidxManifest, err := idx.IndexManifest()\n\t\trequire.NoError(t, err)\n\t\trequire.Len(t, idxManifest.Manifests, 2)\n\n\t\t// First entry: image with platform\n\t\trequire.Equal(t, imgDesc.Digest, idxManifest.Manifests[0].Digest)\n\t\trequire.Equal(t, imgDesc.Size, idxManifest.Manifests[0].Size)\n\t\trequire.Equal(t, \"linux\", idxManifest.Manifests[0].Platform.OS)\n\t\trequire.Equal(t, \"amd64\", idxManifest.Manifests[0].Platform.Architecture)\n\n\t\t// Second entry: weight artifact with unknown platform and annotations\n\t\trequire.Equal(t, weightDesc.Digest, idxManifest.Manifests[1].Digest)\n\t\trequire.Equal(t, weightDesc.Size, idxManifest.Manifests[1].Size)\n\t\trequire.Equal(t, PlatformUnknown, idxManifest.Manifests[1].Platform.OS)\n\t\trequire.Equal(t, PlatformUnknown, idxManifest.Manifests[1].Platform.Architecture)\n\t\trequire.Equal(t, AnnotationValueWeights, idxManifest.Manifests[1].Annotations[AnnotationReferenceType])\n\t\trequire.Equal(t, imgDesc.Digest.String(), idxManifest.Manifests[1].Annotations[AnnotationReferenceDigest])\n\t\trequire.Equal(t, \"model-v1\", idxManifest.Manifests[1].Annotations[AnnotationWeightName])\n\t\trequire.Equal(t, \"/cache/model.safetensors\", idxManifest.Manifests[1].Annotations[AnnotationWeightDest])\n\t})\n\n\tt.Run(\"builds index with multiple weight descriptors\", func(t *testing.T) {\n\t\timgDesc := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      1000,\n\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"img111\"},\n\t\t}\n\t\tweight1 := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      2000,\n\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"w1111\"},\n\t\t}\n\t\tweight2 := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      3000,\n\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"w2222\"},\n\t\t}\n\n\t\tbuilder := NewIndexBuilder()\n\t\tbuilder.SetImageDescriptor(imgDesc, &v1.Platform{OS: \"linux\", Architecture: \"amd64\"})\n\t\tbuilder.AddWeightDescriptor(weight1, imgDesc.Digest.String(), \"weight-1\", \"/weights/w1.bin\")\n\t\tbuilder.AddWeightDescriptor(weight2, imgDesc.Digest.String(), \"weight-2\", \"/weights/w2.bin\")\n\n\t\tidx, err := builder.BuildFromDescriptors()\n\t\trequire.NoError(t, err)\n\n\t\tidxManifest, err := idx.IndexManifest()\n\t\trequire.NoError(t, err)\n\t\trequire.Len(t, idxManifest.Manifests, 3) // 1 image + 2 weights\n\t})\n\n\tt.Run(\"requires image descriptor\", func(t *testing.T) {\n\t\tbuilder := NewIndexBuilder()\n\t\t_, err := builder.BuildFromDescriptors()\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"image descriptor not set\")\n\t})\n\n\tt.Run(\"builds index without weight descriptors\", func(t *testing.T) {\n\t\timgDesc := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      1234,\n\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"aaaa\"},\n\t\t}\n\n\t\tbuilder := NewIndexBuilder()\n\t\tbuilder.SetImageDescriptor(imgDesc, &v1.Platform{OS: \"linux\", Architecture: \"amd64\"})\n\n\t\tidx, err := builder.BuildFromDescriptors()\n\t\trequire.NoError(t, err)\n\n\t\tidxManifest, err := idx.IndexManifest()\n\t\trequire.NoError(t, err)\n\t\trequire.Len(t, idxManifest.Manifests, 1)\n\t})\n\n\tt.Run(\"index has OCI media type\", func(t *testing.T) {\n\t\timgDesc := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      1234,\n\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"aaaa\"},\n\t\t}\n\n\t\tbuilder := NewIndexBuilder()\n\t\tbuilder.SetImageDescriptor(imgDesc, &v1.Platform{OS: \"linux\", Architecture: \"amd64\"})\n\n\t\tidx, err := builder.BuildFromDescriptors()\n\t\trequire.NoError(t, err)\n\n\t\tmt, err := idx.MediaType()\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, types.OCIImageIndex, mt)\n\t})\n}\n"
  },
  {
    "path": "pkg/model/index_test.go",
    "content": "// pkg/model/index_test.go\npackage model\n\nimport (\n\t\"testing\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestModel_IsBundle(t *testing.T) {\n\tt.Run(\"returns false with no artifacts\", func(t *testing.T) {\n\t\tm := &Model{}\n\t\trequire.False(t, m.IsBundle())\n\t})\n\n\tt.Run(\"returns false with only image artifact\", func(t *testing.T) {\n\t\tm := &Model{\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t},\n\t\t}\n\t\trequire.False(t, m.IsBundle())\n\t})\n\n\tt.Run(\"returns true with weight artifacts\", func(t *testing.T) {\n\t\tm := &Model{\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\tNewWeightArtifact(\"w1\", v1.Descriptor{}, \"/tmp/w1\", \"/weights/w1\", WeightConfig{}),\n\t\t\t},\n\t\t}\n\t\trequire.True(t, m.IsBundle())\n\t})\n}\n\nfunc TestManifestType(t *testing.T) {\n\trequire.Equal(t, ManifestType(\"image\"), ManifestTypeImage)\n\trequire.Equal(t, ManifestType(\"weights\"), ManifestTypeWeights)\n}\n"
  },
  {
    "path": "pkg/model/model.go",
    "content": "package model\n\nimport (\n\t\"encoding/json\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\n// Model represents a Cog model extracted from an image.\ntype Model struct {\n\tImage      *ImageArtifact // Underlying OCI image\n\tConfig     *config.Config // Parsed cog.yaml\n\tSchema     *openapi3.T    // OpenAPI schema\n\tCogVersion string         // Version of cog used to build\n\n\t// Index is the OCI Image Index (populated when inspecting a pushed model).\n\tIndex *Index\n\n\t// TODO(md): OCIIndex is a temporary gate. When true, Push() creates an OCI\n\t// Image Index with weight artifacts. When false, Push() does a plain docker push.\n\t// Remove this field once index pushes are validated with all registries.\n\tOCIIndex bool\n\n\t// Artifacts is the collection of all artifacts produced by building this model.\n\t// Populated by Resolver.Build(). Contains ImageArtifact and WeightArtifact instances.\n\tArtifacts []Artifact\n}\n\n// HasGPU returns true if the model requires GPU.\nfunc (m *Model) HasGPU() bool {\n\treturn m.Config != nil && m.Config.Build != nil && m.Config.Build.GPU\n}\n\n// SchemaJSON returns the OpenAPI schema as JSON bytes, or nil if no schema.\nfunc (m *Model) SchemaJSON() ([]byte, error) {\n\tif m.Schema == nil {\n\t\treturn nil, nil\n\t}\n\treturn json.Marshal(m.Schema)\n}\n\n// ImageRef returns the image reference string, or empty string if no image.\nfunc (m *Model) ImageRef() string {\n\tif m.Image == nil {\n\t\treturn \"\"\n\t}\n\treturn m.Image.Reference\n}\n\n// IsBundle returns true if this model has weight artifacts.\nfunc (m *Model) IsBundle() bool {\n\treturn len(m.WeightArtifacts()) > 0\n}\n\n// GetImageArtifact returns the first ImageArtifact from the artifacts collection,\n// or nil if none exists.\nfunc (m *Model) GetImageArtifact() *ImageArtifact {\n\tfor _, a := range m.Artifacts {\n\t\tif ia, ok := a.(*ImageArtifact); ok {\n\t\t\treturn ia\n\t\t}\n\t}\n\treturn nil\n}\n\n// WeightArtifacts returns all WeightArtifact instances from the artifacts collection.\nfunc (m *Model) WeightArtifacts() []*WeightArtifact {\n\tvar weights []*WeightArtifact\n\tfor _, a := range m.Artifacts {\n\t\tif wa, ok := a.(*WeightArtifact); ok {\n\t\t\tweights = append(weights, wa)\n\t\t}\n\t}\n\treturn weights\n}\n\n// ArtifactsByType returns all artifacts matching the given type.\nfunc (m *Model) ArtifactsByType(t ArtifactType) []Artifact {\n\tvar result []Artifact\n\tfor _, a := range m.Artifacts {\n\t\tif a.Type() == t {\n\t\t\tresult = append(result, a)\n\t\t}\n\t}\n\treturn result\n}\n"
  },
  {
    "path": "pkg/model/model_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc TestModel_HasGPU(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\tmodel  *Model\n\t\texpect bool\n\t}{\n\t\t{\n\t\t\tname:   \"nil config\",\n\t\t\tmodel:  &Model{Config: nil},\n\t\t\texpect: false,\n\t\t},\n\t\t{\n\t\t\tname:   \"nil build\",\n\t\t\tmodel:  &Model{Config: &config.Config{Build: nil}},\n\t\t\texpect: false,\n\t\t},\n\t\t{\n\t\t\tname:   \"GPU false\",\n\t\t\tmodel:  &Model{Config: &config.Config{Build: &config.Build{GPU: false}}},\n\t\t\texpect: false,\n\t\t},\n\t\t{\n\t\t\tname:   \"GPU true\",\n\t\t\tmodel:  &Model{Config: &config.Config{Build: &config.Build{GPU: true}}},\n\t\t\texpect: true,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.model.HasGPU()\n\t\t\trequire.Equal(t, tt.expect, result)\n\t\t})\n\t}\n}\n\nfunc TestModel_SchemaJSON(t *testing.T) {\n\ttests := []struct {\n\t\tname       string\n\t\tmodel      *Model\n\t\texpectNil  bool\n\t\texpectJSON string\n\t}{\n\t\t{\n\t\t\tname:      \"nil schema\",\n\t\t\tmodel:     &Model{Schema: nil},\n\t\t\texpectNil: true,\n\t\t},\n\t\t{\n\t\t\tname: \"schema with openapi version\",\n\t\t\tmodel: &Model{\n\t\t\t\tSchema: &openapi3.T{\n\t\t\t\t\tOpenAPI: \"3.0.0\",\n\t\t\t\t},\n\t\t\t},\n\t\t\texpectNil: false,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := tt.model.SchemaJSON()\n\t\t\trequire.NoError(t, err)\n\t\t\tif tt.expectNil {\n\t\t\t\trequire.Nil(t, result)\n\t\t\t} else {\n\t\t\t\trequire.NotNil(t, result)\n\t\t\t\t// Verify it's valid JSON containing expected field\n\t\t\t\trequire.Contains(t, string(result), `\"openapi\"`)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestModel_ImageRef(t *testing.T) {\n\ttests := []struct {\n\t\tname   string\n\t\tmodel  *Model\n\t\texpect string\n\t}{\n\t\t{\n\t\t\tname:   \"nil image\",\n\t\t\tmodel:  &Model{Image: nil},\n\t\t\texpect: \"\",\n\t\t},\n\t\t{\n\t\t\tname: \"with image reference\",\n\t\t\tmodel: &Model{\n\t\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model@sha256:abc123\"},\n\t\t\t},\n\t\t\texpect: \"r8.im/user/model@sha256:abc123\",\n\t\t},\n\t\t{\n\t\t\tname: \"with empty reference\",\n\t\t\tmodel: &Model{\n\t\t\t\tImage: &ImageArtifact{Reference: \"\"},\n\t\t\t},\n\t\t\texpect: \"\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.model.ImageRef()\n\t\t\trequire.Equal(t, tt.expect, result)\n\t\t})\n\t}\n}\n\nfunc TestModel_GetImageArtifact(t *testing.T) {\n\timgArtifact := NewImageArtifact(\"model\",\n\t\tv1.Descriptor{Digest: v1.Hash{Algorithm: \"sha256\", Hex: \"abc123\"}, Size: 1024},\n\t\t\"r8.im/user/model@sha256:abc123\",\n\t)\n\tweightArtifact := NewWeightArtifact(\"weights\",\n\t\tv1.Descriptor{Digest: v1.Hash{Algorithm: \"sha256\", Hex: \"def456\"}, Size: 4096},\n\t\t\"/data/weights.bin\", \"/weights/model.bin\",\n\t\tWeightConfig{SchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"weights\", Target: \"/weights/model.bin\", Created: time.Now()},\n\t)\n\n\ttests := []struct {\n\t\tname      string\n\t\tmodel     *Model\n\t\texpectNil bool\n\t}{\n\t\t{\n\t\t\tname:      \"no artifacts\",\n\t\t\tmodel:     &Model{},\n\t\t\texpectNil: true,\n\t\t},\n\t\t{\n\t\t\tname:      \"nil artifacts\",\n\t\t\tmodel:     &Model{Artifacts: nil},\n\t\t\texpectNil: true,\n\t\t},\n\t\t{\n\t\t\tname:      \"only weight artifacts\",\n\t\t\tmodel:     &Model{Artifacts: []Artifact{weightArtifact}},\n\t\t\texpectNil: true,\n\t\t},\n\t\t{\n\t\t\tname:      \"has image artifact\",\n\t\t\tmodel:     &Model{Artifacts: []Artifact{imgArtifact, weightArtifact}},\n\t\t\texpectNil: false,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.model.GetImageArtifact()\n\t\t\tif tt.expectNil {\n\t\t\t\trequire.Nil(t, result)\n\t\t\t} else {\n\t\t\t\trequire.NotNil(t, result)\n\t\t\t\trequire.Equal(t, ArtifactTypeImage, result.Type())\n\t\t\t\trequire.Equal(t, \"model\", result.Name())\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestModel_WeightArtifacts(t *testing.T) {\n\timgArtifact := NewImageArtifact(\"model\",\n\t\tv1.Descriptor{Digest: v1.Hash{Algorithm: \"sha256\", Hex: \"abc123\"}, Size: 1024},\n\t\t\"r8.im/user/model@sha256:abc123\",\n\t)\n\tw1 := NewWeightArtifact(\"llama\",\n\t\tv1.Descriptor{Digest: v1.Hash{Algorithm: \"sha256\", Hex: \"w1\"}, Size: 4096},\n\t\t\"/data/llama.bin\", \"/weights/llama.bin\",\n\t\tWeightConfig{SchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"llama\", Target: \"/weights/llama.bin\", Created: time.Now()},\n\t)\n\tw2 := NewWeightArtifact(\"embeddings\",\n\t\tv1.Descriptor{Digest: v1.Hash{Algorithm: \"sha256\", Hex: \"w2\"}, Size: 2048},\n\t\t\"/data/embed.bin\", \"/weights/embed.bin\",\n\t\tWeightConfig{SchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"embeddings\", Target: \"/weights/embed.bin\", Created: time.Now()},\n\t)\n\n\ttests := []struct {\n\t\tname   string\n\t\tmodel  *Model\n\t\texpect int\n\t}{\n\t\t{name: \"no artifacts\", model: &Model{}, expect: 0},\n\t\t{name: \"only image\", model: &Model{Artifacts: []Artifact{imgArtifact}}, expect: 0},\n\t\t{name: \"one weight\", model: &Model{Artifacts: []Artifact{imgArtifact, w1}}, expect: 1},\n\t\t{name: \"two weights\", model: &Model{Artifacts: []Artifact{imgArtifact, w1, w2}}, expect: 2},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.model.WeightArtifacts()\n\t\t\trequire.Len(t, result, tt.expect)\n\t\t})\n\t}\n}\n\nfunc TestModel_ArtifactsByType(t *testing.T) {\n\timgArtifact := NewImageArtifact(\"model\",\n\t\tv1.Descriptor{Digest: v1.Hash{Algorithm: \"sha256\", Hex: \"abc123\"}, Size: 1024},\n\t\t\"r8.im/user/model@sha256:abc123\",\n\t)\n\tw1 := NewWeightArtifact(\"llama\",\n\t\tv1.Descriptor{Digest: v1.Hash{Algorithm: \"sha256\", Hex: \"w1\"}, Size: 4096},\n\t\t\"/data/llama.bin\", \"/weights/llama.bin\",\n\t\tWeightConfig{SchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"llama\", Target: \"/weights/llama.bin\", Created: time.Now()},\n\t)\n\n\tm := &Model{Artifacts: []Artifact{imgArtifact, w1}}\n\n\timages := m.ArtifactsByType(ArtifactTypeImage)\n\trequire.Len(t, images, 1)\n\trequire.Equal(t, \"model\", images[0].Name())\n\n\tweights := m.ArtifactsByType(ArtifactTypeWeight)\n\trequire.Len(t, weights, 1)\n\trequire.Equal(t, \"llama\", weights[0].Name())\n}\n"
  },
  {
    "path": "pkg/model/options.go",
    "content": "package model\n\nimport \"github.com/replicate/cog/pkg/config\"\n\n// BuildOptions contains all settings for building a Cog image.\n// This consolidates the many parameters previously passed to image.Build().\ntype BuildOptions struct {\n\t// ImageName is the output image name (required).\n\tImageName string\n\n\t// NoCache disables build cache.\n\tNoCache bool\n\n\t// SeparateWeights builds weights as a separate layer.\n\tSeparateWeights bool\n\n\t// Strip removes debug symbols from binaries.\n\tStrip bool\n\n\t// Precompile precompiles Python bytecode.\n\tPrecompile bool\n\n\t// UseCudaBaseImage controls CUDA base image usage: \"auto\", \"true\", or \"false\".\n\tUseCudaBaseImage string\n\n\t// UseCogBaseImage controls cog base image usage. nil means auto-detect.\n\tUseCogBaseImage *bool\n\n\t// Secrets are build-time secrets to pass to the build.\n\tSecrets []string\n\n\t// ProgressOutput controls build output format: \"auto\", \"plain\", or \"tty\".\n\tProgressOutput string\n\n\t// Annotations are extra labels to add to the image.\n\tAnnotations map[string]string\n\n\t// SchemaFile is a custom OpenAPI schema file path.\n\tSchemaFile string\n\n\t// DockerfileFile is a custom Dockerfile path.\n\tDockerfileFile string\n\n\t// WeightsLockPath is the path to weights.lock file.\n\t// Default: weights.lock in project directory.\n\tWeightsLockPath string\n\n\t// TODO(md): OCIIndex is a temporary gate. When true, builds produce weight\n\t// artifacts and pushes create an OCI Image Index. Set via COG_OCI_INDEX=1.\n\t// Remove this field once index pushes are validated with all registries.\n\tOCIIndex bool\n\n\t// ExcludeSource skips the COPY . /src step in the generated Dockerfile.\n\t// Used by `cog serve` to produce an image identical to `cog build` minus\n\t// the source copy — the source directory is volume-mounted at runtime.\n\t// All other layers (wheel installs, apt, etc.) are shared with `cog build`\n\t// via Docker layer caching.\n\tExcludeSource bool\n\n\t// SkipSchemaValidation skips OpenAPI schema generation and validation.\n\t// Used by `cog run` which executes arbitrary commands and may not have\n\t// a predictor or trainer defined in cog.yaml.\n\tSkipSchemaValidation bool\n\n\t// SkipLabels skips adding Cog metadata labels to the built image.\n\t// Used by `cog run`, `cog predict`, `cog serve`, and `cog train` where\n\t// the image is for local use only and not being distributed.\n\tSkipLabels bool\n}\n\n// WithDefaults returns a copy of BuildOptions with defaults applied from Source.\n// This fills in sensible defaults for any unset fields.\nfunc (o BuildOptions) WithDefaults(src *Source) BuildOptions {\n\t// Default image name from project directory\n\tif o.ImageName == \"\" {\n\t\to.ImageName = config.DockerImageName(src.ProjectDir)\n\t}\n\n\t// Default progress output\n\tif o.ProgressOutput == \"\" {\n\t\to.ProgressOutput = \"auto\"\n\t}\n\n\treturn o\n}\n"
  },
  {
    "path": "pkg/model/options_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc TestBuildOptions_WithDefaults_ImageName(t *testing.T) {\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Build: &config.Build{}},\n\t\tProjectDir: \"/path/to/my-project\",\n\t}\n\n\topts := BuildOptions{}\n\topts = opts.WithDefaults(src)\n\n\t// config.DockerImageName normalizes the name\n\trequire.Equal(t, \"cog-my-project\", opts.ImageName)\n}\n\nfunc TestBuildOptions_WithDefaults_PreservesExplicitImageName(t *testing.T) {\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Build: &config.Build{}},\n\t\tProjectDir: \"/path/to/my-project\",\n\t}\n\n\topts := BuildOptions{ImageName: \"my-custom-image\"}\n\topts = opts.WithDefaults(src)\n\n\trequire.Equal(t, \"my-custom-image\", opts.ImageName)\n}\n\nfunc TestBuildOptions_WithDefaults_ProgressOutput(t *testing.T) {\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Build: &config.Build{}},\n\t\tProjectDir: \"/path/to/project\",\n\t}\n\n\topts := BuildOptions{}\n\topts = opts.WithDefaults(src)\n\n\trequire.Equal(t, \"auto\", opts.ProgressOutput)\n}\n\nfunc TestBuildOptions_WithDefaults_PreservesExplicitProgressOutput(t *testing.T) {\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Build: &config.Build{}},\n\t\tProjectDir: \"/path/to/project\",\n\t}\n\n\topts := BuildOptions{ProgressOutput: \"plain\"}\n\topts = opts.WithDefaults(src)\n\n\trequire.Equal(t, \"plain\", opts.ProgressOutput)\n}\n\nfunc TestBuildOptions_WithDefaults_NilBuildConfig(t *testing.T) {\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Build: nil},\n\t\tProjectDir: \"/path/to/project\",\n\t}\n\n\topts := BuildOptions{}\n\topts = opts.WithDefaults(src)\n\n\t// Should not panic and should apply other defaults\n\trequire.Equal(t, \"cog-project\", opts.ImageName)\n\trequire.Equal(t, \"auto\", opts.ProgressOutput)\n}\n\nfunc TestBuildOptions_WithDefaults_NilConfig(t *testing.T) {\n\tsrc := &Source{\n\t\tConfig:     nil,\n\t\tProjectDir: \"/path/to/project\",\n\t}\n\n\topts := BuildOptions{}\n\topts = opts.WithDefaults(src)\n\n\t// Should not panic and should apply other defaults\n\trequire.Equal(t, \"cog-project\", opts.ImageName)\n\trequire.Equal(t, \"auto\", opts.ProgressOutput)\n}\n\nfunc TestBuildOptions_AllFieldsPreserved(t *testing.T) {\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Build: &config.Build{}},\n\t\tProjectDir: \"/path/to/project\",\n\t}\n\n\tuseCogBase := true\n\topts := BuildOptions{\n\t\tImageName:        \"my-image\",\n\t\tNoCache:          true,\n\t\tSeparateWeights:  true,\n\t\tStrip:            true,\n\t\tPrecompile:       true,\n\t\tUseCudaBaseImage: \"true\",\n\t\tUseCogBaseImage:  &useCogBase,\n\t\tSecrets:          []string{\"secret1\", \"secret2\"},\n\t\tProgressOutput:   \"tty\",\n\t\tAnnotations:      map[string]string{\"key\": \"value\"},\n\t\tSchemaFile:       \"/path/to/schema.json\",\n\t\tDockerfileFile:   \"/path/to/Dockerfile\",\n\t\tWeightsLockPath:  \"/path/to/weights.lock\",\n\t}\n\n\tresult := opts.WithDefaults(src)\n\n\trequire.Equal(t, \"my-image\", result.ImageName)\n\trequire.True(t, result.NoCache)\n\trequire.True(t, result.SeparateWeights)\n\trequire.True(t, result.Strip)\n\trequire.True(t, result.Precompile)\n\trequire.Equal(t, \"true\", result.UseCudaBaseImage)\n\trequire.NotNil(t, result.UseCogBaseImage)\n\trequire.True(t, *result.UseCogBaseImage)\n\trequire.Equal(t, []string{\"secret1\", \"secret2\"}, result.Secrets)\n\trequire.Equal(t, \"tty\", result.ProgressOutput)\n\trequire.Equal(t, map[string]string{\"key\": \"value\"}, result.Annotations)\n\trequire.Equal(t, \"/path/to/schema.json\", result.SchemaFile)\n\trequire.Equal(t, \"/path/to/Dockerfile\", result.DockerfileFile)\n\trequire.Equal(t, \"/path/to/weights.lock\", result.WeightsLockPath)\n}\n\nfunc TestBuildOptions_WeightsLockPath(t *testing.T) {\n\topts := BuildOptions{\n\t\tWeightsLockPath: \"/custom/weights.lock\",\n\t}\n\trequire.Equal(t, \"/custom/weights.lock\", opts.WeightsLockPath)\n}\n"
  },
  {
    "path": "pkg/model/push_helpers.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"os\"\n\t\"strconv\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\nconst (\n\t// DefaultPushConcurrency is the default number of concurrent uploads\n\t// for both image layers and weight artifacts.\n\t// This matches Docker's default concurrency for layer uploads, which is a reasonable baseline for OCI pushes as well.\n\tDefaultPushConcurrency = 5\n\n\t// envPushConcurrency is the environment variable that overrides DefaultPushConcurrency.\n\tenvPushConcurrency = \"COG_PUSH_CONCURRENCY\"\n)\n\n// GetPushConcurrency returns the push concurrency, checking the COG_PUSH_CONCURRENCY\n// environment variable first, then falling back to DefaultPushConcurrency.\nfunc GetPushConcurrency() int {\n\tif v := os.Getenv(envPushConcurrency); v != \"\" {\n\t\tif n, err := strconv.Atoi(v); err == nil && n > 0 {\n\t\t\treturn n\n\t\t}\n\t}\n\treturn DefaultPushConcurrency\n}\n\n// PushProgress reports progress for a layer or blob upload.\n// Used by both ImagePusher (container image layers) and WeightPusher (weight blobs).\ntype PushProgress struct {\n\t// LayerDigest identifies which layer this progress is for.\n\t// Empty for single-layer pushes (e.g., weight uploads).\n\tLayerDigest string\n\t// Complete is the number of bytes uploaded so far.\n\tComplete int64\n\t// Total is the total number of bytes to upload.\n\tTotal int64\n}\n\n// writeLayerWithProgress pushes a layer via registry.WriteLayer, managing the\n// progress channel lifecycle (create, drain, close) on behalf of the caller.\n//\n// onProgress is called for each v1.Update from the registry. If nil, no progress\n// channel is created and no goroutine is spawned.\nfunc writeLayerWithProgress(ctx context.Context, reg registry.Client, opts registry.WriteLayerOptions, onProgress func(v1.Update)) error {\n\tvar progressCh chan v1.Update\n\tvar progressDone chan struct{}\n\n\tif onProgress != nil {\n\t\tprogressCh = make(chan v1.Update, 100)\n\t\tprogressDone = make(chan struct{})\n\t\tgo func() {\n\t\t\tdefer close(progressDone)\n\t\t\tfor update := range progressCh {\n\t\t\t\tonProgress(update)\n\t\t\t}\n\t\t}()\n\t\topts.ProgressCh = progressCh\n\t}\n\n\twriteErr := reg.WriteLayer(ctx, opts)\n\n\t// Close the progress channel ourselves — WriteLayer sends to it but does not close it.\n\tif progressCh != nil {\n\t\tclose(progressCh)\n\t}\n\tif progressDone != nil {\n\t\t<-progressDone\n\t}\n\n\treturn writeErr\n}\n"
  },
  {
    "path": "pkg/model/pusher.go",
    "content": "// pkg/model/pusher.go\npackage model\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"golang.org/x/sync/errgroup\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// PushOptions configures push behavior.\ntype PushOptions struct {\n\t// ProjectDir is the base directory for resolving weight file paths.\n\t//\n\t// Deprecated: Artifacts carry their own file paths.\n\tProjectDir string\n\n\t// FilePaths maps weight name identifiers to their file paths.\n\t//\n\t// Deprecated: Use Model.Artifacts instead — WeightArtifact carries FilePath.\n\tFilePaths map[string]string\n\n\t// Platform specifies the target platform for bundle indexes.\n\t// Default: linux/amd64\n\tPlatform *Platform\n\n\t// ImageProgressFn is an optional callback for reporting per-layer upload progress\n\t// during OCI chunked image push. Each call includes the layer digest, bytes\n\t// completed, and total bytes.\n\tImageProgressFn func(PushProgress)\n\n\t// OnFallback is called when OCI push fails and the push is about to fall\n\t// back to Docker push. This allows the caller to clean up any OCI-specific\n\t// progress display before Docker push starts its own output.\n\tOnFallback func()\n}\n\n// =============================================================================\n// BundlePusher - pushes OCI Index with image + weights\n// =============================================================================\n\n// BundlePusher pushes bundles (OCI Index with image + weight artifacts).\n// It orchestrates ImagePusher and WeightPusher, then assembles the OCI index\n// from the pushed manifest descriptors.\ntype BundlePusher struct {\n\timagePusher  *ImagePusher\n\tweightPusher *WeightPusher\n\tregistry     registry.Client\n}\n\n// NewBundlePusher creates a new BundlePusher from docker and registry clients.\n// Both sub-pushers (image and weight) are created internally to keep\n// construction unified — callers don't need to know about ImagePusher or\n// WeightPusher directly.\nfunc NewBundlePusher(docker command.Command, reg registry.Client) *BundlePusher {\n\treturn &BundlePusher{\n\t\timagePusher:  newImagePusher(docker, reg),\n\t\tweightPusher: NewWeightPusher(reg),\n\t\tregistry:     reg,\n\t}\n}\n\n// Push pushes the model as an OCI Index with weight artifacts.\n// It reads Model.Artifacts to find the image and weight artifacts to push.\nfunc (p *BundlePusher) Push(ctx context.Context, m *Model, opts PushOptions) error {\n\t// Extract artifacts from model\n\timgArtifact := m.GetImageArtifact()\n\tif imgArtifact == nil {\n\t\treturn fmt.Errorf(\"no image artifact in model\")\n\t}\n\n\tweightArtifacts := m.WeightArtifacts()\n\n\t// Derive repo from image reference (strip tag/digest for weight pushes)\n\trepo := repoFromReference(imgArtifact.Reference)\n\n\t// 1. Push image via OCI chunked push (falls back to Docker push on error)\n\tvar imagePushOpts []ImagePushOption\n\tif opts.ImageProgressFn != nil {\n\t\timagePushOpts = append(imagePushOpts, WithProgressFn(opts.ImageProgressFn))\n\t}\n\tif opts.OnFallback != nil {\n\t\timagePushOpts = append(imagePushOpts, WithOnFallback(opts.OnFallback))\n\t}\n\tif err := p.imagePusher.Push(ctx, imgArtifact, imagePushOpts...); err != nil {\n\t\treturn fmt.Errorf(\"push image %q: %w\", imgArtifact.Reference, err)\n\t}\n\n\t// 2. Get image manifest descriptor (lightweight HEAD request)\n\timgDesc, err := p.registry.GetDescriptor(ctx, imgArtifact.Reference)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"get image descriptor: %w\", err)\n\t}\n\n\t// 3. Push weight artifacts concurrently (if any)\n\tvar weightResults []*WeightPushResult\n\tif len(weightArtifacts) > 0 {\n\t\tweightResults, err = p.pushWeights(ctx, repo, weightArtifacts)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\t// 4. Build OCI index from pushed descriptors\n\tplatform := opts.Platform\n\tif platform == nil {\n\t\tplatform = &Platform{OS: \"linux\", Architecture: \"amd64\"}\n\t}\n\n\tbuilder := NewIndexBuilder()\n\tbuilder.SetImageDescriptor(imgDesc, &v1.Platform{\n\t\tOS:           platform.OS,\n\t\tArchitecture: platform.Architecture,\n\t\tVariant:      platform.Variant,\n\t})\n\tfor i, wr := range weightResults {\n\t\tbuilder.AddWeightDescriptor(wr.Descriptor, imgDesc.Digest.String(),\n\t\t\tweightArtifacts[i].Name(), weightArtifacts[i].Target)\n\t}\n\n\tidx, err := builder.BuildFromDescriptors()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"build OCI index: %w\", err)\n\t}\n\n\t// 5. Push OCI index (overwrites the tag with the index)\n\tif err := p.registry.PushIndex(ctx, imgArtifact.Reference, idx); err != nil {\n\t\treturn fmt.Errorf(\"push OCI index: %w\", err)\n\t}\n\n\treturn nil\n}\n\n// pushWeights pushes all weight artifacts concurrently (bounded by GetPushConcurrency)\n// and returns their results in the same order as the input slice.\n// If any weight push fails, remaining pushes are canceled and the first error is returned.\nfunc (p *BundlePusher) pushWeights(ctx context.Context, repo string, weights []*WeightArtifact) ([]*WeightPushResult, error) {\n\tordered := make([]*WeightPushResult, len(weights))\n\n\tg, ctx := errgroup.WithContext(ctx)\n\tg.SetLimit(GetPushConcurrency())\n\n\tfor i, wa := range weights {\n\t\tg.Go(func() error {\n\t\t\tresult, err := p.weightPusher.Push(ctx, repo, wa)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"push weight %q: %w\", wa.Name(), err)\n\t\t\t}\n\t\t\tordered[i] = result\n\t\t\treturn nil\n\t\t})\n\t}\n\n\tif err := g.Wait(); err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn ordered, nil\n}\n\n// repoFromReference extracts the repository (without tag or digest) from an image reference.\n// \"r8.im/user/model:latest\" -> \"r8.im/user/model\"\n// \"r8.im/user/model@sha256:abc\" -> \"r8.im/user/model\"\n// \"localhost:5000/model:latest\" -> \"localhost:5000/model\"\nfunc repoFromReference(ref string) string {\n\tparsed, err := name.ParseReference(ref, name.Insecure)\n\tif err != nil {\n\t\treturn ref // fallback: return as-is if unparseable\n\t}\n\treturn parsed.Context().String()\n}\n"
  },
  {
    "path": "pkg/model/pusher_test.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"sync\"\n\t\"sync/atomic\"\n\t\"testing\"\n\t\"time\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n\t\"github.com/stretchr/testify/require\"\n)\n\n// =============================================================================\n// BundlePusher tests\n// =============================================================================\n\nfunc TestBundlePusher_Push(t *testing.T) {\n\tt.Run(\"returns error when no image artifact in model\", func(t *testing.T) {\n\t\tdocker := &mockDocker{}\n\t\treg := &mockRegistry{}\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage:     nil,\n\t\t\tArtifacts: []Artifact{}, // no image artifact\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{})\n\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"no image artifact\")\n\t})\n\n\tt.Run(\"pushes image-only model as single-entry index\", func(t *testing.T) {\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error { return nil },\n\t\t}\n\n\t\timgDesc := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      1234,\n\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"imgonly\"},\n\t\t}\n\n\t\treg := &mockRegistry{\n\t\t\tgetDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) {\n\t\t\t\treturn imgDesc, nil\n\t\t\t},\n\t\t\tpushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\t\t\t\t// Verify index has exactly 1 entry (image only, no weights)\n\t\t\t\tidxManifest, err := idx.IndexManifest()\n\t\t\t\trequire.NoError(t, err)\n\t\t\t\trequire.Len(t, idxManifest.Manifests, 1)\n\t\t\t\trequire.Equal(t, imgDesc.Digest, idxManifest.Manifests[0].Digest)\n\t\t\t\trequire.Equal(t, \"linux\", idxManifest.Manifests[0].Platform.OS)\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\t// no weight artifacts — image-only model\n\t\t\t},\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{})\n\t\trequire.NoError(t, err)\n\t})\n\n\tt.Run(\"full push flow succeeds with single weight\", func(t *testing.T) {\n\t\t// Create temp weight file\n\t\tdir := t.TempDir()\n\t\tweightPath := filepath.Join(dir, \"model.safetensors\")\n\t\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"fake weight data\"), 0o644))\n\n\t\t// Track call sequence (mutex-protected for goroutine safety)\n\t\tvar mu sync.Mutex\n\t\tvar callOrder []string\n\t\ttrack := func(entry string) {\n\t\t\tmu.Lock()\n\t\t\tcallOrder = append(callOrder, entry)\n\t\t\tmu.Unlock()\n\t\t}\n\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error {\n\t\t\t\ttrack(\"docker:push:\" + ref)\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\timgDesc := v1.Descriptor{\n\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\tSize:      1234,\n\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"imgdigest\"},\n\t\t}\n\n\t\treg := &mockRegistry{\n\t\t\tgetDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) {\n\t\t\t\ttrack(\"registry:getDescriptor:\" + ref)\n\t\t\t\treturn imgDesc, nil\n\t\t\t},\n\t\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\t\ttrack(\"registry:pushImage:\" + ref)\n\t\t\t\treturn nil\n\t\t\t},\n\t\t\tpushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\t\t\t\ttrack(\"registry:pushIndex:\" + ref)\n\n\t\t\t\t// Verify the index structure\n\t\t\t\tidxManifest, err := idx.IndexManifest()\n\t\t\t\trequire.NoError(t, err)\n\t\t\t\trequire.Len(t, idxManifest.Manifests, 2) // image + 1 weight\n\n\t\t\t\t// First manifest: image with platform\n\t\t\t\trequire.Equal(t, imgDesc.Digest, idxManifest.Manifests[0].Digest)\n\t\t\t\trequire.Equal(t, \"linux\", idxManifest.Manifests[0].Platform.OS)\n\t\t\t\trequire.Equal(t, \"amd64\", idxManifest.Manifests[0].Platform.Architecture)\n\n\t\t\t\t// Second manifest: weight with annotations\n\t\t\t\trequire.Equal(t, PlatformUnknown, idxManifest.Manifests[1].Platform.OS)\n\t\t\t\trequire.Equal(t, AnnotationValueWeights, idxManifest.Manifests[1].Annotations[AnnotationReferenceType])\n\t\t\t\trequire.Equal(t, imgDesc.Digest.String(), idxManifest.Manifests[1].Annotations[AnnotationReferenceDigest])\n\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\tNewWeightArtifact(\"model-v1\", v1.Descriptor{\n\t\t\t\t\tDigest: v1.Hash{Algorithm: \"sha256\", Hex: \"aabbccddee112233445566778899aabb\"},\n\t\t\t\t}, weightPath, \"/weights/model.safetensors\", WeightConfig{\n\t\t\t\t\tSchemaVersion: \"1.0\",\n\t\t\t\t\tCogVersion:    \"0.15.0\",\n\t\t\t\t\tName:          \"model-v1\",\n\t\t\t\t\tTarget:        \"/weights/model.safetensors\",\n\t\t\t\t\tCreated:       time.Now().UTC(),\n\t\t\t\t}),\n\t\t\t},\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{\n\t\t\tPlatform: &Platform{OS: \"linux\", Architecture: \"amd64\"},\n\t\t})\n\n\t\trequire.NoError(t, err)\n\n\t\t// Verify the call sequence:\n\t\t// 1. Push image via docker\n\t\t// 2. Get image descriptor from registry (lightweight HEAD)\n\t\t// 3. Push weight via registry (single combined tag)\n\t\t// 4. Push OCI index to registry\n\t\trequire.Len(t, callOrder, 4)\n\t\trequire.Equal(t, \"docker:push:r8.im/user/model:latest\", callOrder[0])\n\t\trequire.Equal(t, \"registry:getDescriptor:r8.im/user/model:latest\", callOrder[1])\n\t\trequire.Equal(t, \"registry:pushImage:r8.im/user/model:weights-model-v1-aabbccddee11\", callOrder[2])\n\t\trequire.Equal(t, \"registry:pushIndex:r8.im/user/model:latest\", callOrder[3])\n\t})\n\n\tt.Run(\"uses default platform when not specified\", func(t *testing.T) {\n\t\tdir := t.TempDir()\n\t\tweightPath := filepath.Join(dir, \"model.bin\")\n\t\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error { return nil },\n\t\t}\n\n\t\treg := &mockRegistry{\n\t\t\tgetDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) {\n\t\t\t\treturn v1.Descriptor{\n\t\t\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\t\t\tSize:      100,\n\t\t\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"abc\"},\n\t\t\t\t}, nil\n\t\t\t},\n\t\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error { return nil },\n\t\t\tpushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\t\t\t\tidxManifest, _ := idx.IndexManifest()\n\t\t\t\t// Default platform should be linux/amd64\n\t\t\t\trequire.Equal(t, \"linux\", idxManifest.Manifests[0].Platform.OS)\n\t\t\t\trequire.Equal(t, \"amd64\", idxManifest.Manifests[0].Platform.Architecture)\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\tNewWeightArtifact(\"w1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\t\t\t\tSchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"w1\",\n\t\t\t\t\tTarget: \"/weights/model.bin\", Created: time.Now().UTC(),\n\t\t\t\t}),\n\t\t\t},\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{})\n\t\trequire.NoError(t, err)\n\t})\n\n\tt.Run(\"returns error when image push fails\", func(t *testing.T) {\n\t\tdir := t.TempDir()\n\t\tweightPath := filepath.Join(dir, \"model.bin\")\n\t\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error {\n\t\t\t\treturn errors.New(\"unauthorized: authentication required\")\n\t\t\t},\n\t\t}\n\t\treg := &mockRegistry{}\n\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\tNewWeightArtifact(\"w1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\t\t\t\tSchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"w1\",\n\t\t\t\t\tTarget: \"/weights/model.bin\", Created: time.Now().UTC(),\n\t\t\t\t}),\n\t\t\t},\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{})\n\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"push image\")\n\t\trequire.Contains(t, err.Error(), \"unauthorized\")\n\t})\n\n\tt.Run(\"returns error when get descriptor fails\", func(t *testing.T) {\n\t\tdir := t.TempDir()\n\t\tweightPath := filepath.Join(dir, \"model.bin\")\n\t\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error { return nil },\n\t\t}\n\t\treg := &mockRegistry{\n\t\t\tgetDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) {\n\t\t\t\treturn v1.Descriptor{}, errors.New(\"manifest not found\")\n\t\t\t},\n\t\t}\n\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\tNewWeightArtifact(\"w1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\t\t\t\tSchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"w1\",\n\t\t\t\t\tTarget: \"/weights/model.bin\", Created: time.Now().UTC(),\n\t\t\t\t}),\n\t\t\t},\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{})\n\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"get image descriptor\")\n\t})\n\n\tt.Run(\"returns error when weight push fails\", func(t *testing.T) {\n\t\tdir := t.TempDir()\n\t\tweightPath := filepath.Join(dir, \"model.bin\")\n\t\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error { return nil },\n\t\t}\n\t\treg := &mockRegistry{\n\t\t\tgetDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) {\n\t\t\t\treturn v1.Descriptor{\n\t\t\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\t\t\tSize:      100,\n\t\t\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"abc\"},\n\t\t\t\t}, nil\n\t\t\t},\n\t\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\t\treturn errors.New(\"weight push failed: quota exceeded\")\n\t\t\t},\n\t\t}\n\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\tNewWeightArtifact(\"w1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\t\t\t\tSchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"w1\",\n\t\t\t\t\tTarget: \"/weights/model.bin\", Created: time.Now().UTC(),\n\t\t\t\t}),\n\t\t\t},\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{})\n\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"push weight\")\n\t\trequire.Contains(t, err.Error(), \"w1\")\n\t})\n\n\tt.Run(\"returns error when index push fails\", func(t *testing.T) {\n\t\tdir := t.TempDir()\n\t\tweightPath := filepath.Join(dir, \"model.bin\")\n\t\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error { return nil },\n\t\t}\n\t\treg := &mockRegistry{\n\t\t\tgetDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) {\n\t\t\t\treturn v1.Descriptor{\n\t\t\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\t\t\tSize:      100,\n\t\t\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"abc\"},\n\t\t\t\t}, nil\n\t\t\t},\n\t\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error { return nil },\n\t\t\tpushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\t\t\t\treturn errors.New(\"index push failed\")\n\t\t\t},\n\t\t}\n\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\tNewWeightArtifact(\"w1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\t\t\t\tSchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"w1\",\n\t\t\t\t\tTarget: \"/weights/model.bin\", Created: time.Now().UTC(),\n\t\t\t\t}),\n\t\t\t},\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{})\n\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"push OCI index\")\n\t})\n\n\tt.Run(\"pushes multiple weights concurrently\", func(t *testing.T) {\n\t\tdir := t.TempDir()\n\t\tweight1Path := filepath.Join(dir, \"model1.bin\")\n\t\tweight2Path := filepath.Join(dir, \"model2.bin\")\n\t\trequire.NoError(t, os.WriteFile(weight1Path, []byte(\"weight 1 data\"), 0o644))\n\t\trequire.NoError(t, os.WriteFile(weight2Path, []byte(\"weight 2 data\"), 0o644))\n\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error { return nil },\n\t\t}\n\n\t\t// Use atomic counter — safe for concurrent access from goroutines\n\t\tvar pushedWeightCount atomic.Int32\n\t\treg := &mockRegistry{\n\t\t\tgetDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) {\n\t\t\t\treturn v1.Descriptor{\n\t\t\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\t\t\tSize:      100,\n\t\t\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"abc\"},\n\t\t\t\t}, nil\n\t\t\t},\n\t\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\t\tpushedWeightCount.Add(1)\n\t\t\t\treturn nil\n\t\t\t},\n\t\t\tpushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\t\t\t\tidxManifest, _ := idx.IndexManifest()\n\t\t\t\trequire.Len(t, idxManifest.Manifests, 3) // 1 image + 2 weights\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\n\t\tpusher := NewBundlePusher(docker, reg)\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t\tNewWeightArtifact(\"w1\", v1.Descriptor{\n\t\t\t\t\tDigest: v1.Hash{Algorithm: \"sha256\", Hex: \"aaaa111122223333444455556666777788889999aaaabbbbccccddddeeee0000\"},\n\t\t\t\t}, weight1Path, \"/weights/model1.bin\", WeightConfig{\n\t\t\t\t\tSchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"w1\",\n\t\t\t\t\tTarget: \"/weights/model1.bin\", Created: time.Now().UTC(),\n\t\t\t\t}),\n\t\t\t\tNewWeightArtifact(\"w2\", v1.Descriptor{\n\t\t\t\t\tDigest: v1.Hash{Algorithm: \"sha256\", Hex: \"bbbb111122223333444455556666777788889999aaaabbbbccccddddeeee0000\"},\n\t\t\t\t}, weight2Path, \"/weights/model2.bin\", WeightConfig{\n\t\t\t\t\tSchemaVersion: \"1.0\", CogVersion: \"0.15.0\", Name: \"w2\",\n\t\t\t\t\tTarget: \"/weights/model2.bin\", Created: time.Now().UTC(),\n\t\t\t\t}),\n\t\t\t},\n\t\t}\n\n\t\terr := pusher.Push(context.Background(), m, PushOptions{})\n\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, int32(2), pushedWeightCount.Load()) // both weights pushed (1 tag each)\n\t})\n}\n\n// =============================================================================\n// Resolver.Push tests\n// =============================================================================\n\nfunc TestResolver_Push(t *testing.T) {\n\tt.Run(\"default uses docker push\", func(t *testing.T) {\n\t\tvar dockerPushed bool\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error {\n\t\t\t\tdockerPushed = true\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\t\treg := &mockRegistry{}\n\t\tresolver := NewResolver(docker, reg)\n\n\t\tm := &Model{\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t},\n\t\t}\n\n\t\terr := resolver.Push(context.Background(), m, PushOptions{})\n\t\trequire.NoError(t, err)\n\t\trequire.True(t, dockerPushed, \"standalone should use docker push\")\n\t})\n\n\tt.Run(\"OCIIndex false uses docker push\", func(t *testing.T) {\n\t\tvar dockerPushed bool\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error {\n\t\t\t\tdockerPushed = true\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\t\treg := &mockRegistry{}\n\t\tresolver := NewResolver(docker, reg)\n\n\t\tm := &Model{\n\t\t\t// OCIIndex not set (false by default)\n\t\t\tImage: &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t},\n\t\t}\n\n\t\terr := resolver.Push(context.Background(), m, PushOptions{})\n\t\trequire.NoError(t, err)\n\t\trequire.True(t, dockerPushed, \"default format should use docker push\")\n\t})\n\n\tt.Run(\"OCIIndex true produces an OCI index\", func(t *testing.T) {\n\t\tvar indexPushed bool\n\t\tdocker := &mockDocker{\n\t\t\tpushFunc: func(ctx context.Context, ref string) error { return nil },\n\t\t}\n\t\treg := &mockRegistry{\n\t\t\tgetDescriptorFunc: func(ctx context.Context, ref string) (v1.Descriptor, error) {\n\t\t\t\treturn v1.Descriptor{\n\t\t\t\t\tMediaType: types.OCIManifestSchema1,\n\t\t\t\t\tSize:      100,\n\t\t\t\t\tDigest:    v1.Hash{Algorithm: \"sha256\", Hex: \"abc\"},\n\t\t\t\t}, nil\n\t\t\t},\n\t\t\tpushIndexFunc: func(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\t\t\t\tindexPushed = true\n\t\t\t\treturn nil\n\t\t\t},\n\t\t}\n\t\tresolver := NewResolver(docker, reg)\n\n\t\tm := &Model{\n\t\t\tOCIIndex: true,\n\t\t\tImage:    &ImageArtifact{Reference: \"r8.im/user/model:latest\"},\n\t\t\tArtifacts: []Artifact{\n\t\t\t\t&ImageArtifact{name: \"model\", Reference: \"r8.im/user/model:latest\"},\n\t\t\t},\n\t\t}\n\n\t\terr := resolver.Push(context.Background(), m, PushOptions{})\n\t\trequire.NoError(t, err)\n\t\trequire.True(t, indexPushed, \"OCIIndex=true should push an OCI index\")\n\t})\n\n\tt.Run(\"standalone returns error when image nil\", func(t *testing.T) {\n\t\tdocker := &mockDocker{}\n\t\treg := &mockRegistry{}\n\t\tresolver := NewResolver(docker, reg)\n\n\t\tm := &Model{\n\t\t\tImage:     nil,\n\t\t\tArtifacts: []Artifact{},\n\t\t}\n\n\t\terr := resolver.Push(context.Background(), m, PushOptions{})\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"no image artifact\")\n\t})\n\n\tt.Run(\"OCIIndex true returns error when no image artifact\", func(t *testing.T) {\n\t\tdocker := &mockDocker{}\n\t\treg := &mockRegistry{}\n\t\tresolver := NewResolver(docker, reg)\n\n\t\tm := &Model{\n\t\t\tOCIIndex:  true,\n\t\t\tImage:     nil,\n\t\t\tArtifacts: []Artifact{},\n\t\t}\n\n\t\terr := resolver.Push(context.Background(), m, PushOptions{})\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"no image artifact\")\n\t})\n}\n\n// =============================================================================\n// PushOptions tests\n// =============================================================================\n\nfunc TestPushOptions(t *testing.T) {\n\tt.Run(\"ProjectDir field\", func(t *testing.T) {\n\t\topts := PushOptions{\n\t\t\tProjectDir: \"/path/to/project\",\n\t\t}\n\t\trequire.Equal(t, \"/path/to/project\", opts.ProjectDir)\n\t})\n\n\tt.Run(\"Platform field\", func(t *testing.T) {\n\t\topts := PushOptions{\n\t\t\tPlatform: &Platform{OS: \"linux\", Architecture: \"arm64\"},\n\t\t}\n\t\trequire.Equal(t, \"linux\", opts.Platform.OS)\n\t\trequire.Equal(t, \"arm64\", opts.Platform.Architecture)\n\t})\n}\n\n// =============================================================================\n// repoFromReference tests\n// =============================================================================\n\nfunc TestRepoFromReference(t *testing.T) {\n\ttests := []struct {\n\t\tinput string\n\t\twant  string\n\t}{\n\t\t{\"r8.im/user/model:latest\", \"r8.im/user/model\"},\n\t\t{\"r8.im/user/model:v1.0\", \"r8.im/user/model\"},\n\t\t{\"r8.im/user/model@sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855\", \"r8.im/user/model\"},\n\t\t{\"r8.im/user/model\", \"r8.im/user/model\"},\n\t\t{\"registry.example.com/org/model:tag\", \"registry.example.com/org/model\"},\n\t\t{\"localhost:5000/model:latest\", \"localhost:5000/model\"},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.input, func(t *testing.T) {\n\t\t\tgot := repoFromReference(tt.input)\n\t\t\trequire.Equal(t, tt.want, got)\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "pkg/model/ref.go",
    "content": "package model\n\nimport (\n\t\"fmt\"\n\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n)\n\n// ParseOption configures how image references are parsed.\ntype ParseOption func(*parseOptions)\n\ntype parseOptions struct {\n\tnameOpts []name.Option\n}\n\n// Insecure allows parsing references to registries that use HTTP\n// or have invalid/self-signed certificates.\n// Use this for local development registries like localhost:5000.\nfunc Insecure() ParseOption {\n\treturn func(o *parseOptions) {\n\t\to.nameOpts = append(o.nameOpts, name.Insecure)\n\t}\n}\n\n// WithDefaultRegistry sets the registry to use when the reference\n// doesn't include one. By default, Docker Hub (index.docker.io) is used.\nfunc WithDefaultRegistry(registry string) ParseOption {\n\treturn func(o *parseOptions) {\n\t\to.nameOpts = append(o.nameOpts, name.WithDefaultRegistry(registry))\n\t}\n}\n\n// WithDefaultTag sets the tag to use when the reference doesn't\n// include one. By default, \"latest\" is used.\nfunc WithDefaultTag(tag string) ParseOption {\n\treturn func(o *parseOptions) {\n\t\to.nameOpts = append(o.nameOpts, name.WithDefaultTag(tag))\n\t}\n}\n\n// ParsedRef wraps a validated and parsed image reference.\ntype ParsedRef struct {\n\t// Original is the input string before parsing.\n\tOriginal string\n\n\t// Ref is the underlying parsed reference from go-containerregistry.\n\tRef name.Reference\n}\n\n// ParseRef validates and parses an image reference.\nfunc ParseRef(ref string, opts ...ParseOption) (*ParsedRef, error) {\n\tvar po parseOptions\n\tfor _, opt := range opts {\n\t\topt(&po)\n\t}\n\n\tparsed, err := name.ParseReference(ref, po.nameOpts...)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"invalid image reference %q: %w\", ref, err)\n\t}\n\n\treturn &ParsedRef{\n\t\tOriginal: ref,\n\t\tRef:      parsed,\n\t}, nil\n}\n\n// String returns the fully-qualified canonical reference string.\nfunc (p *ParsedRef) String() string {\n\treturn p.Ref.Name()\n}\n\n// Registry returns the registry host (e.g., \"r8.im\", \"index.docker.io\").\nfunc (p *ParsedRef) Registry() string {\n\treturn p.Ref.Context().RegistryStr()\n}\n\n// Repository returns the repository path (e.g., \"user/model\", \"library/nginx\").\nfunc (p *ParsedRef) Repository() string {\n\treturn p.Ref.Context().RepositoryStr()\n}\n\n// Tag returns the tag (e.g., \"v1\", \"latest\") or empty string if this is a digest reference.\nfunc (p *ParsedRef) Tag() string {\n\tif t, ok := p.Ref.(name.Tag); ok {\n\t\treturn t.TagStr()\n\t}\n\treturn \"\"\n}\n\n// Digest returns the digest (e.g., \"sha256:...\") or empty string if this is a tag reference.\nfunc (p *ParsedRef) Digest() string {\n\tif d, ok := p.Ref.(name.Digest); ok {\n\t\treturn d.DigestStr()\n\t}\n\treturn \"\"\n}\n\n// IsTag returns true if the reference includes a tag.\nfunc (p *ParsedRef) IsTag() bool {\n\t_, ok := p.Ref.(name.Tag)\n\treturn ok\n}\n\n// IsDigest returns true if the reference includes a digest.\nfunc (p *ParsedRef) IsDigest() bool {\n\t_, ok := p.Ref.(name.Digest)\n\treturn ok\n}\n\n// IsReplicate returns true if the registry is the Replicate registry (r8.im).\nfunc (p *ParsedRef) IsReplicate() bool {\n\treturn p.Registry() == global.ReplicateRegistryHost\n}\n"
  },
  {
    "path": "pkg/model/ref_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestParseRef(t *testing.T) {\n\ttests := []struct {\n\t\tname          string\n\t\tref           string\n\t\topts          []ParseOption\n\t\twantRegistry  string\n\t\twantRepo      string\n\t\twantTag       string\n\t\twantDigest    string\n\t\twantReplicate bool\n\t\twantErr       bool\n\t\terrContains   string\n\t}{\n\t\t{\n\t\t\tname:          \"basic tag\",\n\t\t\tref:           \"nginx:latest\",\n\t\t\twantRegistry:  \"index.docker.io\",\n\t\t\twantRepo:      \"library/nginx\",\n\t\t\twantTag:       \"latest\",\n\t\t\twantReplicate: false,\n\t\t},\n\t\t{\n\t\t\tname:          \"implicit latest tag\",\n\t\t\tref:           \"nginx\",\n\t\t\twantRegistry:  \"index.docker.io\",\n\t\t\twantRepo:      \"library/nginx\",\n\t\t\twantTag:       \"latest\",\n\t\t\twantReplicate: false,\n\t\t},\n\t\t{\n\t\t\tname:          \"replicate registry\",\n\t\t\tref:           \"r8.im/user/model:v1\",\n\t\t\twantRegistry:  \"r8.im\",\n\t\t\twantRepo:      \"user/model\",\n\t\t\twantTag:       \"v1\",\n\t\t\twantReplicate: true,\n\t\t},\n\t\t{\n\t\t\tname:          \"replicate registry implicit latest\",\n\t\t\tref:           \"r8.im/user/model\",\n\t\t\twantRegistry:  \"r8.im\",\n\t\t\twantRepo:      \"user/model\",\n\t\t\twantTag:       \"latest\",\n\t\t\twantReplicate: true,\n\t\t},\n\t\t{\n\t\t\tname:          \"non-replicate registry\",\n\t\t\tref:           \"ghcr.io/foo/bar:v1\",\n\t\t\twantRegistry:  \"ghcr.io\",\n\t\t\twantRepo:      \"foo/bar\",\n\t\t\twantTag:       \"v1\",\n\t\t\twantReplicate: false,\n\t\t},\n\t\t{\n\t\t\tname:          \"digest reference\",\n\t\t\tref:           \"nginx@sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef\",\n\t\t\twantRegistry:  \"index.docker.io\",\n\t\t\twantRepo:      \"library/nginx\",\n\t\t\twantTag:       \"\",\n\t\t\twantDigest:    \"sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef\",\n\t\t\twantReplicate: false,\n\t\t},\n\t\t{\n\t\t\tname:          \"replicate with digest\",\n\t\t\tref:           \"r8.im/user/model@sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef\",\n\t\t\twantRegistry:  \"r8.im\",\n\t\t\twantRepo:      \"user/model\",\n\t\t\twantTag:       \"\",\n\t\t\twantDigest:    \"sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef\",\n\t\t\twantReplicate: true,\n\t\t},\n\t\t{\n\t\t\tname:          \"custom registry with port\",\n\t\t\tref:           \"localhost:5000/myimage:test\",\n\t\t\topts:          []ParseOption{Insecure()},\n\t\t\twantRegistry:  \"localhost:5000\",\n\t\t\twantRepo:      \"myimage\",\n\t\t\twantTag:       \"test\",\n\t\t\twantReplicate: false,\n\t\t},\n\t\t{\n\t\t\tname:          \"with default registry option\",\n\t\t\tref:           \"user/model:v1\",\n\t\t\topts:          []ParseOption{WithDefaultRegistry(\"r8.im\")},\n\t\t\twantRegistry:  \"r8.im\",\n\t\t\twantRepo:      \"user/model\",\n\t\t\twantTag:       \"v1\",\n\t\t\twantReplicate: true,\n\t\t},\n\t\t{\n\t\t\tname:          \"with default tag option\",\n\t\t\tref:           \"nginx\",\n\t\t\topts:          []ParseOption{WithDefaultTag(\"stable\")},\n\t\t\twantRegistry:  \"index.docker.io\",\n\t\t\twantRepo:      \"library/nginx\",\n\t\t\twantTag:       \"stable\",\n\t\t\twantReplicate: false,\n\t\t},\n\t\t{\n\t\t\tname:        \"invalid reference\",\n\t\t\tref:         \":::invalid\",\n\t\t\twantErr:     true,\n\t\t\terrContains: `invalid image reference \":::invalid\"`,\n\t\t},\n\t\t{\n\t\t\tname:        \"empty reference\",\n\t\t\tref:         \"\",\n\t\t\twantErr:     true,\n\t\t\terrContains: `invalid image reference \"\"`,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tparsed, err := ParseRef(tt.ref, tt.opts...)\n\n\t\t\tif tt.wantErr {\n\t\t\t\trequire.Error(t, err)\n\t\t\t\tif tt.errContains != \"\" {\n\t\t\t\t\trequire.Contains(t, err.Error(), tt.errContains)\n\t\t\t\t}\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\trequire.NoError(t, err)\n\t\t\trequire.NotNil(t, parsed)\n\n\t\t\trequire.Equal(t, tt.ref, parsed.Original, \"Original should preserve input\")\n\t\t\trequire.Equal(t, tt.wantRegistry, parsed.Registry(), \"Registry mismatch\")\n\t\t\trequire.Equal(t, tt.wantRepo, parsed.Repository(), \"Repository mismatch\")\n\t\t\trequire.Equal(t, tt.wantTag, parsed.Tag(), \"Tag mismatch\")\n\t\t\trequire.Equal(t, tt.wantDigest, parsed.Digest(), \"Digest mismatch\")\n\t\t\trequire.Equal(t, tt.wantReplicate, parsed.IsReplicate(), \"IsReplicate mismatch\")\n\t\t})\n\t}\n}\n\nfunc TestParsedRef_IsTag(t *testing.T) {\n\ttagRef, err := ParseRef(\"nginx:latest\")\n\trequire.NoError(t, err)\n\trequire.True(t, tagRef.IsTag())\n\trequire.False(t, tagRef.IsDigest())\n\n\tdigestRef, err := ParseRef(\"nginx@sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef\")\n\trequire.NoError(t, err)\n\trequire.False(t, digestRef.IsTag())\n\trequire.True(t, digestRef.IsDigest())\n}\n\nfunc TestParsedRef_String(t *testing.T) {\n\ttests := []struct {\n\t\tname    string\n\t\tref     string\n\t\twantStr string\n\t}{\n\t\t{\n\t\t\tname:    \"bare image gets fully qualified\",\n\t\t\tref:     \"nginx\",\n\t\t\twantStr: \"index.docker.io/library/nginx:latest\",\n\t\t},\n\t\t{\n\t\t\tname:    \"replicate ref\",\n\t\t\tref:     \"r8.im/user/model:v1\",\n\t\t\twantStr: \"r8.im/user/model:v1\",\n\t\t},\n\t\t{\n\t\t\tname:    \"digest ref\",\n\t\t\tref:     \"nginx@sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef\",\n\t\t\twantStr: \"index.docker.io/library/nginx@sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tparsed, err := ParseRef(tt.ref)\n\t\t\trequire.NoError(t, err)\n\t\t\trequire.Equal(t, tt.wantStr, parsed.String())\n\t\t})\n\t}\n}\n\nfunc TestParseOptions(t *testing.T) {\n\tt.Run(\"multiple options can be combined\", func(t *testing.T) {\n\t\tparsed, err := ParseRef(\"myimage\",\n\t\t\tWithDefaultRegistry(\"r8.im\"),\n\t\t\tWithDefaultTag(\"v2\"),\n\t\t)\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"r8.im\", parsed.Registry())\n\t\trequire.Equal(t, \"v2\", parsed.Tag())\n\t\trequire.True(t, parsed.IsReplicate())\n\t})\n\n\tt.Run(\"insecure allows localhost registries\", func(t *testing.T) {\n\t\tparsed, err := ParseRef(\"localhost:5000/test:v1\", Insecure())\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"localhost:5000\", parsed.Registry())\n\t\trequire.Equal(t, \"test\", parsed.Repository())\n\t\trequire.Equal(t, \"v1\", parsed.Tag())\n\t})\n}\n"
  },
  {
    "path": "pkg/model/ref_types.go",
    "content": "package model\n\nimport \"context\"\n\n// Ref represents something that can be resolved to a Model.\n// This interface enables declarative model resolution - callers specify\n// \"what they have\" (a tag, local image, or source to build) and the\n// Resolver figures out how to produce a Model.\ntype Ref interface {\n\t// resolve is unexported to keep the interface internal.\n\t// External code uses Resolver.Resolve() instead of calling this directly.\n\tresolve(ctx context.Context, r *Resolver) (*Model, error)\n}\n\n// Resolve resolves any Ref to a Model.\n// This is the main entry point for declarative model resolution.\nfunc (r *Resolver) Resolve(ctx context.Context, ref Ref) (*Model, error) {\n\treturn ref.resolve(ctx, r)\n}\n\n// =============================================================================\n// TagRef - resolves an image by tag/digest, trying local then remote\n// =============================================================================\n\n// TagRef resolves an image by tag or digest reference.\n// It uses the default Load behavior: try remote registry first,\n// then fall back to local docker daemon if not found remotely.\ntype TagRef struct {\n\tParsed *ParsedRef\n}\n\n// FromTag parses and validates a tag reference.\n// Returns an error immediately if the reference is invalid.\nfunc FromTag(ref string) (*TagRef, error) {\n\tparsed, err := ParseRef(ref)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn &TagRef{Parsed: parsed}, nil\n}\n\nfunc (t *TagRef) resolve(ctx context.Context, r *Resolver) (*Model, error) {\n\t// Use default Inspect behavior (PreferRemote)\n\treturn r.Inspect(ctx, t.Parsed)\n}\n\n// =============================================================================\n// LocalRef - explicitly loads from docker daemon only\n// =============================================================================\n\n// LocalRef resolves an image from the local docker daemon only.\n// It will not fall back to remote registry if the image is not found locally.\ntype LocalRef struct {\n\tParsed *ParsedRef\n}\n\n// FromLocal parses and validates a reference for local resolution.\n// Returns an error immediately if the reference is invalid.\nfunc FromLocal(ref string) (*LocalRef, error) {\n\tparsed, err := ParseRef(ref)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn &LocalRef{Parsed: parsed}, nil\n}\n\nfunc (l *LocalRef) resolve(ctx context.Context, r *Resolver) (*Model, error) {\n\treturn r.Inspect(ctx, l.Parsed, LocalOnly())\n}\n\n// =============================================================================\n// RemoteRef - explicitly loads from registry only\n// =============================================================================\n\n// RemoteRef resolves an image from a remote registry only.\n// It will not check the local docker daemon.\ntype RemoteRef struct {\n\tParsed *ParsedRef\n}\n\n// FromRemote parses and validates a reference for remote resolution.\n// Returns an error immediately if the reference is invalid.\nfunc FromRemote(ref string) (*RemoteRef, error) {\n\tparsed, err := ParseRef(ref)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn &RemoteRef{Parsed: parsed}, nil\n}\n\nfunc (rr *RemoteRef) resolve(ctx context.Context, r *Resolver) (*Model, error) {\n\treturn r.Inspect(ctx, rr.Parsed, RemoteOnly())\n}\n\n// =============================================================================\n// BuildRef - creates a Model by building from source\n// =============================================================================\n\n// BuildRef resolves to a Model by building from source.\n// This wraps a Source and BuildOptions for deferred building.\ntype BuildRef struct {\n\tSource  *Source\n\tOptions BuildOptions\n}\n\n// FromBuild creates a BuildRef from source and options.\n// Unlike the other From* functions, this does not validate eagerly -\n// validation happens at build time.\nfunc FromBuild(src *Source, opts BuildOptions) *BuildRef {\n\treturn &BuildRef{Source: src, Options: opts}\n}\n\nfunc (b *BuildRef) resolve(ctx context.Context, r *Resolver) (*Model, error) {\n\treturn r.Build(ctx, b.Source, b.Options)\n}\n"
  },
  {
    "path": "pkg/model/ref_types_test.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"testing\"\n\n\t\"github.com/docker/docker/api/types/image\"\n\tdockerspec \"github.com/moby/docker-image-spec/specs-go/v1\"\n\tocispec \"github.com/opencontainers/image-spec/specs-go/v1\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// =============================================================================\n// FromTag tests\n// =============================================================================\n\nfunc TestFromTag_ValidRef(t *testing.T) {\n\tref, err := FromTag(\"my-image:latest\")\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, ref)\n\trequire.NotNil(t, ref.Parsed)\n\trequire.Equal(t, \"my-image:latest\", ref.Parsed.Original)\n}\n\nfunc TestFromTag_ValidRefWithRegistry(t *testing.T) {\n\tref, err := FromTag(\"r8.im/user/model:v1\")\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, ref)\n\trequire.Equal(t, \"r8.im\", ref.Parsed.Registry())\n\trequire.Equal(t, \"v1\", ref.Parsed.Tag())\n}\n\nfunc TestFromTag_InvalidRef(t *testing.T) {\n\tref, err := FromTag(\"INVALID::REF\")\n\n\trequire.Error(t, err)\n\trequire.Nil(t, ref)\n\trequire.Contains(t, err.Error(), \"invalid image reference\")\n}\n\n// =============================================================================\n// FromLocal tests\n// =============================================================================\n\nfunc TestFromLocal_ValidRef(t *testing.T) {\n\tref, err := FromLocal(\"my-image:latest\")\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, ref)\n\trequire.NotNil(t, ref.Parsed)\n\trequire.Equal(t, \"my-image:latest\", ref.Parsed.Original)\n}\n\nfunc TestFromLocal_InvalidRef(t *testing.T) {\n\tref, err := FromLocal(\"INVALID::REF\")\n\n\trequire.Error(t, err)\n\trequire.Nil(t, ref)\n\trequire.Contains(t, err.Error(), \"invalid image reference\")\n}\n\n// =============================================================================\n// FromRemote tests\n// =============================================================================\n\nfunc TestFromRemote_ValidRef(t *testing.T) {\n\tref, err := FromRemote(\"r8.im/user/model\")\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, ref)\n\trequire.NotNil(t, ref.Parsed)\n\trequire.Equal(t, \"r8.im\", ref.Parsed.Registry())\n}\n\nfunc TestFromRemote_InvalidRef(t *testing.T) {\n\tref, err := FromRemote(\"INVALID::REF\")\n\n\trequire.Error(t, err)\n\trequire.Nil(t, ref)\n\trequire.Contains(t, err.Error(), \"invalid image reference\")\n}\n\n// =============================================================================\n// FromBuild tests\n// =============================================================================\n\nfunc TestFromBuild(t *testing.T) {\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Predict: \"predict.py:Predictor\"},\n\t\tProjectDir: \"/path/to/project\",\n\t}\n\topts := BuildOptions{\n\t\tImageName: \"my-built-image:latest\",\n\t\tNoCache:   true,\n\t}\n\n\tref := FromBuild(src, opts)\n\n\trequire.NotNil(t, ref)\n\trequire.Same(t, src, ref.Source)\n\trequire.Equal(t, \"my-built-image:latest\", ref.Options.ImageName)\n\trequire.True(t, ref.Options.NoCache)\n}\n\nfunc TestFromBuild_NilSource(t *testing.T) {\n\t// FromBuild should accept nil source - validation happens at resolve time\n\tref := FromBuild(nil, BuildOptions{ImageName: \"test\"})\n\n\trequire.NotNil(t, ref)\n\trequire.Nil(t, ref.Source)\n}\n\n// =============================================================================\n// TagRef.resolve tests\n// =============================================================================\n\nfunc TestTagRef_Resolve_Success(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := FromTag(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Resolve(context.Background(), ref)\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.Equal(t, \"0.10.0\", model.CogVersion)\n}\n\nfunc TestTagRef_Resolve_FallsBackToLocal(t *testing.T) {\n\tlocalCalled := false\n\tremoteCalled := false\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tlocalCalled = true\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:local123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tremoteCalled = true\n\t\t\treturn nil, registry.NotFoundError\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := FromTag(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Resolve(context.Background(), ref)\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.True(t, remoteCalled, \"TagRef should try remote first\")\n\trequire.True(t, localCalled, \"TagRef should fall back to local\")\n\trequire.Equal(t, ImageSourceLocal, model.Image.Source)\n}\n\n// =============================================================================\n// LocalRef.resolve tests\n// =============================================================================\n\nfunc TestLocalRef_Resolve_Success(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:local123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.9.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := FromLocal(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Resolve(context.Background(), ref)\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.Equal(t, ImageSourceLocal, model.Image.Source)\n\trequire.Equal(t, \"0.9.0\", model.CogVersion)\n}\n\nfunc TestLocalRef_Resolve_NotFound(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"No such image: my-image:latest\")\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tt.Fatal(\"LocalRef should not fall back to remote\")\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := FromLocal(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Resolve(context.Background(), ref)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"not found locally\")\n}\n\n// =============================================================================\n// RemoteRef.resolve tests\n// =============================================================================\n\nfunc TestRemoteRef_Resolve_Success(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tt.Fatal(\"RemoteRef should not check local docker\")\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\treturn ®istry.ManifestResult{\n\t\t\t\tSchemaVersion: 2,\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := FromRemote(\"r8.im/user/model\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Resolve(context.Background(), ref)\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.Equal(t, ImageSourceRemote, model.Image.Source)\n\trequire.Equal(t, \"0.10.0\", model.CogVersion)\n}\n\nfunc TestRemoteRef_Resolve_NotFound(t *testing.T) {\n\tdocker := &mockDocker{}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\treturn nil, registry.NotFoundError\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := FromRemote(\"r8.im/user/model\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Resolve(context.Background(), ref)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"not found in registry\")\n}\n\n// =============================================================================\n// BuildRef.resolve tests\n// =============================================================================\n\nfunc TestBuildRef_Resolve_Success(t *testing.T) {\n\tbuildCalled := false\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelVersion: \"0.11.0\",\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"gpu\":true}}`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tfactory := &mockFactory{\n\t\tname: \"test\",\n\t\tbuildFunc: func(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error) {\n\t\t\tbuildCalled = true\n\t\t\trequire.Equal(t, \"my-built-image\", opts.ImageName)\n\t\t\treturn &ImageArtifact{Reference: opts.ImageName, Source: ImageSourceBuild}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{}).WithFactory(factory)\n\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Predict: \"predict.py:Predictor\"},\n\t\tProjectDir: \"/tmp/test\",\n\t}\n\tref := FromBuild(src, BuildOptions{ImageName: \"my-built-image\"})\n\n\tmodel, err := resolver.Resolve(context.Background(), ref)\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.True(t, buildCalled, \"BuildRef should call factory.Build\")\n\trequire.Equal(t, \"0.11.0\", model.CogVersion)\n}\n\nfunc TestBuildRef_Resolve_BuildError(t *testing.T) {\n\tfactory := &mockFactory{\n\t\tname: \"test\",\n\t\tbuildFunc: func(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error) {\n\t\t\treturn nil, errors.New(\"build failed: missing dependencies\")\n\t\t},\n\t}\n\n\tresolver := NewResolver(&mockDocker{}, &mockRegistry{}).WithFactory(factory)\n\n\tsrc := &Source{\n\t\tConfig:     &config.Config{},\n\t\tProjectDir: \"/tmp/test\",\n\t}\n\tref := FromBuild(src, BuildOptions{ImageName: \"my-image\"})\n\n\t_, err := resolver.Resolve(context.Background(), ref)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"build failed\")\n}\n\n// =============================================================================\n// Resolver.Resolve dispatch tests\n// =============================================================================\n\nfunc TestResolver_Resolve_DispatchesCorrectly(t *testing.T) {\n\t// This test verifies that Resolver.Resolve correctly dispatches to each Ref type\n\ttests := []struct {\n\t\tname        string\n\t\tref         Ref\n\t\texpectLocal bool\n\t}{\n\t\t{\n\t\t\tname: \"TagRef dispatches to Load (default behavior)\",\n\t\t\tref: func() Ref {\n\t\t\t\tr, _ := FromTag(\"my-image:latest\")\n\t\t\t\treturn r\n\t\t\t}(),\n\t\t\texpectLocal: true, // TagRef tries remote first, falls back to local\n\t\t},\n\t\t{\n\t\t\tname: \"LocalRef dispatches to local only\",\n\t\t\tref: func() Ref {\n\t\t\t\tr, _ := FromLocal(\"my-image:latest\")\n\t\t\t\treturn r\n\t\t\t}(),\n\t\t\texpectLocal: true,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tlocalCalled := false\n\t\t\tdocker := &mockDocker{\n\t\t\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\t\t\tlocalCalled = true\n\t\t\t\t\treturn &image.InspectResponse{\n\t\t\t\t\t\tID: \"sha256:test\",\n\t\t\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\t\t\tLabels: map[string]string{\n\t\t\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t\t\t},\n\t\t\t\t\t\t\t},\n\t\t\t\t\t\t},\n\t\t\t\t\t}, nil\n\t\t\t\t},\n\t\t\t}\n\n\t\t\tresolver := NewResolver(docker, &mockRegistry{})\n\t\t\t_, err := resolver.Resolve(context.Background(), tt.ref)\n\n\t\t\trequire.NoError(t, err)\n\t\t\trequire.Equal(t, tt.expectLocal, localCalled)\n\t\t})\n\t}\n}\n\n// =============================================================================\n// Ref interface compile-time checks\n// =============================================================================\n\n// Compile-time check that all types implement Ref interface\nvar (\n\t_ Ref = (*TagRef)(nil)\n\t_ Ref = (*LocalRef)(nil)\n\t_ Ref = (*RemoteRef)(nil)\n\t_ Ref = (*BuildRef)(nil)\n)\n"
  },
  {
    "path": "pkg/model/resolver.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"path/filepath\"\n\t\"strings\"\n\n\t\"github.com/docker/docker/api/types/image\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// Option configures how Resolver methods behave.\ntype Option func(*options)\n\ntype options struct {\n\tlocalOnly   bool\n\tremoteOnly  bool\n\tpreferLocal bool // default: true\n\tplatform    *registry.Platform\n}\n\nfunc defaultOptions() *options {\n\treturn &options{} // Default: preferRemote (try registry first, fall back to local)\n}\n\n// LocalOnly loads only from the local docker daemon.\n// Returns an error if the image is not found locally.\nfunc LocalOnly() Option {\n\treturn func(o *options) {\n\t\to.localOnly = true\n\t\to.remoteOnly = false\n\t\to.preferLocal = false\n\t}\n}\n\n// RemoteOnly loads only from the remote registry.\n// Does not check the local docker daemon.\nfunc RemoteOnly() Option {\n\treturn func(o *options) {\n\t\to.remoteOnly = true\n\t\to.localOnly = false\n\t\to.preferLocal = false\n\t}\n}\n\n// PreferLocal tries local docker daemon first, falls back to remote on not-found.\nfunc PreferLocal() Option {\n\treturn func(o *options) {\n\t\to.preferLocal = true\n\t\to.localOnly = false\n\t\to.remoteOnly = false\n\t}\n}\n\n// PreferRemote tries remote registry first, falls back to local on not-found.\n// This is the default behavior.\nfunc PreferRemote() Option {\n\treturn func(o *options) {\n\t\to.preferLocal = false\n\t\to.localOnly = false\n\t\to.remoteOnly = false\n\t}\n}\n\n// WithPlatform sets the platform for remote registry queries.\nfunc WithPlatform(p *registry.Platform) Option {\n\treturn func(o *options) {\n\t\to.platform = p\n\t}\n}\n\n// Resolver orchestrates building and loading Models.\ntype Resolver struct {\n\tdocker      command.Command\n\tregistry    registry.Client\n\tfactory     Factory\n\timagePusher *ImagePusher\n}\n\n// NewResolver creates a Resolver with the default factory.\nfunc NewResolver(docker command.Command, reg registry.Client) *Resolver {\n\treturn &Resolver{\n\t\tdocker:      docker,\n\t\tregistry:    reg,\n\t\tfactory:     defaultFactory(docker, reg),\n\t\timagePusher: newImagePusher(docker, reg),\n\t}\n}\n\n// WithFactory sets a custom factory and returns the Resolver for chaining.\nfunc (r *Resolver) WithFactory(f Factory) *Resolver {\n\tr.factory = f\n\treturn r\n}\n\n// Inspect returns Model metadata for a parsed ref without pulling.\n// By default (PreferLocal), tries local docker daemon first, then remote registry.\n// Only falls back on \"not found\" errors; real errors (docker down, auth) are surfaced.\n// Returns ErrNotCogModel if the image is not a valid Cog model.\nfunc (r *Resolver) Inspect(ctx context.Context, ref *ParsedRef, opts ...Option) (*Model, error) {\n\to := defaultOptions()\n\tfor _, opt := range opts {\n\t\topt(o)\n\t}\n\n\tswitch {\n\tcase o.localOnly:\n\t\treturn r.loadLocal(ctx, ref)\n\tcase o.remoteOnly:\n\t\treturn r.loadRemote(ctx, ref, o.platform)\n\tcase o.preferLocal:\n\t\tmodel, localErr := r.loadLocal(ctx, ref)\n\t\tif localErr == nil {\n\t\t\treturn model, nil\n\t\t}\n\t\t// Check the underlying error before the wrapper adds \"not found\" text\n\t\tif !isNotFoundError(errors.Unwrap(localErr)) {\n\t\t\treturn nil, localErr // Real error, don't mask it\n\t\t}\n\t\treturn r.loadRemote(ctx, ref, o.platform)\n\tdefault:\n\t\t// PreferRemote\n\t\tmodel, remoteErr := r.loadRemote(ctx, ref, o.platform)\n\t\tif remoteErr == nil {\n\t\t\treturn model, nil\n\t\t}\n\t\t// Check the underlying error before the wrapper adds \"not found\" text\n\t\tif !isNotFoundError(errors.Unwrap(remoteErr)) {\n\t\t\treturn nil, remoteErr\n\t\t}\n\t\treturn r.loadLocal(ctx, ref)\n\t}\n}\n\n// InspectByID returns Model metadata from the local docker daemon by image ID.\n// This supports both full IDs (sha256:...) and short IDs (e.g., \"9056219a5fb2\").\n// Use this when you have an image ID rather than a tagged reference.\n// Returns ErrNotCogModel if the image is not a valid Cog model.\nfunc (r *Resolver) InspectByID(ctx context.Context, id string) (*Model, error) {\n\tresp, err := r.docker.Inspect(ctx, id)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to load image by ID %s: %w\", id, err)\n\t}\n\n\t// Use the canonical ID from the response as the reference\n\timg := &ImageArtifact{\n\t\tReference: resp.ID,\n\t\tDigest:    resp.ID,\n\t\tLabels:    resp.Config.Labels,\n\t\tSource:    ImageSourceLocal,\n\t}\n\n\tmodel, err := img.ToModel()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"image %s: %w\", id, err)\n\t}\n\treturn model, nil\n}\n\n// Pull ensures a Model is locally available for running.\n// It first checks if the image exists locally. If not, it pulls from the registry.\n// Returns ErrNotCogModel if the image is not a valid Cog model.\n// Returns ErrNotFound if the image cannot be found locally or remotely.\nfunc (r *Resolver) Pull(ctx context.Context, ref *ParsedRef, opts ...Option) (*Model, error) {\n\to := defaultOptions()\n\tfor _, opt := range opts {\n\t\topt(o)\n\t}\n\n\t// First, try to inspect locally\n\tmodel, err := r.loadLocal(ctx, ref)\n\tif err == nil {\n\t\treturn model, nil\n\t}\n\n\t// If local-only mode, don't try to pull\n\tif o.localOnly {\n\t\treturn nil, fmt.Errorf(\"image %s: %w\", ref.Original, ErrNotFound)\n\t}\n\n\t// If local image exists but isn't a Cog model, don't try to pull\n\t// (pulling won't change the existing image)\n\tif errors.Is(err, ErrNotCogModel) {\n\t\treturn nil, err\n\t}\n\n\t// Check if it's a \"not found\" error (safe to try pull)\n\tif !isNotFoundError(errors.Unwrap(err)) {\n\t\t// Real error (connection refused, etc.) - don't mask it\n\t\treturn nil, err\n\t}\n\n\t// Pull the image\n\t// TODO: Support platform option for multi-platform images\n\t_, err = r.docker.Pull(ctx, ref.String(), false)\n\tif err != nil {\n\t\tif isNotFoundError(err) {\n\t\t\treturn nil, fmt.Errorf(\"image %s: %w\", ref.Original, ErrNotFound)\n\t\t}\n\t\treturn nil, fmt.Errorf(\"failed to pull image %s: %w\", ref.Original, err)\n\t}\n\n\t// Inspect the now-local image\n\treturn r.loadLocal(ctx, ref)\n}\n\n// Build creates a Model by building from source.\nfunc (r *Resolver) Build(ctx context.Context, src *Source, opts BuildOptions) (*Model, error) {\n\tif src == nil {\n\t\treturn nil, fmt.Errorf(\"source is required for Build\")\n\t}\n\tif src.Config == nil {\n\t\treturn nil, fmt.Errorf(\"source.Config is required for Build\")\n\t}\n\tif src.ProjectDir == \"\" {\n\t\treturn nil, fmt.Errorf(\"source.ProjectDir is required for Build\")\n\t}\n\topts = opts.WithDefaults(src)\n\n\t// Build image artifact via ImageBuilder\n\tib := NewImageBuilder(r.factory, r.docker, src, opts)\n\timageSpec := NewImageSpec(\"model\", opts.ImageName)\n\timgResult, err := ib.Build(ctx, imageSpec)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tia, ok := imgResult.(*ImageArtifact)\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"unexpected artifact type from image builder: %T\", imgResult)\n\t}\n\n\tm, err := r.modelFromImage(ia, src.Config)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tm.OCIIndex = opts.OCIIndex\n\tm.Artifacts = []Artifact{ia}\n\n\t// Build weight artifacts if OCI index mode is enabled\n\tlockPath := opts.WeightsLockPath\n\tif lockPath == \"\" {\n\t\tlockPath = filepath.Join(src.ProjectDir, WeightsLockFilename)\n\t}\n\n\tif opts.OCIIndex && len(src.Config.Weights) > 0 {\n\t\twb := NewWeightBuilder(src, m.CogVersion, lockPath)\n\t\tfor _, ws := range src.Config.Weights {\n\t\t\tspec := NewWeightSpec(ws.Name, ws.Source, ws.Target)\n\t\t\tartifact, buildErr := wb.Build(ctx, spec)\n\t\t\tif buildErr != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"build weight %q: %w\", ws.Name, buildErr)\n\t\t\t}\n\t\t\tm.Artifacts = append(m.Artifacts, artifact)\n\t\t}\n\n\t}\n\n\treturn m, nil\n}\n\n// Push pushes a Model to a container registry.\n//\n// Uses the OCI chunked push path (via ImagePusher) which bypasses Docker's\n// monolithic push and supports layers of any size through chunked uploads.\n// Falls back to legacy Docker push if OCI push is not available.\nfunc (r *Resolver) Push(ctx context.Context, m *Model, opts PushOptions) error {\n\tif m.OCIIndex {\n\t\tpusher := NewBundlePusher(r.docker, r.registry)\n\t\treturn pusher.Push(ctx, m, opts)\n\t}\n\n\timgArtifact := m.GetImageArtifact()\n\tif imgArtifact == nil {\n\t\treturn fmt.Errorf(\"no image artifact in model\")\n\t}\n\n\tvar imagePushOpts []ImagePushOption\n\tif opts.ImageProgressFn != nil {\n\t\timagePushOpts = append(imagePushOpts, WithProgressFn(opts.ImageProgressFn))\n\t}\n\tif opts.OnFallback != nil {\n\t\timagePushOpts = append(imagePushOpts, WithOnFallback(opts.OnFallback))\n\t}\n\treturn r.imagePusher.Push(ctx, imgArtifact, imagePushOpts...)\n}\n\n// loadLocal loads a Model from the local docker daemon.\nfunc (r *Resolver) loadLocal(ctx context.Context, ref *ParsedRef) (*Model, error) {\n\tresp, err := r.docker.Inspect(ctx, ref.String())\n\tif err != nil {\n\t\tif isNotFoundError(err) {\n\t\t\treturn nil, fmt.Errorf(\"image %s not found locally: %w\", ref.Original, err)\n\t\t}\n\t\treturn nil, fmt.Errorf(\"failed to inspect local image %s: %w\", ref.Original, err)\n\t}\n\treturn r.modelFromInspect(ref, resp, ImageSourceLocal)\n}\n\n// loadRemote loads a Model from the remote registry.\nfunc (r *Resolver) loadRemote(ctx context.Context, ref *ParsedRef, platform *registry.Platform) (*Model, error) {\n\tmanifest, err := r.registry.Inspect(ctx, ref.String(), platform)\n\tif err != nil {\n\t\tif errors.Is(err, registry.NotFoundError) {\n\t\t\treturn nil, fmt.Errorf(\"image %s not found in registry: %w\", ref.Original, err)\n\t\t}\n\t\treturn nil, fmt.Errorf(\"failed to inspect remote image %s: %w\", ref.Original, err)\n\t}\n\treturn r.modelFromManifest(ref, manifest, ImageSourceRemote)\n}\n\n// modelFromImage creates a Model from ImageArtifact with a known config (post-build).\n// Uses the provided config rather than parsing from labels.\nfunc (r *Resolver) modelFromImage(img *ImageArtifact, cfg *config.Config) (*Model, error) {\n\tschema, err := img.ParsedOpenAPISchema()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to parse schema from image labels: %w\", err)\n\t}\n\n\treturn &Model{\n\t\tImage:      img,\n\t\tConfig:     cfg,\n\t\tSchema:     schema,\n\t\tCogVersion: img.CogVersion(),\n\t}, nil\n}\n\n// modelFromInspect creates a Model from docker inspect response.\n// Returns ErrNotCogModel if the image is not a valid Cog model.\nfunc (r *Resolver) modelFromInspect(ref *ParsedRef, resp *image.InspectResponse, source ImageSource) (*Model, error) {\n\timg := &ImageArtifact{\n\t\tReference: ref.String(),\n\t\tDigest:    resp.ID,\n\t\tLabels:    resp.Config.Labels,\n\t\tSource:    source,\n\t}\n\n\tmodel, err := img.ToModel()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"image %s: %w\", ref.Original, err)\n\t}\n\treturn model, nil\n}\n\n// modelFromManifest creates a Model from registry manifest.\n// Returns ErrNotCogModel if the image is not a valid Cog model.\nfunc (r *Resolver) modelFromManifest(ref *ParsedRef, manifest *registry.ManifestResult, source ImageSource) (*Model, error) {\n\t// Check if this is an OCI Index (v2 format)\n\tif isOCIIndex(manifest) {\n\t\treturn r.modelFromIndex(ref, manifest, source)\n\t}\n\n\t// Standard image (v1 format)\n\timg := &ImageArtifact{\n\t\tReference: ref.String(),\n\t\tDigest:    manifest.Config, // Config digest serves as image ID\n\t\tLabels:    manifest.Labels,\n\t\tSource:    source,\n\t}\n\n\tm, err := img.ToModel()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"image %s: %w\", ref.Original, err)\n\t}\n\treturn m, nil\n}\n\n// modelFromIndex creates a Model from an OCI Image Index.\n// It extracts the image manifest and weights manifest from the index.\nfunc (r *Resolver) modelFromIndex(ref *ParsedRef, manifest *registry.ManifestResult, source ImageSource) (*Model, error) {\n\t// Find the image manifest (skip unknown/unknown platform artifacts)\n\timgManifest := findImageManifest(manifest.Manifests, nil)\n\tif imgManifest == nil {\n\t\treturn nil, fmt.Errorf(\"no image manifest found in index %s\", ref.Original)\n\t}\n\n\t// Create ImageArtifact from the image manifest\n\timg := &ImageArtifact{\n\t\tReference: ref.String(),\n\t\tDigest:    imgManifest.Digest,\n\t\tLabels:    manifest.Labels, // Labels come from the index inspection\n\t\tSource:    source,\n\t\tPlatform: &Platform{\n\t\t\tOS:           imgManifest.OS,\n\t\t\tArchitecture: imgManifest.Architecture,\n\t\t\tVariant:      imgManifest.Variant,\n\t\t},\n\t}\n\n\tm, err := img.ToModel()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"image %s: %w\", ref.Original, err)\n\t}\n\n\tm.Index = &Index{\n\t\tDigest:    manifest.Digest, // Content-addressable digest from registry\n\t\tReference: ref.String(),\n\t\tMediaType: manifest.MediaType,\n\t\tManifests: make([]IndexManifest, len(manifest.Manifests)),\n\t}\n\n\t// Populate index manifests\n\tfor i, pm := range manifest.Manifests {\n\t\tim := IndexManifest{\n\t\t\tDigest:      pm.Digest,\n\t\t\tMediaType:   pm.MediaType,\n\t\t\tSize:        pm.Size,\n\t\t\tAnnotations: pm.Annotations,\n\t\t}\n\t\tif pm.OS != \"\" {\n\t\t\tim.Platform = &Platform{\n\t\t\t\tOS:           pm.OS,\n\t\t\t\tArchitecture: pm.Architecture,\n\t\t\t\tVariant:      pm.Variant,\n\t\t\t}\n\t\t}\n\t\t// Determine manifest type\n\t\tif pm.OS == PlatformUnknown && pm.Annotations != nil && pm.Annotations[AnnotationReferenceType] == AnnotationValueWeights {\n\t\t\tim.Type = ManifestTypeWeights\n\t\t} else {\n\t\t\tim.Type = ManifestTypeImage\n\t\t}\n\t\tm.Index.Manifests[i] = im\n\t}\n\n\treturn m, nil\n}\n\n// isOCIIndex checks if the manifest result is an OCI Image Index.\nfunc isOCIIndex(mr *registry.ManifestResult) bool {\n\treturn mr.IsIndex()\n}\n\n// findWeightsManifest finds the weights manifest in an index.\n// Returns nil if no weights manifest is found.\nfunc findWeightsManifest(manifests []registry.PlatformManifest) *registry.PlatformManifest {\n\tfor i := range manifests {\n\t\tm := &manifests[i]\n\t\tif m.Annotations != nil && m.Annotations[AnnotationReferenceType] == AnnotationValueWeights {\n\t\t\treturn m\n\t\t}\n\t}\n\treturn nil\n}\n\n// findImageManifest finds the model image manifest in an index.\n// If platform is specified, matches on OS/Architecture.\n// Skips artifacts (platform: unknown/unknown).\nfunc findImageManifest(manifests []registry.PlatformManifest, platform *registry.Platform) *registry.PlatformManifest {\n\tfor i := range manifests {\n\t\tm := &manifests[i]\n\t\t// Skip artifacts (unknown platform)\n\t\tif m.OS == PlatformUnknown {\n\t\t\tcontinue\n\t\t}\n\t\t// Match platform if specified\n\t\tif platform != nil {\n\t\t\tif m.OS != platform.OS || m.Architecture != platform.Architecture {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t}\n\t\treturn m\n\t}\n\treturn nil\n}\n\n// isNotFoundError checks if an error indicates \"not found\" vs a real error.\n// Only \"not found\" errors should trigger fallback to alternative source.\nfunc isNotFoundError(err error) bool {\n\tif err == nil {\n\t\treturn false\n\t}\n\n\t// Don't treat context errors as \"not found\"\n\tif errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {\n\t\treturn false\n\t}\n\n\t// Check for registry NotFoundError\n\tif errors.Is(err, registry.NotFoundError) {\n\t\treturn true\n\t}\n\n\t// Check for common not-found patterns in error strings\n\terrStr := err.Error()\n\treturn strings.Contains(errStr, \"not found\") ||\n\t\tstrings.Contains(errStr, \"No such image\") ||\n\t\tstrings.Contains(errStr, \"manifest unknown\") ||\n\t\tstrings.Contains(errStr, \"NAME_UNKNOWN\")\n}\n"
  },
  {
    "path": "pkg/model/resolver_test.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"io\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/docker/api/types/image\"\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\tdockerspec \"github.com/moby/docker-image-spec/specs-go/v1\"\n\tocispec \"github.com/opencontainers/image-spec/specs-go/v1\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// mockDocker implements command.Command for testing.\ntype mockDocker struct {\n\tinspectFunc   func(ctx context.Context, ref string) (*image.InspectResponse, error)\n\tpullFunc      func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error)\n\tpushFunc      func(ctx context.Context, ref string) error\n\timageSaveFunc func(ctx context.Context, imageRef string) (io.ReadCloser, error)\n}\n\nfunc (m *mockDocker) Inspect(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\tif m.inspectFunc != nil {\n\t\treturn m.inspectFunc(ctx, ref)\n\t}\n\treturn nil, errors.New(\"not implemented\")\n}\n\nfunc (m *mockDocker) Pull(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) {\n\tif m.pullFunc != nil {\n\t\treturn m.pullFunc(ctx, ref, force)\n\t}\n\treturn nil, errors.New(\"mockDocker.Pull not implemented\")\n}\n\nfunc (m *mockDocker) Push(ctx context.Context, ref string) error {\n\tif m.pushFunc != nil {\n\t\treturn m.pushFunc(ctx, ref)\n\t}\n\treturn errors.New(\"mockDocker.Push not implemented\")\n}\n\nfunc (m *mockDocker) LoadUserInformation(ctx context.Context, registryHost string) (*command.UserInfo, error) {\n\treturn nil, errors.New(\"mockDocker.LoadUserInformation not implemented\")\n}\n\nfunc (m *mockDocker) ImageExists(ctx context.Context, ref string) (bool, error) {\n\treturn false, errors.New(\"mockDocker.ImageExists not implemented\")\n}\n\nfunc (m *mockDocker) ContainerLogs(ctx context.Context, containerID string, w io.Writer) error {\n\treturn errors.New(\"mockDocker.ContainerLogs not implemented\")\n}\n\nfunc (m *mockDocker) ContainerInspect(ctx context.Context, id string) (*container.InspectResponse, error) {\n\treturn nil, errors.New(\"mockDocker.ContainerInspect not implemented\")\n}\n\nfunc (m *mockDocker) ContainerStop(ctx context.Context, containerID string) error {\n\treturn errors.New(\"mockDocker.ContainerStop not implemented\")\n}\n\nfunc (m *mockDocker) RemoveImage(ctx context.Context, ref string) error {\n\treturn errors.New(\"mockDocker.RemoveImage not implemented\")\n}\n\nfunc (m *mockDocker) ImageBuild(ctx context.Context, options command.ImageBuildOptions) (string, error) {\n\treturn \"\", errors.New(\"mockDocker.ImageBuild not implemented\")\n}\n\nfunc (m *mockDocker) Run(ctx context.Context, options command.RunOptions) error {\n\treturn errors.New(\"mockDocker.Run not implemented\")\n}\n\nfunc (m *mockDocker) ContainerStart(ctx context.Context, options command.RunOptions) (string, error) {\n\treturn \"\", errors.New(\"mockDocker.ContainerStart not implemented\")\n}\n\nfunc (m *mockDocker) ImageSave(ctx context.Context, imageRef string) (io.ReadCloser, error) {\n\tif m.imageSaveFunc != nil {\n\t\treturn m.imageSaveFunc(ctx, imageRef)\n\t}\n\treturn nil, errors.New(\"mockDocker.ImageSave not implemented\")\n}\n\n// mockRegistry implements registry.Client for testing.\ntype mockRegistry struct {\n\tinspectFunc       func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error)\n\tgetImageFunc      func(ctx context.Context, ref string, platform *registry.Platform) (v1.Image, error)\n\tgetDescriptorFunc func(ctx context.Context, ref string) (v1.Descriptor, error)\n\tpushImageFunc     func(ctx context.Context, ref string, img v1.Image) error\n\tpushIndexFunc     func(ctx context.Context, ref string, idx v1.ImageIndex) error\n\twriteLayerFunc    func(ctx context.Context, opts registry.WriteLayerOptions) error\n}\n\nfunc (m *mockRegistry) Inspect(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\tif m.inspectFunc != nil {\n\t\treturn m.inspectFunc(ctx, ref, platform)\n\t}\n\treturn nil, registry.NotFoundError\n}\n\nfunc (m *mockRegistry) GetImage(ctx context.Context, ref string, platform *registry.Platform) (v1.Image, error) {\n\tif m.getImageFunc != nil {\n\t\treturn m.getImageFunc(ctx, ref, platform)\n\t}\n\treturn nil, errors.New(\"mockRegistry.GetImage not implemented\")\n}\n\nfunc (m *mockRegistry) GetDescriptor(ctx context.Context, ref string) (v1.Descriptor, error) {\n\tif m.getDescriptorFunc != nil {\n\t\treturn m.getDescriptorFunc(ctx, ref)\n\t}\n\treturn v1.Descriptor{}, errors.New(\"mockRegistry.GetDescriptor not implemented\")\n}\n\nfunc (m *mockRegistry) Exists(ctx context.Context, ref string) (bool, error) {\n\treturn false, errors.New(\"mockRegistry.Exists not implemented\")\n}\n\nfunc (m *mockRegistry) PushImage(ctx context.Context, ref string, img v1.Image) error {\n\tif m.pushImageFunc != nil {\n\t\treturn m.pushImageFunc(ctx, ref, img)\n\t}\n\treturn errors.New(\"mockRegistry.PushImage not implemented\")\n}\n\nfunc (m *mockRegistry) PushIndex(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\tif m.pushIndexFunc != nil {\n\t\treturn m.pushIndexFunc(ctx, ref, idx)\n\t}\n\treturn errors.New(\"mockRegistry.PushIndex not implemented\")\n}\n\nfunc (m *mockRegistry) WriteLayer(ctx context.Context, opts registry.WriteLayerOptions) error {\n\tif m.writeLayerFunc != nil {\n\t\treturn m.writeLayerFunc(ctx, opts)\n\t}\n\t// Default: no-op. The caller (WeightPusher) owns closing ProgressCh.\n\treturn nil\n}\n\n// mockFactory implements Factory for testing.\ntype mockFactory struct {\n\tname      string\n\tbuildFunc func(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error)\n}\n\nfunc (f *mockFactory) Build(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error) {\n\tif f.buildFunc != nil {\n\t\treturn f.buildFunc(ctx, src, opts)\n\t}\n\treturn &ImageArtifact{Reference: opts.ImageName, Source: ImageSourceBuild}, nil\n}\n\nfunc (f *mockFactory) Name() string {\n\treturn f.name\n}\n\nfunc TestNewResolver(t *testing.T) {\n\tdocker := &mockDocker{}\n\treg := &mockRegistry{}\n\n\tresolver := NewResolver(docker, reg)\n\n\trequire.NotNil(t, resolver)\n\trequire.Equal(t, \"dockerfile\", resolver.factory.Name())\n}\n\nfunc TestResolver_WithFactory(t *testing.T) {\n\tdocker := &mockDocker{}\n\treg := &mockRegistry{}\n\n\tresolver := NewResolver(docker, reg)\n\trequire.Equal(t, \"dockerfile\", resolver.factory.Name())\n\n\tcustomFactory := &mockFactory{name: \"custom\"}\n\tresult := resolver.WithFactory(customFactory)\n\n\t// WithFactory returns the same resolver for chaining\n\trequire.Same(t, resolver, result)\n\trequire.Equal(t, \"custom\", resolver.factory.Name())\n}\n\nfunc TestResolver_Inspect_LocalOnly_Found(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\treg := &mockRegistry{}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Inspect(context.Background(), ref, LocalOnly())\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.Equal(t, ImageSourceLocal, model.Image.Source)\n\trequire.Equal(t, \"0.10.0\", model.CogVersion)\n\trequire.NotNil(t, model.Config)\n\trequire.Equal(t, \"3.11\", model.Config.Build.PythonVersion)\n}\n\nfunc TestResolver_Inspect_LocalOnly_NotFound(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"No such image: my-image:latest\")\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tt.Fatal(\"should not call registry when LocalOnly\")\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Inspect(context.Background(), ref, LocalOnly())\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"not found locally\")\n}\n\nfunc TestResolver_Inspect_RemoteOnly_Found(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tt.Fatal(\"should not call docker.Inspect when RemoteOnly\")\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\treturn ®istry.ManifestResult{\n\t\t\t\tSchemaVersion: 2,\n\t\t\t\tMediaType:     \"application/vnd.docker.distribution.manifest.v2+json\",\n\t\t\t\tConfig:        \"sha256:configdigest\",\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"r8.im/user/model\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Inspect(context.Background(), ref, RemoteOnly())\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.Equal(t, ImageSourceRemote, model.Image.Source)\n\trequire.Equal(t, \"0.10.0\", model.CogVersion)\n}\n\nfunc TestResolver_Inspect_RemoteOnly_NotFound(t *testing.T) {\n\tdocker := &mockDocker{}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\treturn nil, registry.NotFoundError\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"r8.im/user/model\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Inspect(context.Background(), ref, RemoteOnly())\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"not found in registry\")\n}\n\nfunc TestResolver_Inspect_RemoteOnly_NotCogModel(t *testing.T) {\n\tdocker := &mockDocker{}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\treturn ®istry.ManifestResult{\n\t\t\t\tSchemaVersion: 2,\n\t\t\t\tMediaType:     \"application/vnd.docker.distribution.manifest.v2+json\",\n\t\t\t\tConfig:        \"sha256:configdigest\",\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\t// No Cog labels - just a regular image\n\t\t\t\t\t\"maintainer\": \"someone@example.com\",\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"nginx:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Inspect(context.Background(), ref, RemoteOnly())\n\n\trequire.Error(t, err)\n\trequire.ErrorIs(t, err, ErrNotCogModel)\n}\n\nfunc TestResolver_Inspect_PreferLocal_FoundLocally(t *testing.T) {\n\tlocalCalled := false\n\tremoteCalled := false\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tlocalCalled = true\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:local123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.9.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tremoteCalled = true\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Inspect(context.Background(), ref, PreferLocal())\n\n\trequire.NoError(t, err)\n\trequire.True(t, localCalled, \"should try local first\")\n\trequire.False(t, remoteCalled, \"should not call remote when local succeeds\")\n\trequire.Equal(t, ImageSourceLocal, model.Image.Source)\n}\n\nfunc TestResolver_Inspect_PreferLocal_Fallback(t *testing.T) {\n\tlocalCalled := false\n\tremoteCalled := false\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tlocalCalled = true\n\t\t\treturn nil, errors.New(\"No such image\")\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tremoteCalled = true\n\t\t\treturn ®istry.ManifestResult{\n\t\t\t\tSchemaVersion: 2,\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Inspect(context.Background(), ref, PreferLocal())\n\n\trequire.NoError(t, err)\n\trequire.True(t, localCalled, \"should try local first\")\n\trequire.True(t, remoteCalled, \"should fall back to remote\")\n\trequire.Equal(t, ImageSourceRemote, model.Image.Source)\n}\n\nfunc TestResolver_Inspect_PreferLocal_NoFallbackOnRealError(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"connection refused\") // Real error, not \"not found\"\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tt.Fatal(\"should not fall back to remote on real error\")\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Inspect(context.Background(), ref, PreferLocal())\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"failed to inspect local image\")\n\trequire.Contains(t, err.Error(), \"connection refused\")\n}\n\nfunc TestResolver_Inspect_PreferRemote_FoundRemotely(t *testing.T) {\n\tlocalCalled := false\n\tremoteCalled := false\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tlocalCalled = true\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tremoteCalled = true\n\t\t\treturn ®istry.ManifestResult{\n\t\t\t\tSchemaVersion: 2,\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Inspect(context.Background(), ref, PreferRemote())\n\n\trequire.NoError(t, err)\n\trequire.False(t, localCalled, \"should not try local when remote succeeds\")\n\trequire.True(t, remoteCalled, \"should try remote first\")\n\trequire.Equal(t, ImageSourceRemote, model.Image.Source)\n}\n\nfunc TestResolver_Inspect_PreferRemote_Fallback(t *testing.T) {\n\tlocalCalled := false\n\tremoteCalled := false\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tlocalCalled = true\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:local123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tremoteCalled = true\n\t\t\treturn nil, errors.New(\"manifest unknown\")\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Inspect(context.Background(), ref, PreferRemote())\n\n\trequire.NoError(t, err)\n\trequire.True(t, remoteCalled, \"should try remote first\")\n\trequire.True(t, localCalled, \"should fall back to local\")\n\trequire.Equal(t, ImageSourceLocal, model.Image.Source)\n}\n\nfunc TestResolver_Inspect_PreferRemote_NoFallbackOnRealError(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tt.Fatal(\"should not fall back to local on real error\")\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\treturn nil, errors.New(\"authentication required\")\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Inspect(context.Background(), ref, PreferRemote())\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"failed to inspect remote image\")\n\trequire.Contains(t, err.Error(), \"authentication required\")\n}\n\nfunc TestResolver_Inspect_WithPlatform(t *testing.T) {\n\tvar capturedPlatform *registry.Platform\n\n\tdocker := &mockDocker{}\n\treg := &mockRegistry{\n\t\tinspectFunc: func(ctx context.Context, ref string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\t\t\tcapturedPlatform = platform\n\t\t\treturn ®istry.ManifestResult{\n\t\t\t\tSchemaVersion: 2,\n\t\t\t\tLabels: map[string]string{\n\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, reg)\n\tref, err := ParseRef(\"my-image\")\n\trequire.NoError(t, err)\n\n\tplatform := ®istry.Platform{OS: \"linux\", Architecture: \"amd64\"}\n\t_, err = resolver.Inspect(context.Background(), ref, RemoteOnly(), WithPlatform(platform))\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, capturedPlatform)\n\trequire.Equal(t, \"linux\", capturedPlatform.OS)\n\trequire.Equal(t, \"amd64\", capturedPlatform.Architecture)\n}\n\nfunc TestResolver_Inspect_ParsesConfigFromLabels(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"gpu\":true,\"python_version\":\"3.12\"},\"predict\":\"predict.py:Predictor\"}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.11.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Inspect(context.Background(), ref, LocalOnly())\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model.Config)\n\trequire.NotNil(t, model.Config.Build)\n\trequire.True(t, model.Config.Build.GPU)\n\trequire.Equal(t, \"3.12\", model.Config.Build.PythonVersion)\n\trequire.Equal(t, \"predict.py:Predictor\", model.Config.Predict)\n}\n\nfunc TestResolver_Inspect_InvalidConfigJSON(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig: `{invalid json`,\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Inspect(context.Background(), ref, LocalOnly())\n\n\trequire.Error(t, err)\n\t// Error should contain the JSON parse error message\n\trequire.Contains(t, err.Error(), \"invalid character\")\n}\n\nfunc TestResolver_Inspect_NoConfigLabel_ReturnsErrNotCogModel(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\t// No LabelConfig - just version label\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Inspect(context.Background(), ref, LocalOnly())\n\n\t// Without LabelConfig, image is not a valid Cog model\n\trequire.Error(t, err)\n\trequire.ErrorIs(t, err, ErrNotCogModel)\n}\n\nfunc TestResolver_Inspect_NotCogModel(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\t// No Cog labels at all - just some random image\n\t\t\t\t\t\t\t\"maintainer\": \"someone@example.com\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"nginx:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Inspect(context.Background(), ref, LocalOnly())\n\n\trequire.Error(t, err)\n\trequire.ErrorIs(t, err, ErrNotCogModel)\n\trequire.Contains(t, err.Error(), \"nginx:latest\")\n}\n\nfunc TestResolver_InspectByID_Found(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\t// Verify the ID was passed directly (not mangled by ParseRef)\n\t\t\trequire.Equal(t, \"9056219a5fb2\", ref)\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:9056219a5fb2abc123def456\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\n\tmodel, err := resolver.InspectByID(context.Background(), \"9056219a5fb2\")\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.Equal(t, ImageSourceLocal, model.Image.Source)\n\trequire.Equal(t, \"sha256:9056219a5fb2abc123def456\", model.Image.Digest)\n\trequire.Equal(t, \"sha256:9056219a5fb2abc123def456\", model.Image.Reference)\n\trequire.Equal(t, \"0.10.0\", model.CogVersion)\n\trequire.NotNil(t, model.Config)\n\trequire.Equal(t, \"3.11\", model.Config.Build.PythonVersion)\n}\n\nfunc TestResolver_InspectByID_FullSHA(t *testing.T) {\n\tfullID := \"sha256:9056219a5fb2abc123def456789\"\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\trequire.Equal(t, fullID, ref)\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: fullID,\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\n\tmodel, err := resolver.InspectByID(context.Background(), fullID)\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, model)\n\trequire.Equal(t, fullID, model.Image.Digest)\n}\n\nfunc TestResolver_InspectByID_NotCogModel(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\t// No Cog labels\n\t\t\t\t\t\t\t\"maintainer\": \"someone\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\n\t_, err := resolver.InspectByID(context.Background(), \"abc123\")\n\n\trequire.Error(t, err)\n\trequire.ErrorIs(t, err, ErrNotCogModel)\n}\n\nfunc TestResolver_InspectByID_NotFound(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"No such image: abc123\")\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\n\t_, err := resolver.InspectByID(context.Background(), \"abc123\")\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"failed to load image by ID\")\n}\n\n// =============================================================================\n// Pull tests\n// =============================================================================\n\nfunc TestResolver_Pull_AlreadyLocal(t *testing.T) {\n\tpullCalled := false\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"gpu\":false}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t\tpullFunc: func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) {\n\t\t\tpullCalled = true\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Pull(context.Background(), ref)\n\n\trequire.NoError(t, err)\n\trequire.False(t, pullCalled, \"should not pull when image exists locally\")\n\trequire.NotNil(t, model)\n\trequire.Equal(t, \"0.10.0\", model.CogVersion)\n}\n\nfunc TestResolver_Pull_NotLocal_PullsAndReturns(t *testing.T) {\n\tpullCalled := false\n\tinspectCalls := 0\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tinspectCalls++\n\t\t\tif inspectCalls == 1 {\n\t\t\t\t// First call: not found locally\n\t\t\t\treturn nil, errors.New(\"No such image\")\n\t\t\t}\n\t\t\t// After pull: found\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"gpu\":true}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t\tpullFunc: func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) {\n\t\t\tpullCalled = true\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"gpu\":true}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"r8.im/user/model:latest\")\n\trequire.NoError(t, err)\n\n\tmodel, err := resolver.Pull(context.Background(), ref)\n\n\trequire.NoError(t, err)\n\trequire.True(t, pullCalled, \"should call Pull when image not local\")\n\trequire.NotNil(t, model)\n\trequire.True(t, model.HasGPU())\n}\n\nfunc TestResolver_Pull_NotCogModel(t *testing.T) {\n\tinspectCalls := 0\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\tinspectCalls++\n\t\t\tif inspectCalls == 1 {\n\t\t\t\t// First call: not found locally\n\t\t\t\treturn nil, errors.New(\"No such image\")\n\t\t\t}\n\t\t\t// After pull: found but not a Cog model\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\t// Not a Cog model\n\t\t\t\t\t\t\t\"some.label\": \"value\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t\tpullFunc: func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: \"sha256:abc123\",\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\t\"some.label\": \"value\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"not-cog:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Pull(context.Background(), ref)\n\n\trequire.Error(t, err)\n\trequire.ErrorIs(t, err, ErrNotCogModel)\n}\n\nfunc TestResolver_Pull_LocalOnly_NotFound(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"No such image\")\n\t\t},\n\t\tpullFunc: func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) {\n\t\t\tt.Fatal(\"should not pull when LocalOnly\")\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Pull(context.Background(), ref, LocalOnly())\n\n\trequire.Error(t, err)\n\trequire.ErrorIs(t, err, ErrNotFound)\n}\n\nfunc TestResolver_Pull_PullFails(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"No such image\")\n\t\t},\n\t\tpullFunc: func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"manifest unknown\")\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"nonexistent:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Pull(context.Background(), ref)\n\n\trequire.Error(t, err)\n\trequire.ErrorIs(t, err, ErrNotFound)\n}\n\nfunc TestResolver_Pull_LocalInspectRealError(t *testing.T) {\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn nil, errors.New(\"connection refused\")\n\t\t},\n\t\tpullFunc: func(ctx context.Context, ref string, force bool) (*image.InspectResponse, error) {\n\t\t\tt.Fatal(\"should not pull when local inspect has real error\")\n\t\t\treturn nil, nil\n\t\t},\n\t}\n\n\tresolver := NewResolver(docker, &mockRegistry{})\n\tref, err := ParseRef(\"my-image:latest\")\n\trequire.NoError(t, err)\n\n\t_, err = resolver.Pull(context.Background(), ref)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"connection refused\")\n}\n\n// =============================================================================\n// Build tests\n// =============================================================================\n\nfunc TestResolver_Build_NoWeightsManifestWithoutWeights(t *testing.T) {\n\tvalidDigest := \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\"\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: validDigest,\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.10.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tfactory := &mockFactory{}\n\tresolver := NewResolver(docker, &mockRegistry{}).WithFactory(factory)\n\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Build: &config.Build{}},\n\t\tProjectDir: t.TempDir(),\n\t}\n\n\tm, err := resolver.Build(context.Background(), src, BuildOptions{\n\t\tImageName: \"test-image\",\n\t})\n\n\trequire.NoError(t, err)\n\trequire.False(t, m.IsBundle())\n\trequire.Empty(t, m.WeightArtifacts())\n}\n\nfunc TestResolver_Build_PopulatesArtifacts(t *testing.T) {\n\timageDigest := \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\"\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: imageDigest,\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.15.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tfactory := &mockFactory{\n\t\tbuildFunc: func(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error) {\n\t\t\treturn &ImageArtifact{\n\t\t\t\tReference: opts.ImageName,\n\t\t\t\tDigest:    imageDigest,\n\t\t\t\tSource:    ImageSourceBuild,\n\t\t\t}, nil\n\t\t},\n\t}\n\tresolver := NewResolver(docker, &mockRegistry{}).WithFactory(factory)\n\n\tsrc := &Source{\n\t\tConfig:     &config.Config{Build: &config.Build{}},\n\t\tProjectDir: t.TempDir(),\n\t}\n\n\tm, err := resolver.Build(context.Background(), src, BuildOptions{\n\t\tImageName: \"test-image:latest\",\n\t})\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, m.Artifacts, \"Build should populate Artifacts\")\n\trequire.Len(t, m.Artifacts, 1, \"should have exactly one artifact (image)\")\n\n\t// Verify it's an ImageArtifact with correct data\n\timgArtifact := m.GetImageArtifact()\n\trequire.NotNil(t, imgArtifact, \"should contain an ImageArtifact\")\n\trequire.Equal(t, \"model\", imgArtifact.Name())\n\trequire.Equal(t, ArtifactTypeImage, imgArtifact.Type())\n\trequire.Equal(t, \"test-image:latest\", imgArtifact.Reference)\n\trequire.Equal(t, imageDigest, imgArtifact.Descriptor().Digest.String())\n}\n\nfunc TestResolver_Build_PopulatesWeightArtifacts(t *testing.T) {\n\timageDigest := \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\"\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: imageDigest,\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.15.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tfactory := &mockFactory{\n\t\tbuildFunc: func(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error) {\n\t\t\treturn &ImageArtifact{\n\t\t\t\tReference: opts.ImageName,\n\t\t\t\tDigest:    imageDigest,\n\t\t\t\tSource:    ImageSourceBuild,\n\t\t\t}, nil\n\t\t},\n\t}\n\tresolver := NewResolver(docker, &mockRegistry{}).WithFactory(factory)\n\n\t// Create a temp directory with a real weight file\n\tdir := t.TempDir()\n\tweightContent := []byte(\"test weight for resolver build\")\n\trequire.NoError(t, os.WriteFile(filepath.Join(dir, \"model.safetensors\"), weightContent, 0o644))\n\n\tsrc := &Source{\n\t\tConfig: &config.Config{\n\t\t\tBuild: &config.Build{},\n\t\t\tWeights: []config.WeightSource{\n\t\t\t\t{Name: \"my-model\", Source: \"model.safetensors\", Target: \"/srv/weights/model.safetensors\"},\n\t\t\t},\n\t\t},\n\t\tProjectDir: dir,\n\t}\n\n\tm, err := resolver.Build(context.Background(), src, BuildOptions{\n\t\tImageName: \"test-image:latest\",\n\t\tOCIIndex:  true,\n\t})\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, m.Artifacts)\n\n\t// Should have 2 artifacts: 1 image + 1 weight\n\trequire.Len(t, m.Artifacts, 2, \"should have image + weight artifacts\")\n\n\t// Verify image artifact\n\timgArtifact := m.GetImageArtifact()\n\trequire.NotNil(t, imgArtifact)\n\trequire.Equal(t, \"model\", imgArtifact.Name())\n\n\t// Verify weight artifact\n\tweightArtifacts := m.WeightArtifacts()\n\trequire.Len(t, weightArtifacts, 1)\n\twa := weightArtifacts[0]\n\trequire.Equal(t, \"my-model\", wa.Name())\n\trequire.Equal(t, ArtifactTypeWeight, wa.Type())\n\trequire.Equal(t, \"/srv/weights/model.safetensors\", wa.Target)\n\trequire.Equal(t, filepath.Join(dir, \"model.safetensors\"), wa.FilePath)\n\n\t// Weight config should be populated\n\trequire.Equal(t, \"1.0\", wa.Config.SchemaVersion)\n\trequire.Equal(t, \"my-model\", wa.Config.Name)\n\trequire.Equal(t, \"/srv/weights/model.safetensors\", wa.Config.Target)\n\trequire.False(t, wa.Config.Created.IsZero())\n}\n\nfunc TestResolver_Build_WithWeightsLoadsManifest(t *testing.T) {\n\timageDigest := \"sha256:a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2\"\n\n\tdocker := &mockDocker{\n\t\tinspectFunc: func(ctx context.Context, ref string) (*image.InspectResponse, error) {\n\t\t\treturn &image.InspectResponse{\n\t\t\t\tID: imageDigest,\n\t\t\t\tConfig: &dockerspec.DockerOCIImageConfig{\n\t\t\t\t\tImageConfig: ocispec.ImageConfig{\n\t\t\t\t\t\tLabels: map[string]string{\n\n\t\t\t\t\t\t\tLabelConfig:  `{\"build\":{\"python_version\":\"3.11\"}}`,\n\t\t\t\t\t\t\tLabelVersion: \"0.15.0\",\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t}, nil\n\t\t},\n\t}\n\n\tfactory := &mockFactory{\n\t\tbuildFunc: func(ctx context.Context, src *Source, opts BuildOptions) (*ImageArtifact, error) {\n\t\t\treturn &ImageArtifact{\n\t\t\t\tReference: opts.ImageName,\n\t\t\t\tDigest:    imageDigest,\n\t\t\t\tSource:    ImageSourceBuild,\n\t\t\t}, nil\n\t\t},\n\t}\n\tresolver := NewResolver(docker, &mockRegistry{}).WithFactory(factory)\n\n\tdir := t.TempDir()\n\trequire.NoError(t, os.WriteFile(filepath.Join(dir, \"model.bin\"), []byte(\"test weights\"), 0o644))\n\n\tsrc := &Source{\n\t\tConfig: &config.Config{\n\t\t\tBuild: &config.Build{},\n\t\t\tWeights: []config.WeightSource{\n\t\t\t\t{Name: \"my-model\", Source: \"model.bin\", Target: \"/weights/model.bin\"},\n\t\t\t},\n\t\t},\n\t\tProjectDir: dir,\n\t}\n\n\tm, err := resolver.Build(context.Background(), src, BuildOptions{\n\t\tImageName: \"test-image:latest\",\n\t\tOCIIndex:  true,\n\t})\n\n\trequire.NoError(t, err)\n\trequire.True(t, m.IsBundle())\n\trequire.True(t, m.OCIIndex)\n\n\t// Should have 2 artifacts: image + weight\n\trequire.Len(t, m.Artifacts, 2)\n\trequire.NotNil(t, m.GetImageArtifact())\n\trequire.Len(t, m.WeightArtifacts(), 1)\n\n\t// Weight artifacts should be populated\n\trequire.Len(t, m.WeightArtifacts(), 1)\n}\n\nfunc TestIndexDetectionHelpers(t *testing.T) {\n\tt.Run(\"findWeightsManifest\", func(t *testing.T) {\n\t\tmanifests := []registry.PlatformManifest{\n\t\t\t{Digest: \"sha256:image123\", OS: \"linux\", Architecture: \"amd64\"},\n\t\t\t{\n\t\t\t\tDigest:       \"sha256:weights456\",\n\t\t\t\tOS:           PlatformUnknown,\n\t\t\t\tArchitecture: PlatformUnknown,\n\t\t\t\tAnnotations: map[string]string{\n\t\t\t\t\tAnnotationReferenceType: AnnotationValueWeights,\n\t\t\t\t},\n\t\t\t},\n\t\t}\n\n\t\twm := findWeightsManifest(manifests)\n\t\trequire.NotNil(t, wm)\n\t\trequire.Equal(t, \"sha256:weights456\", wm.Digest)\n\t})\n\n\tt.Run(\"findWeightsManifest not found\", func(t *testing.T) {\n\t\tmanifests := []registry.PlatformManifest{\n\t\t\t{Digest: \"sha256:image123\", OS: \"linux\", Architecture: \"amd64\"},\n\t\t}\n\n\t\twm := findWeightsManifest(manifests)\n\t\trequire.Nil(t, wm)\n\t})\n\n\tt.Run(\"findImageManifest\", func(t *testing.T) {\n\t\tmanifests := []registry.PlatformManifest{\n\t\t\t{Digest: \"sha256:image123\", OS: \"linux\", Architecture: \"amd64\"},\n\t\t\t{Digest: \"sha256:weights456\", OS: PlatformUnknown, Architecture: PlatformUnknown},\n\t\t}\n\n\t\tplatform := ®istry.Platform{OS: \"linux\", Architecture: \"amd64\"}\n\t\tim := findImageManifest(manifests, platform)\n\t\trequire.NotNil(t, im)\n\t\trequire.Equal(t, \"sha256:image123\", im.Digest)\n\t})\n\n\tt.Run(\"findImageManifest skips unknown\", func(t *testing.T) {\n\t\tmanifests := []registry.PlatformManifest{\n\t\t\t{Digest: \"sha256:weights456\", OS: PlatformUnknown, Architecture: PlatformUnknown},\n\t\t}\n\n\t\tim := findImageManifest(manifests, nil)\n\t\trequire.Nil(t, im)\n\t})\n\n\tt.Run(\"findImageManifest no platform filter\", func(t *testing.T) {\n\t\tmanifests := []registry.PlatformManifest{\n\t\t\t{Digest: \"sha256:arm123\", OS: \"linux\", Architecture: \"arm64\"},\n\t\t\t{Digest: \"sha256:weights456\", OS: PlatformUnknown, Architecture: PlatformUnknown},\n\t\t}\n\n\t\tim := findImageManifest(manifests, nil)\n\t\trequire.NotNil(t, im)\n\t\trequire.Equal(t, \"sha256:arm123\", im.Digest)\n\t})\n\n\tt.Run(\"findImageManifest platform mismatch\", func(t *testing.T) {\n\t\tmanifests := []registry.PlatformManifest{\n\t\t\t{Digest: \"sha256:arm123\", OS: \"linux\", Architecture: \"arm64\"},\n\t\t\t{Digest: \"sha256:weights456\", OS: PlatformUnknown, Architecture: PlatformUnknown},\n\t\t}\n\n\t\tplatform := ®istry.Platform{OS: \"linux\", Architecture: \"amd64\"}\n\t\tim := findImageManifest(manifests, platform)\n\t\trequire.Nil(t, im)\n\t})\n\n\tt.Run(\"isOCIIndex with index\", func(t *testing.T) {\n\t\tmr := ®istry.ManifestResult{\n\t\t\tMediaType: \"application/vnd.oci.image.index.v1+json\",\n\t\t}\n\t\trequire.True(t, isOCIIndex(mr))\n\t})\n\n\tt.Run(\"isOCIIndex with single manifest\", func(t *testing.T) {\n\t\tmr := ®istry.ManifestResult{\n\t\t\tMediaType: \"application/vnd.oci.image.manifest.v1+json\",\n\t\t}\n\t\trequire.False(t, isOCIIndex(mr))\n\t})\n}\n\nfunc TestIsNotFoundError(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\terr      error\n\t\texpected bool\n\t}{\n\t\t{\n\t\t\tname:     \"nil error\",\n\t\t\terr:      nil,\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"No such image\",\n\t\t\terr:      errors.New(\"No such image: my-image:latest\"),\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"not found\",\n\t\t\terr:      errors.New(\"image not found\"),\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"manifest unknown\",\n\t\t\terr:      errors.New(\"manifest unknown: repository does not exist\"),\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"NAME_UNKNOWN\",\n\t\t\terr:      errors.New(\"NAME_UNKNOWN: repository name not known to registry\"),\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"connection refused\",\n\t\t\terr:      errors.New(\"connection refused\"),\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"authentication required\",\n\t\t\terr:      errors.New(\"authentication required\"),\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"context canceled\",\n\t\t\terr:      context.Canceled,\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"context deadline exceeded\",\n\t\t\terr:      context.DeadlineExceeded,\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"registry NotFoundError\",\n\t\t\terr:      registry.NotFoundError,\n\t\t\texpected: true,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := isNotFoundError(tt.err)\n\t\t\trequire.Equal(t, tt.expected, result)\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "pkg/model/source.go",
    "content": "package model\n\nimport (\n\t\"os\"\n\t\"path/filepath\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// Source represents a Cog project ready to build.\n// It combines the parsed configuration with the project directory location.\ntype Source struct {\n\tConfig         *config.Config\n\tProjectDir     string\n\tConfigFilename string // Base filename like \"cog.yaml\" or \"my-config.yaml\"\n\tWarnings       []config.DeprecationWarning\n}\n\n// NewSource loads configuration from the given path and returns a Source.\n// The configPath can be a filename (e.g., \"cog.yaml\") which will be searched\n// for in the current directory and parent directories.\nfunc NewSource(configPath string) (*Source, error) {\n\tif configPath == \"\" {\n\t\tconfigPath = \"cog.yaml\"\n\t}\n\n\t// Find the root project directory\n\trootDir, err := config.GetProjectDir(configPath)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Open and read the config file\n\tfullPath := filepath.Join(rootDir, configPath)\n\tf, err := os.Open(fullPath)\n\tif err != nil {\n\t\treturn nil, &config.ParseError{Filename: configPath, Err: err}\n\t}\n\tdefer f.Close()\n\n\tresult, err := config.Load(f, rootDir)\n\tif err != nil {\n\t\t// Add filename context to parse errors if not already present\n\t\tif parseErr, ok := err.(*config.ParseError); ok && parseErr.Filename == \"\" {\n\t\t\tparseErr.Filename = configPath\n\t\t}\n\t\treturn nil, err\n\t}\n\n\t// Display deprecation warnings\n\tfor _, w := range result.Warnings {\n\t\tconsole.Warnf(\"%s\", w.Error())\n\t}\n\n\treturn &Source{\n\t\tConfig:         result.Config,\n\t\tProjectDir:     result.RootDir,\n\t\tConfigFilename: filepath.Base(configPath),\n\t\tWarnings:       result.Warnings,\n\t}, nil\n}\n\n// NewSourceFromConfig creates a Source from an existing Config.\n// Use this when you already have a parsed config and know the project directory.\nfunc NewSourceFromConfig(cfg *config.Config, projectDir string) *Source {\n\treturn &Source{\n\t\tConfig:         cfg,\n\t\tProjectDir:     projectDir,\n\t\tConfigFilename: \"cog.yaml\",\n\t}\n}\n\n// ArtifactSpecs returns the artifact declarations derived from this source.\n// Always produces at least one ImageSpec. Produces a WeightSpec for each\n// weight declared in the config. Returns nil if Config is nil.\nfunc (s *Source) ArtifactSpecs() []ArtifactSpec {\n\tif s.Config == nil {\n\t\treturn nil\n\t}\n\n\tvar specs []ArtifactSpec\n\n\t// Always have an image artifact\n\tspecs = append(specs, NewImageSpec(\"model\", s.Config.Image))\n\n\t// Add weight specs from config\n\tfor _, w := range s.Config.Weights {\n\t\tspecs = append(specs, NewWeightSpec(w.Name, w.Source, w.Target))\n\t}\n\n\treturn specs\n}\n"
  },
  {
    "path": "pkg/model/source_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc TestNewSourceFromConfig(t *testing.T) {\n\tcfg := &config.Config{\n\t\tBuild: &config.Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.11\",\n\t\t},\n\t}\n\tprojectDir := \"/path/to/project\"\n\n\tsrc := NewSourceFromConfig(cfg, projectDir)\n\n\trequire.NotNil(t, src)\n\trequire.Equal(t, cfg, src.Config)\n\trequire.Equal(t, projectDir, src.ProjectDir)\n}\n\nfunc TestNewSourceFromConfig_NilConfig(t *testing.T) {\n\tsrc := NewSourceFromConfig(nil, \"/path/to/project\")\n\n\trequire.NotNil(t, src)\n\trequire.Nil(t, src.Config)\n\trequire.Equal(t, \"/path/to/project\", src.ProjectDir)\n}\n\nfunc TestSource_ArtifactSpecs_NoWeights(t *testing.T) {\n\tcfg := &config.Config{\n\t\tImage: \"r8.im/user/model\",\n\t\tBuild: &config.Build{\n\t\t\tGPU:           true,\n\t\t\tPythonVersion: \"3.11\",\n\t\t},\n\t}\n\tsrc := NewSourceFromConfig(cfg, \"/path/to/project\")\n\n\tspecs := src.ArtifactSpecs()\n\n\trequire.Len(t, specs, 1)\n\n\t// First spec should be an ImageSpec\n\timgSpec, ok := specs[0].(*ImageSpec)\n\trequire.True(t, ok, \"first spec should be *ImageSpec\")\n\trequire.Equal(t, ArtifactTypeImage, imgSpec.Type())\n\trequire.Equal(t, \"model\", imgSpec.Name())\n\trequire.Equal(t, \"r8.im/user/model\", imgSpec.ImageName)\n}\n\nfunc TestSource_ArtifactSpecs_WithWeights(t *testing.T) {\n\tcfg := &config.Config{\n\t\tImage: \"r8.im/user/model\",\n\t\tBuild: &config.Build{PythonVersion: \"3.11\"},\n\t\tWeights: []config.WeightSource{\n\t\t\t{Name: \"llama-7b\", Source: \"/data/llama-7b.safetensors\", Target: \"/weights/llama-7b.safetensors\"},\n\t\t\t{Name: \"embeddings\", Source: \"/data/embeddings.bin\", Target: \"/weights/embeddings.bin\"},\n\t\t},\n\t}\n\tsrc := NewSourceFromConfig(cfg, \"/path/to/project\")\n\n\tspecs := src.ArtifactSpecs()\n\n\trequire.Len(t, specs, 3) // 1 image + 2 weights\n\n\t// First is always the image\n\timgSpec, ok := specs[0].(*ImageSpec)\n\trequire.True(t, ok, \"first spec should be *ImageSpec\")\n\trequire.Equal(t, ArtifactTypeImage, imgSpec.Type())\n\n\t// Remaining are weight specs in order\n\tw1, ok := specs[1].(*WeightSpec)\n\trequire.True(t, ok, \"second spec should be *WeightSpec\")\n\trequire.Equal(t, ArtifactTypeWeight, w1.Type())\n\trequire.Equal(t, \"llama-7b\", w1.Name())\n\trequire.Equal(t, \"/data/llama-7b.safetensors\", w1.Source)\n\trequire.Equal(t, \"/weights/llama-7b.safetensors\", w1.Target)\n\n\tw2, ok := specs[2].(*WeightSpec)\n\trequire.True(t, ok, \"third spec should be *WeightSpec\")\n\trequire.Equal(t, \"embeddings\", w2.Name())\n\trequire.Equal(t, \"/data/embeddings.bin\", w2.Source)\n\trequire.Equal(t, \"/weights/embeddings.bin\", w2.Target)\n}\n\nfunc TestSource_ArtifactSpecs_EmptyImageName(t *testing.T) {\n\tcfg := &config.Config{\n\t\tBuild: &config.Build{PythonVersion: \"3.11\"},\n\t}\n\tsrc := NewSourceFromConfig(cfg, \"/path/to/project\")\n\n\tspecs := src.ArtifactSpecs()\n\n\trequire.Len(t, specs, 1)\n\timgSpec, ok := specs[0].(*ImageSpec)\n\trequire.True(t, ok)\n\trequire.Equal(t, \"\", imgSpec.ImageName) // empty is fine; BuildOptions fills it later\n}\n\nfunc TestSource_ArtifactSpecs_NilConfig(t *testing.T) {\n\tsrc := NewSourceFromConfig(nil, \"/path/to/project\")\n\n\tspecs := src.ArtifactSpecs()\n\n\trequire.Nil(t, specs)\n}\n"
  },
  {
    "path": "pkg/model/weight_builder.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"time\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n)\n\n// WeightBuilder builds WeightArtifact from WeightSpec.\n// It hashes the source file, creates a WeightConfig, and manages a lockfile as build cache.\ntype WeightBuilder struct {\n\tsource     *Source\n\tcogVersion string\n\tlockPath   string\n}\n\n// NewWeightBuilder creates a WeightBuilder.\n// lockPath is where the weights.lock file is read/written as a build cache.\nfunc NewWeightBuilder(source *Source, cogVersion, lockPath string) *WeightBuilder {\n\treturn &WeightBuilder{\n\t\tsource:     source,\n\t\tcogVersion: cogVersion,\n\t\tlockPath:   lockPath,\n\t}\n}\n\n// Build builds a WeightArtifact from a WeightSpec.\n// It resolves the source file, computes its SHA256 digest, and creates the artifact\n// with a versioned WeightConfig.\nfunc (b *WeightBuilder) Build(ctx context.Context, spec ArtifactSpec) (Artifact, error) {\n\tws, ok := spec.(*WeightSpec)\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"weight builder: expected *WeightSpec, got %T\", spec)\n\t}\n\n\tselect {\n\tcase <-ctx.Done():\n\t\treturn nil, ctx.Err()\n\tdefault:\n\t}\n\n\t// Resolve file path\n\tabsPath := filepath.Join(b.source.ProjectDir, ws.Source)\n\n\t// Stat the file to check existence and size\n\tfi, err := os.Stat(absPath)\n\tif err != nil {\n\t\tif os.IsNotExist(err) {\n\t\t\treturn nil, fmt.Errorf(\"weight source not found: %s\", ws.Source)\n\t\t}\n\t\treturn nil, fmt.Errorf(\"stat weight file %s: %w\", ws.Source, err)\n\t}\n\n\t// Check lockfile cache: if we have a matching entry (name + size), skip hashing.\n\t// NOTE: This cache only checks name + file size. Same-size modifications (rare for\n\t// weight files) won't be detected. Delete the lockfile to force re-hashing.\n\t// TODO: Consider adding mtime to the cache key for stronger invalidation.\n\tvar digestStr string\n\tvar size int64\n\tif cached := b.findCachedEntry(ws.Name(), fi.Size()); cached != nil {\n\t\tdigestStr = cached.Digest\n\t\tsize = cached.Size\n\t} else {\n\t\t// Cache miss: hash the file\n\t\tdigestStr, size, err = hashFile(absPath)\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"hash weight file %s: %w\", ws.Source, err)\n\t\t}\n\t}\n\n\t// Parse as v1.Hash for the descriptor\n\tdigest, err := v1.NewHash(digestStr)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"parse digest: %w\", err)\n\t}\n\n\t// Build the WeightConfig\n\tcfg := WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    b.cogVersion,\n\t\tName:          ws.Name(),\n\t\tTarget:        ws.Target,\n\t\tCreated:       time.Now().UTC(),\n\t}\n\n\t// Build the descriptor\n\tdesc := v1.Descriptor{\n\t\tDigest:    digest,\n\t\tSize:      size,\n\t\tMediaType: MediaTypeWeightLayer,\n\t}\n\n\t// Update lockfile\n\tif err := b.updateLockfile(ws, digestStr, size); err != nil {\n\t\treturn nil, fmt.Errorf(\"update lockfile: %w\", err)\n\t}\n\n\treturn NewWeightArtifact(ws.Name(), desc, absPath, ws.Target, cfg), nil\n}\n\n// findCachedEntry checks the lockfile for an entry matching name and fileSize.\n// Returns the cached WeightFile if found and size matches, nil otherwise.\nfunc (b *WeightBuilder) findCachedEntry(name string, fileSize int64) *WeightFile {\n\tif _, err := os.Stat(b.lockPath); err != nil {\n\t\treturn nil\n\t}\n\tlock, err := LoadWeightsLock(b.lockPath)\n\tif err != nil {\n\t\treturn nil\n\t}\n\tfor i, f := range lock.Files {\n\t\tif f.Name == name && f.Size == fileSize {\n\t\t\treturn &lock.Files[i]\n\t\t}\n\t}\n\treturn nil\n}\n\n// updateLockfile loads the existing lockfile (if any), adds or updates\n// the entry for the given weight, and saves it back.\nfunc (b *WeightBuilder) updateLockfile(ws *WeightSpec, digest string, size int64) error {\n\t// Load existing lockfile, or start fresh.\n\t// LoadWeightsLock wraps the underlying error, so we check the raw file first.\n\tlock := &WeightsLock{\n\t\tVersion: \"1.0\",\n\t\tCreated: time.Now().UTC(),\n\t}\n\tif _, err := os.Stat(b.lockPath); err == nil {\n\t\texisting, loadErr := LoadWeightsLock(b.lockPath)\n\t\tif loadErr != nil {\n\t\t\treturn fmt.Errorf(\"load existing lockfile: %w\", loadErr)\n\t\t}\n\t\tlock = existing\n\t}\n\n\tentry := WeightFile{\n\t\tName:             ws.Name(),\n\t\tDest:             ws.Target,\n\t\tDigest:           digest,\n\t\tDigestOriginal:   digest,\n\t\tSize:             size,\n\t\tSizeUncompressed: size,\n\t\tMediaType:        MediaTypeWeightLayer,\n\t}\n\n\t// Update existing entry or append\n\tupdated := false\n\tfor i, f := range lock.Files {\n\t\tif f.Name == ws.Name() {\n\t\t\tlock.Files[i] = entry\n\t\t\tupdated = true\n\t\t\tbreak\n\t\t}\n\t}\n\tif !updated {\n\t\tlock.Files = append(lock.Files, entry)\n\t}\n\n\treturn lock.Save(b.lockPath)\n}\n"
  },
  {
    "path": "pkg/model/weight_builder_test.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"crypto/sha256\"\n\t\"encoding/hex\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc TestWeightBuilder_HappyPath(t *testing.T) {\n\t// Setup: real temp file as a weight source\n\ttmpDir := t.TempDir()\n\tweightContent := []byte(\"test weight data for builder\")\n\tweightFile := filepath.Join(tmpDir, \"model.safetensors\")\n\terr := os.WriteFile(weightFile, weightContent, 0o644)\n\trequire.NoError(t, err)\n\n\t// Compute expected digest\n\thash := sha256.Sum256(weightContent)\n\texpectedDigest := \"sha256:\" + hex.EncodeToString(hash[:])\n\n\t// Create source with config that has one weight\n\tsrc := NewSourceFromConfig(&config.Config{\n\t\tWeights: []config.WeightSource{\n\t\t\t{Name: \"my-model\", Source: \"model.safetensors\", Target: \"/srv/weights/model.safetensors\"},\n\t\t},\n\t}, tmpDir)\n\n\t// Create a WeightBuilder\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\twb := NewWeightBuilder(src, \"0.15.0\", lockPath)\n\n\t// Build from the weight spec\n\tspec := NewWeightSpec(\"my-model\", \"model.safetensors\", \"/srv/weights/model.safetensors\")\n\n\tartifact, err := wb.Build(context.Background(), spec)\n\trequire.NoError(t, err)\n\trequire.NotNil(t, artifact)\n\n\t// Type assertion: should be a *WeightArtifact\n\twa, ok := artifact.(*WeightArtifact)\n\trequire.True(t, ok, \"expected *WeightArtifact, got %T\", artifact)\n\n\t// Check artifact interface\n\trequire.Equal(t, ArtifactTypeWeight, wa.Type())\n\trequire.Equal(t, \"my-model\", wa.Name())\n\n\t// Check descriptor\n\tdesc := wa.Descriptor()\n\trequire.Equal(t, expectedDigest, desc.Digest.String())\n\trequire.Equal(t, int64(len(weightContent)), desc.Size)\n\n\t// Check weight-specific fields\n\trequire.Equal(t, weightFile, wa.FilePath)\n\trequire.Equal(t, \"/srv/weights/model.safetensors\", wa.Target)\n\n\t// Check WeightConfig\n\trequire.Equal(t, \"1.0\", wa.Config.SchemaVersion)\n\trequire.Equal(t, \"0.15.0\", wa.Config.CogVersion)\n\trequire.Equal(t, \"my-model\", wa.Config.Name)\n\trequire.Equal(t, \"/srv/weights/model.safetensors\", wa.Config.Target)\n\trequire.False(t, wa.Config.Created.IsZero(), \"Created should be set\")\n}\n\nfunc TestWeightBuilder_WritesLockfile(t *testing.T) {\n\t// After Build(), a weights.lock should be written/updated at lockPath.\n\ttmpDir := t.TempDir()\n\tweightContent := []byte(\"lockfile test weight\")\n\terr := os.WriteFile(filepath.Join(tmpDir, \"model.bin\"), weightContent, 0o644)\n\trequire.NoError(t, err)\n\n\thash := sha256.Sum256(weightContent)\n\texpectedDigest := \"sha256:\" + hex.EncodeToString(hash[:])\n\n\tsrc := NewSourceFromConfig(&config.Config{\n\t\tWeights: []config.WeightSource{\n\t\t\t{Name: \"my-model\", Source: \"model.bin\", Target: \"/weights/model.bin\"},\n\t\t},\n\t}, tmpDir)\n\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\twb := NewWeightBuilder(src, \"0.15.0\", lockPath)\n\n\tspec := NewWeightSpec(\"my-model\", \"model.bin\", \"/weights/model.bin\")\n\t_, err = wb.Build(context.Background(), spec)\n\trequire.NoError(t, err)\n\n\t// Lockfile should exist\n\t_, err = os.Stat(lockPath)\n\trequire.NoError(t, err, \"lockfile should be created\")\n\n\t// Load and verify lockfile contents\n\tlock, err := LoadWeightsLock(lockPath)\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"1.0\", lock.Version)\n\trequire.Len(t, lock.Files, 1)\n\n\twf := lock.Files[0]\n\trequire.Equal(t, \"my-model\", wf.Name)\n\trequire.Equal(t, \"/weights/model.bin\", wf.Dest)\n\trequire.Equal(t, expectedDigest, wf.Digest)\n\trequire.Equal(t, int64(len(weightContent)), wf.Size)\n}\n\nfunc TestWeightBuilder_UpdatesExistingLockfile(t *testing.T) {\n\t// If a lockfile already exists with entries, Build() should add/update the entry\n\t// for the built weight without losing other entries.\n\ttmpDir := t.TempDir()\n\n\t// Create two weight files\n\tcontent1 := []byte(\"weight one data\")\n\tcontent2 := []byte(\"weight two data\")\n\terr := os.WriteFile(filepath.Join(tmpDir, \"w1.bin\"), content1, 0o644)\n\trequire.NoError(t, err)\n\terr = os.WriteFile(filepath.Join(tmpDir, \"w2.bin\"), content2, 0o644)\n\trequire.NoError(t, err)\n\n\tsrc := NewSourceFromConfig(&config.Config{\n\t\tWeights: []config.WeightSource{\n\t\t\t{Name: \"weight-1\", Source: \"w1.bin\", Target: \"/weights/w1.bin\"},\n\t\t\t{Name: \"weight-2\", Source: \"w2.bin\", Target: \"/weights/w2.bin\"},\n\t\t},\n\t}, tmpDir)\n\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\twb := NewWeightBuilder(src, \"0.15.0\", lockPath)\n\n\t// Build first weight\n\tspec1 := NewWeightSpec(\"weight-1\", \"w1.bin\", \"/weights/w1.bin\")\n\t_, err = wb.Build(context.Background(), spec1)\n\trequire.NoError(t, err)\n\n\t// Build second weight\n\tspec2 := NewWeightSpec(\"weight-2\", \"w2.bin\", \"/weights/w2.bin\")\n\t_, err = wb.Build(context.Background(), spec2)\n\trequire.NoError(t, err)\n\n\t// Lockfile should contain both entries\n\tlock, err := LoadWeightsLock(lockPath)\n\trequire.NoError(t, err)\n\trequire.Len(t, lock.Files, 2)\n\n\tnames := map[string]bool{}\n\tfor _, f := range lock.Files {\n\t\tnames[f.Name] = true\n\t}\n\trequire.True(t, names[\"weight-1\"])\n\trequire.True(t, names[\"weight-2\"])\n}\n\nfunc TestWeightBuilder_CacheHit(t *testing.T) {\n\t// When a lockfile entry exists with matching name and size,\n\t// the builder should use the cached digest without re-hashing.\n\ttmpDir := t.TempDir()\n\tweightContent := []byte(\"cached weight data\")\n\terr := os.WriteFile(filepath.Join(tmpDir, \"model.bin\"), weightContent, 0o644)\n\trequire.NoError(t, err)\n\n\thash := sha256.Sum256(weightContent)\n\texpectedDigest := \"sha256:\" + hex.EncodeToString(hash[:])\n\n\tsrc := NewSourceFromConfig(&config.Config{\n\t\tWeights: []config.WeightSource{\n\t\t\t{Name: \"my-model\", Source: \"model.bin\", Target: \"/weights/model.bin\"},\n\t\t},\n\t}, tmpDir)\n\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\twb := NewWeightBuilder(src, \"0.15.0\", lockPath)\n\n\t// First build — populates lockfile\n\tspec := NewWeightSpec(\"my-model\", \"model.bin\", \"/weights/model.bin\")\n\tartifact1, err := wb.Build(context.Background(), spec)\n\trequire.NoError(t, err)\n\n\t// Second build — should hit cache\n\tartifact2, err := wb.Build(context.Background(), spec)\n\trequire.NoError(t, err)\n\n\t// Both builds should produce the same digest\n\twa1 := artifact1.(*WeightArtifact)\n\twa2 := artifact2.(*WeightArtifact)\n\trequire.Equal(t, expectedDigest, wa1.Descriptor().Digest.String())\n\trequire.Equal(t, expectedDigest, wa2.Descriptor().Digest.String())\n\n\t// Lockfile should still have exactly one entry (not duplicated)\n\tlock, err := LoadWeightsLock(lockPath)\n\trequire.NoError(t, err)\n\trequire.Len(t, lock.Files, 1)\n\trequire.Equal(t, \"my-model\", lock.Files[0].Name)\n}\n\nfunc TestWeightBuilder_CacheMiss_SizeChanged(t *testing.T) {\n\t// When the file size changes, the builder should re-hash.\n\ttmpDir := t.TempDir()\n\tcontent1 := []byte(\"original content\")\n\terr := os.WriteFile(filepath.Join(tmpDir, \"model.bin\"), content1, 0o644)\n\trequire.NoError(t, err)\n\n\tsrc := NewSourceFromConfig(&config.Config{\n\t\tWeights: []config.WeightSource{\n\t\t\t{Name: \"my-model\", Source: \"model.bin\", Target: \"/weights/model.bin\"},\n\t\t},\n\t}, tmpDir)\n\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\twb := NewWeightBuilder(src, \"0.15.0\", lockPath)\n\n\tspec := NewWeightSpec(\"my-model\", \"model.bin\", \"/weights/model.bin\")\n\n\t// First build\n\t_, err = wb.Build(context.Background(), spec)\n\trequire.NoError(t, err)\n\n\t// Change the file (different size)\n\tcontent2 := []byte(\"modified content with different size!!\")\n\terr = os.WriteFile(filepath.Join(tmpDir, \"model.bin\"), content2, 0o644)\n\trequire.NoError(t, err)\n\n\t// Second build — should detect size change and re-hash\n\tartifact2, err := wb.Build(context.Background(), spec)\n\trequire.NoError(t, err)\n\n\twa2 := artifact2.(*WeightArtifact)\n\thash2 := sha256.Sum256(content2)\n\texpectedDigest2 := \"sha256:\" + hex.EncodeToString(hash2[:])\n\trequire.Equal(t, expectedDigest2, wa2.Descriptor().Digest.String())\n\trequire.Equal(t, int64(len(content2)), wa2.Descriptor().Size)\n}\n\nfunc TestWeightBuilder_ErrorWrongSpecType(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tsrc := NewSourceFromConfig(&config.Config{}, tmpDir)\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\twb := NewWeightBuilder(src, \"0.15.0\", lockPath)\n\n\t// Pass an ImageSpec instead of WeightSpec\n\timageSpec := NewImageSpec(\"model\", \"test-image\")\n\t_, err := wb.Build(context.Background(), imageSpec)\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"expected *WeightSpec\")\n}\n\nfunc TestWeightBuilder_ErrorFileNotFound(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tsrc := NewSourceFromConfig(&config.Config{}, tmpDir)\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\twb := NewWeightBuilder(src, \"0.15.0\", lockPath)\n\n\tspec := NewWeightSpec(\"missing\", \"nonexistent.bin\", \"/weights/nonexistent.bin\")\n\t_, err := wb.Build(context.Background(), spec)\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"weight source not found\")\n}\n\nfunc TestWeightBuilder_ErrorContextCancelled(t *testing.T) {\n\ttmpDir := t.TempDir()\n\terr := os.WriteFile(filepath.Join(tmpDir, \"model.bin\"), []byte(\"data\"), 0o644)\n\trequire.NoError(t, err)\n\n\tsrc := NewSourceFromConfig(&config.Config{}, tmpDir)\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\twb := NewWeightBuilder(src, \"0.15.0\", lockPath)\n\n\tctx, cancel := context.WithCancel(context.Background())\n\tcancel() // Cancel immediately\n\n\tspec := NewWeightSpec(\"model\", \"model.bin\", \"/weights/model.bin\")\n\t_, err = wb.Build(ctx, spec)\n\trequire.Error(t, err)\n\trequire.ErrorIs(t, err, context.Canceled)\n}\n\nfunc TestWeightBuilder_ImplementsBuilderInterface(t *testing.T) {\n\ttmpDir := t.TempDir()\n\tsrc := NewSourceFromConfig(&config.Config{}, tmpDir)\n\tlockPath := filepath.Join(tmpDir, \"weights.lock\")\n\n\t// Compile-time check\n\tvar _ Builder = NewWeightBuilder(src, \"0.1.0\", lockPath)\n}\n"
  },
  {
    "path": "pkg/model/weight_pusher.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"crypto/sha256\"\n\t\"encoding/hex\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"strings\"\n\t\"sync\"\n\t\"time\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/empty\"\n\t\"github.com/google/go-containerregistry/pkg/v1/mutate\"\n\t\"github.com/google/go-containerregistry/pkg/v1/tarball\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\n// WeightPushOptions configures optional behavior for WeightPusher.Push.\ntype WeightPushOptions struct {\n\t// ProgressFn is an optional callback for reporting upload progress.\n\tProgressFn func(PushProgress)\n\t// RetryFn is an optional callback for reporting retry attempts.\n\t// Return false to abort the retry.\n\tRetryFn func(WeightRetryEvent) bool\n}\n\n// WeightRetryEvent reports a retry attempt for a weight file upload.\ntype WeightRetryEvent struct {\n\t// Name identifies which file is being retried.\n\tName string\n\t// Attempt is the current retry attempt number (1-indexed).\n\tAttempt int\n\t// MaxAttempts is the maximum number of retry attempts.\n\tMaxAttempts int\n\t// Err is the error that caused the retry.\n\tErr error\n\t// NextRetryIn is the duration until the next retry attempt.\n\tNextRetryIn time.Duration\n}\n\n// WeightPushResult contains the result of pushing a single weight artifact.\ntype WeightPushResult struct {\n\t// Ref is the full image reference for the pushed weight manifest (e.g., \"registry/repo:weights-name-abc123\").\n\tRef string\n\t// Descriptor is the OCI descriptor for the pushed weight manifest.\n\tDescriptor v1.Descriptor\n}\n\n// WeightPusher pushes a WeightArtifact as a proper OCI artifact manifest\n// with config blob and tarball layers. The layer blob is pushed via\n// registry.WriteLayer (which supports multipart uploads, progress, and retry),\n// followed by the manifest via PushImage.\ntype WeightPusher struct {\n\tregistry registry.Client\n}\n\n// NewWeightPusher creates a new WeightPusher.\nfunc NewWeightPusher(reg registry.Client) *WeightPusher {\n\treturn &WeightPusher{registry: reg}\n}\n\n// Push pushes a WeightArtifact to the registry as an OCI artifact manifest.\n// The layer blob is pushed first via WriteLayer (multipart uploads, progress, retry),\n// then the manifest is pushed via PushImage.\n// Returns the descriptor of the pushed manifest.\nfunc (p *WeightPusher) Push(ctx context.Context, repo string, artifact *WeightArtifact, opts ...WeightPushOptions) (*WeightPushResult, error) {\n\tif artifact == nil {\n\t\treturn nil, fmt.Errorf(\"artifact is nil\")\n\t}\n\tif repo == \"\" {\n\t\treturn nil, fmt.Errorf(\"repo is required\")\n\t}\n\n\t// Merge options (use first if provided)\n\tvar opt WeightPushOptions\n\tif len(opts) > 0 {\n\t\topt = opts[0]\n\t}\n\n\t// Verify the weight file exists\n\tif _, err := os.Stat(artifact.FilePath); err != nil {\n\t\treturn nil, fmt.Errorf(\"weight file %q: %w\", artifact.FilePath, err)\n\t}\n\n\t// Build the OCI artifact image (config blob + tarball layer)\n\timg, err := buildWeightImage(artifact)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"build weight image: %w\", err)\n\t}\n\n\t// Extract the layer to push via WriteLayer (gets multipart + progress + retry)\n\tlayers, err := img.Layers()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"get image layers: %w\", err)\n\t}\n\tif len(layers) != 1 {\n\t\treturn nil, fmt.Errorf(\"expected 1 layer, got %d\", len(layers))\n\t}\n\tlayer := layers[0]\n\n\t// Build progress callback\n\tvar onProgress func(v1.Update)\n\tif opt.ProgressFn != nil {\n\t\tonProgress = func(update v1.Update) {\n\t\t\topt.ProgressFn(PushProgress{\n\t\t\t\tComplete: update.Complete,\n\t\t\t\tTotal:    update.Total,\n\t\t\t})\n\t\t}\n\t}\n\n\t// Build retry configuration if callback is provided\n\tvar retryConfig *registry.RetryConfig\n\tif opt.RetryFn != nil {\n\t\tretryConfig = ®istry.RetryConfig{\n\t\t\tOnRetry: func(event registry.RetryEvent) bool {\n\t\t\t\treturn opt.RetryFn(WeightRetryEvent{\n\t\t\t\t\tName:        artifact.Name(),\n\t\t\t\t\tAttempt:     event.Attempt,\n\t\t\t\t\tMaxAttempts: event.MaxAttempts,\n\t\t\t\t\tErr:         event.Err,\n\t\t\t\t\tNextRetryIn: event.NextRetryIn,\n\t\t\t\t})\n\t\t\t},\n\t\t}\n\t}\n\n\t// 1. Push layer blob via WriteLayer (multipart uploads, progress, retry)\n\twriteErr := writeLayerWithProgress(ctx, p.registry, registry.WriteLayerOptions{\n\t\tRepo:  repo,\n\t\tLayer: layer,\n\t\tRetry: retryConfig,\n\t}, onProgress)\n\n\tif writeErr != nil {\n\t\treturn nil, fmt.Errorf(\"push weight layer: %w\", writeErr)\n\t}\n\n\t// 2. Push manifest via PushImage with a single tag combining name and digest.\n\t// The layer blob is already in the registry, so PushImage will skip re-uploading it.\n\t// Tag format: :weights--<12chars> (e.g., :weights-model-v1-383d1f4afa43)\n\t//\n\t// We use the artifact's descriptor digest (original file hash from the lock file),\n\t// NOT the tarball layer digest. This ensures that `weights inspect` can look up the tag\n\t// using the same digest stored in weights.lock, independent of the transport format.\n\ttag := WeightTag(artifact.Name(), artifact.Descriptor().Digest.String())\n\tref := repo + \":\" + tag\n\tif err := p.registry.PushImage(ctx, ref, img); err != nil {\n\t\treturn nil, fmt.Errorf(\"push weight manifest (%s): %w\", tag, err)\n\t}\n\n\t// Build result descriptor from the pushed image\n\tdesc, err := descriptorFromImage(img)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"compute manifest descriptor: %w\", err)\n\t}\n\n\treturn &WeightPushResult{Ref: ref, Descriptor: desc}, nil\n}\n\n// buildWeightImage creates an OCI artifact image with a config blob (WeightConfig JSON)\n// and a tarball layer for the weight file.\nfunc buildWeightImage(artifact *WeightArtifact) (v1.Image, error) {\n\t// 1. Create the base image with OCI manifest media type\n\timg := mutate.MediaType(empty.Image, types.OCIManifestSchema1)\n\n\t// 2. Create tarball layer from the weight file.\n\t// WithCompressedCaching memoizes the compressed output so that Digest() and\n\t// Compressed() see identical bytes. Without this, gzip non-determinism between\n\t// separate passes causes DIGEST_INVALID errors on large uploads.\n\tlayer, err := tarball.LayerFromFile(artifact.FilePath,\n\t\ttarball.WithMediaType(types.MediaType(MediaTypeWeightLayer)),\n\t\ttarball.WithCompressedCaching,\n\t)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"create tarball layer: %w\", err)\n\t}\n\n\t// 3. Append the layer\n\timg, err = mutate.AppendLayers(img, layer)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"append weight layer: %w\", err)\n\t}\n\n\t// 4. Serialize the WeightConfig as the config blob\n\tconfigJSON, err := json.Marshal(artifact.Config)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"marshal weight config: %w\", err)\n\t}\n\n\t// 5. Wrap to set custom config blob, config media type, and artifactType\n\treturn &weightManifestImage{\n\t\tImage:      img,\n\t\tconfigBlob: configJSON,\n\t}, nil\n}\n\n// descriptorFromImage computes the v1.Descriptor for a built image manifest.\nfunc descriptorFromImage(img v1.Image) (v1.Descriptor, error) {\n\tdigest, err := img.Digest()\n\tif err != nil {\n\t\treturn v1.Descriptor{}, fmt.Errorf(\"get digest: %w\", err)\n\t}\n\n\trawManifest, err := img.RawManifest()\n\tif err != nil {\n\t\treturn v1.Descriptor{}, fmt.Errorf(\"get raw manifest: %w\", err)\n\t}\n\n\treturn v1.Descriptor{\n\t\tMediaType: types.OCIManifestSchema1,\n\t\tSize:      int64(len(rawManifest)),\n\t\tDigest:    digest,\n\t}, nil\n}\n\n// weightOCIManifest extends v1.Manifest with artifactType for OCI 1.1 support.\n// go-containerregistry v0.20.5's v1.Manifest struct does not include artifactType,\n// so we serialize it ourselves.\ntype weightOCIManifest struct {\n\tSchemaVersion int64             `json:\"schemaVersion\"`\n\tMediaType     types.MediaType   `json:\"mediaType,omitempty\"`\n\tConfig        v1.Descriptor     `json:\"config\"`\n\tLayers        []v1.Descriptor   `json:\"layers\"`\n\tAnnotations   map[string]string `json:\"annotations,omitempty\"`\n\tArtifactType  string            `json:\"artifactType,omitempty\"`\n}\n\n// weightManifestImage wraps a v1.Image to set a custom config blob with\n// the correct media type and artifactType. This produces a proper OCI 1.1\n// artifact manifest for weight data.\n//\n// The raw manifest is cached on first computation to ensure deterministic\n// digests across multiple calls (e.g., during remote.Write which calls\n// both RawManifest and Digest).\ntype weightManifestImage struct {\n\tv1.Image\n\tconfigBlob     []byte\n\trawManifest    []byte\n\trawManifestErr error\n\trawOnce        sync.Once\n}\n\n// RawConfigFile returns the WeightConfig JSON as the config blob.\nfunc (w *weightManifestImage) RawConfigFile() ([]byte, error) {\n\treturn w.configBlob, nil\n}\n\n// Digest computes the digest from the cached raw manifest.\nfunc (w *weightManifestImage) Digest() (v1.Hash, error) {\n\traw, err := w.RawManifest()\n\tif err != nil {\n\t\treturn v1.Hash{}, err\n\t}\n\th := sha256.Sum256(raw)\n\treturn v1.Hash{\n\t\tAlgorithm: \"sha256\",\n\t\tHex:       hex.EncodeToString(h[:]),\n\t}, nil\n}\n\n// ArtifactType implements the withArtifactType interface used by partial.Descriptor.\nfunc (w *weightManifestImage) ArtifactType() (string, error) {\n\treturn MediaTypeWeightArtifact, nil\n}\n\n// Manifest returns the modified manifest with custom config descriptor.\nfunc (w *weightManifestImage) Manifest() (*v1.Manifest, error) {\n\tm, err := w.Image.Manifest()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\t// Make a copy to avoid mutating the original\n\tmCopy := m.DeepCopy()\n\n\t// Set config to point to our custom config blob\n\tconfigDigest := sha256.Sum256(w.configBlob)\n\tmCopy.Config = v1.Descriptor{\n\t\tMediaType: types.MediaType(MediaTypeWeightConfig),\n\t\tSize:      int64(len(w.configBlob)),\n\t\tDigest: v1.Hash{\n\t\t\tAlgorithm: \"sha256\",\n\t\t\tHex:       hex.EncodeToString(configDigest[:]),\n\t\t},\n\t}\n\n\treturn mCopy, nil\n}\n\n// RawManifest serializes our modified manifest with artifactType field.\n// The result is cached to ensure deterministic digests across multiple calls.\nfunc (w *weightManifestImage) RawManifest() ([]byte, error) {\n\tw.rawOnce.Do(func() {\n\t\tm, err := w.Manifest()\n\t\tif err != nil {\n\t\t\tw.rawManifestErr = err\n\t\t\treturn\n\t\t}\n\n\t\t// Build the OCI manifest with artifactType (not in v1.Manifest struct)\n\t\tociManifest := weightOCIManifest{\n\t\t\tSchemaVersion: m.SchemaVersion,\n\t\t\tMediaType:     m.MediaType,\n\t\t\tConfig:        m.Config,\n\t\t\tLayers:        m.Layers,\n\t\t\tAnnotations:   m.Annotations,\n\t\t\tArtifactType:  MediaTypeWeightArtifact,\n\t\t}\n\n\t\tw.rawManifest, w.rawManifestErr = json.Marshal(ociManifest)\n\t})\n\n\treturn w.rawManifest, w.rawManifestErr\n}\n\n// =============================================================================\n// Weight tag helpers\n// =============================================================================\n\nconst weightTagPrefix = \"weights-\"\n\n// WeightTag returns the tag for a weight manifest combining name and digest.\n// The digest should be in \"sha256:abc123...\" format.\n// Returns e.g., \"weights-model-v1-abc123def456\" (12-char hex suffix).\n// Falls back to \"weights-\" if digest is empty or invalid.\nfunc WeightTag(name, digest string) string {\n\t_, hex, ok := strings.Cut(digest, \":\")\n\tif !ok || hex == \"\" {\n\t\treturn weightTagPrefix + name\n\t}\n\tshort := hex\n\tif len(short) > 12 {\n\t\tshort = short[:12]\n\t}\n\treturn weightTagPrefix + name + \"-\" + short\n}\n"
  },
  {
    "path": "pkg/model/weight_pusher_test.go",
    "content": "package model\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"sync\"\n\t\"testing\"\n\t\"time\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\nfunc TestWeightPusher_Push_ReturnsErrorForNilArtifact(t *testing.T) {\n\treg := &mockRegistry{}\n\tpusher := NewWeightPusher(reg)\n\n\t_, err := pusher.Push(context.Background(), \"r8.im/user/model\", nil)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"artifact is nil\")\n}\n\nfunc TestWeightPusher_Push_ReturnsErrorForMissingFile(t *testing.T) {\n\treg := &mockRegistry{}\n\tpusher := NewWeightPusher(reg)\n\n\tartifact := NewWeightArtifact(\"model-v1\", v1.Descriptor{}, \"/nonexistent/path/weights.bin\", \"/weights/model.bin\", WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Now().UTC(),\n\t})\n\n\t_, err := pusher.Push(context.Background(), \"r8.im/user/model\", artifact)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"weight file\")\n}\n\nfunc TestWeightPusher_Push_PushesCorrectOCIArtifact(t *testing.T) {\n\t// Create a temp weight file\n\tdir := t.TempDir()\n\tweightPath := filepath.Join(dir, \"model.safetensors\")\n\tweightContent := []byte(\"fake weight data for testing tarball layer creation\")\n\trequire.NoError(t, os.WriteFile(weightPath, weightContent, 0o644))\n\n\tcreated := time.Date(2026, 2, 5, 12, 0, 0, 0, time.UTC)\n\tcfg := WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.safetensors\",\n\t\tCreated:       created,\n\t}\n\n\tdesc := v1.Descriptor{\n\t\tDigest: v1.Hash{Algorithm: \"sha256\", Hex: \"aabbccddee112233445566778899aabb00112233445566778899aabbccddeeff\"},\n\t}\n\tartifact := NewWeightArtifact(\"model-v1\", desc, weightPath, \"/weights/model.safetensors\", cfg)\n\n\t// Capture what gets pushed\n\tvar pushedRefs []string\n\tvar pushedImg v1.Image\n\treg := &mockRegistry{\n\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\tpushedRefs = append(pushedRefs, ref)\n\t\t\tpushedImg = img\n\t\t\treturn nil\n\t\t},\n\t}\n\n\tpusher := NewWeightPusher(reg)\n\tresult, err := pusher.Push(context.Background(), \"r8.im/user/model\", artifact)\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, result)\n\n\t// Verify the image was pushed with a single combined tag\n\trequire.Len(t, pushedRefs, 1)\n\trequire.Equal(t, \"r8.im/user/model:weights-model-v1-aabbccddee11\", pushedRefs[0])\n\trequire.NotNil(t, pushedImg)\n\n\t// Verify manifest structure\n\tmanifest, err := pushedImg.Manifest()\n\trequire.NoError(t, err)\n\trequire.Equal(t, types.OCIManifestSchema1, manifest.MediaType)\n\n\t// Verify config blob has correct media type\n\trequire.Equal(t, types.MediaType(MediaTypeWeightConfig), manifest.Config.MediaType)\n\n\t// Verify config blob content is correct WeightConfig JSON\n\tconfigBlob, err := pushedImg.RawConfigFile()\n\trequire.NoError(t, err)\n\tvar parsedConfig WeightConfig\n\trequire.NoError(t, json.Unmarshal(configBlob, &parsedConfig))\n\trequire.Equal(t, \"1.0\", parsedConfig.SchemaVersion)\n\trequire.Equal(t, \"0.15.0\", parsedConfig.CogVersion)\n\trequire.Equal(t, \"model-v1\", parsedConfig.Name)\n\trequire.Equal(t, \"/weights/model.safetensors\", parsedConfig.Target)\n\trequire.Equal(t, created, parsedConfig.Created)\n\n\t// Verify there's exactly one layer (single file = single layer)\n\trequire.Len(t, manifest.Layers, 1)\n\n\t// Verify layer media type\n\trequire.Equal(t, types.MediaType(MediaTypeWeightLayer), manifest.Layers[0].MediaType)\n\n\t// Verify layer size matches the tarball wrapping of the weight file\n\t// (tarball will be larger than raw content due to tar headers)\n\trequire.Greater(t, manifest.Layers[0].Size, int64(0))\n\n\t// Verify the result contains a valid descriptor\n\trequire.NotEmpty(t, result.Descriptor.Digest.String())\n\trequire.Greater(t, result.Descriptor.Size, int64(0))\n}\n\nfunc TestWeightPusher_Push_PropagatesPushError(t *testing.T) {\n\tdir := t.TempDir()\n\tweightPath := filepath.Join(dir, \"model.bin\")\n\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\tartifact := NewWeightArtifact(\"model-v1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Now().UTC(),\n\t})\n\n\treg := &mockRegistry{\n\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\treturn fmt.Errorf(\"unauthorized: authentication required\")\n\t\t},\n\t}\n\n\tpusher := NewWeightPusher(reg)\n\t_, err := pusher.Push(context.Background(), \"r8.im/user/model\", artifact)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"push weight manifest\")\n\trequire.Contains(t, err.Error(), \"unauthorized\")\n}\n\nfunc TestWeightPusher_Push_RawManifestContainsArtifactType(t *testing.T) {\n\tdir := t.TempDir()\n\tweightPath := filepath.Join(dir, \"model.bin\")\n\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test weight data\"), 0o644))\n\n\tartifact := NewWeightArtifact(\"model-v1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Date(2026, 2, 5, 12, 0, 0, 0, time.UTC),\n\t})\n\n\tvar pushedImg v1.Image\n\treg := &mockRegistry{\n\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\tpushedImg = img\n\t\t\treturn nil\n\t\t},\n\t}\n\n\tpusher := NewWeightPusher(reg)\n\t_, err := pusher.Push(context.Background(), \"r8.im/user/model\", artifact)\n\trequire.NoError(t, err)\n\n\t// Parse raw manifest JSON to verify artifactType field\n\trawManifest, err := pushedImg.RawManifest()\n\trequire.NoError(t, err)\n\n\tvar manifestJSON map[string]any\n\trequire.NoError(t, json.Unmarshal(rawManifest, &manifestJSON))\n\n\t// artifactType must be present at the manifest level (OCI 1.1)\n\trequire.Equal(t, MediaTypeWeightArtifact, manifestJSON[\"artifactType\"])\n\n\t// config.mediaType must be the weight config type\n\tconfigMap, ok := manifestJSON[\"config\"].(map[string]any)\n\trequire.True(t, ok, \"config should be an object\")\n\trequire.Equal(t, MediaTypeWeightConfig, configMap[\"mediaType\"])\n\n\t// layers should have exactly one entry with the weight layer media type\n\tlayers, ok := manifestJSON[\"layers\"].([]any)\n\trequire.True(t, ok, \"layers should be an array\")\n\trequire.Len(t, layers, 1)\n\n\tlayerMap, ok := layers[0].(map[string]any)\n\trequire.True(t, ok, \"layer should be an object\")\n\trequire.Equal(t, MediaTypeWeightLayer, layerMap[\"mediaType\"])\n}\n\nfunc TestWeightPusher_Push_ReturnsErrorForEmptyRepo(t *testing.T) {\n\tdir := t.TempDir()\n\tweightPath := filepath.Join(dir, \"model.bin\")\n\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\tartifact := NewWeightArtifact(\"model-v1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Now().UTC(),\n\t})\n\n\treg := &mockRegistry{}\n\tpusher := NewWeightPusher(reg)\n\n\t_, err := pusher.Push(context.Background(), \"\", artifact)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"repo is required\")\n}\n\nfunc TestWeightPusher_Push_ReportsProgressViaWriteLayer(t *testing.T) {\n\tdir := t.TempDir()\n\tweightPath := filepath.Join(dir, \"model.bin\")\n\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test weight data for progress tracking\"), 0o644))\n\n\tartifact := NewWeightArtifact(\"model-v1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Now().UTC(),\n\t})\n\n\t// Track progress updates received via callback\n\tvar (\n\t\tmu       sync.Mutex\n\t\tprogress []PushProgress\n\t)\n\n\t// Mock WriteLayer to simulate progress updates (caller owns closing the channel)\n\treg := &mockRegistry{\n\t\twriteLayerFunc: func(ctx context.Context, opts registry.WriteLayerOptions) error {\n\t\t\t// Simulate progress updates like the real registry client\n\t\t\tif opts.ProgressCh != nil {\n\t\t\t\topts.ProgressCh <- v1.Update{Complete: 500, Total: 1000}\n\t\t\t\topts.ProgressCh <- v1.Update{Complete: 1000, Total: 1000}\n\t\t\t}\n\t\t\treturn nil\n\t\t},\n\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\treturn nil\n\t\t},\n\t}\n\n\tpusher := NewWeightPusher(reg)\n\tresult, err := pusher.Push(context.Background(), \"r8.im/user/model\", artifact, WeightPushOptions{\n\t\tProgressFn: func(p PushProgress) {\n\t\t\tmu.Lock()\n\t\t\tdefer mu.Unlock()\n\t\t\tprogress = append(progress, p)\n\t\t},\n\t})\n\n\trequire.NoError(t, err)\n\trequire.NotNil(t, result)\n\n\t// Verify we received progress updates\n\tmu.Lock()\n\tdefer mu.Unlock()\n\trequire.GreaterOrEqual(t, len(progress), 2, \"should receive at least 2 progress updates\")\n\n\t// Verify progress updates contain expected values\n\trequire.Equal(t, int64(500), progress[0].Complete)\n\trequire.Equal(t, int64(1000), progress[0].Total)\n\trequire.Equal(t, int64(1000), progress[1].Complete)\n\trequire.Equal(t, int64(1000), progress[1].Total)\n}\n\nfunc TestWeightPusher_Push_ForwardsRetryCallback(t *testing.T) {\n\tdir := t.TempDir()\n\tweightPath := filepath.Join(dir, \"model.bin\")\n\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test weight data\"), 0o644))\n\n\tartifact := NewWeightArtifact(\"model-v1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Now().UTC(),\n\t})\n\n\t// Mock WriteLayer to capture the retry config and invoke it\n\tvar retryEvents []WeightRetryEvent\n\treg := &mockRegistry{\n\t\twriteLayerFunc: func(ctx context.Context, opts registry.WriteLayerOptions) error {\n\t\t\t// Simulate the registry invoking the retry callback\n\t\t\tif opts.Retry != nil && opts.Retry.OnRetry != nil {\n\t\t\t\topts.Retry.OnRetry(registry.RetryEvent{\n\t\t\t\t\tAttempt:     1,\n\t\t\t\t\tMaxAttempts: 3,\n\t\t\t\t\tErr:         fmt.Errorf(\"connection reset\"),\n\t\t\t\t\tNextRetryIn: 2 * time.Second,\n\t\t\t\t})\n\t\t\t}\n\t\t\treturn nil\n\t\t},\n\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\treturn nil\n\t\t},\n\t}\n\n\tpusher := NewWeightPusher(reg)\n\t_, err := pusher.Push(context.Background(), \"r8.im/user/model\", artifact, WeightPushOptions{\n\t\tRetryFn: func(event WeightRetryEvent) bool {\n\t\t\tretryEvents = append(retryEvents, event)\n\t\t\treturn true\n\t\t},\n\t})\n\n\trequire.NoError(t, err)\n\trequire.Len(t, retryEvents, 1)\n\trequire.Equal(t, \"model-v1\", retryEvents[0].Name)\n\trequire.Equal(t, 1, retryEvents[0].Attempt)\n\trequire.Equal(t, 3, retryEvents[0].MaxAttempts)\n\trequire.Contains(t, retryEvents[0].Err.Error(), \"connection reset\")\n\trequire.Equal(t, 2*time.Second, retryEvents[0].NextRetryIn)\n}\n\nfunc TestWeightPusher_Push_WriteLayerErrorReported(t *testing.T) {\n\tdir := t.TempDir()\n\tweightPath := filepath.Join(dir, \"model.bin\")\n\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\tartifact := NewWeightArtifact(\"model-v1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Now().UTC(),\n\t})\n\n\treg := &mockRegistry{\n\t\twriteLayerFunc: func(ctx context.Context, opts registry.WriteLayerOptions) error {\n\t\t\treturn fmt.Errorf(\"upload failed: 503 Service Unavailable\")\n\t\t},\n\t}\n\n\tpusher := NewWeightPusher(reg)\n\t_, err := pusher.Push(context.Background(), \"r8.im/user/model\", artifact)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"push weight layer\")\n\trequire.Contains(t, err.Error(), \"503 Service Unavailable\")\n}\n\nfunc TestWeightPusher_Push_PropagatesContextCancellation(t *testing.T) {\n\tdir := t.TempDir()\n\tweightPath := filepath.Join(dir, \"model.bin\")\n\trequire.NoError(t, os.WriteFile(weightPath, []byte(\"test\"), 0o644))\n\n\tartifact := NewWeightArtifact(\"model-v1\", v1.Descriptor{}, weightPath, \"/weights/model.bin\", WeightConfig{\n\t\tSchemaVersion: \"1.0\",\n\t\tCogVersion:    \"0.15.0\",\n\t\tName:          \"model-v1\",\n\t\tTarget:        \"/weights/model.bin\",\n\t\tCreated:       time.Now().UTC(),\n\t})\n\n\tctx, cancel := context.WithCancel(context.Background())\n\tcancel() // Cancel immediately\n\n\treg := &mockRegistry{\n\t\tpushImageFunc: func(ctx context.Context, ref string, img v1.Image) error {\n\t\t\treturn ctx.Err()\n\t\t},\n\t}\n\n\tpusher := NewWeightPusher(reg)\n\t_, err := pusher.Push(ctx, \"r8.im/user/model\", artifact)\n\n\trequire.Error(t, err)\n\trequire.Contains(t, err.Error(), \"context canceled\")\n}\n"
  },
  {
    "path": "pkg/model/weights.go",
    "content": "package model\n\n// WeightFile represents a single weight file entry in a weights lockfile or manifest.\n// The Name field is an identifier/handle (like a Docker volume name), not a filename.\ntype WeightFile struct {\n\t// Name is the identifier/handle for this weight (e.g., \"personaplex-7b-v1\", \"model-v42.5\").\n\t// This is a logical name that maps to deployment blob metadata, not a file path.\n\tName string `json:\"name\"`\n\t// Dest is the mount path in the container (e.g., /cache/model.safetensors).\n\tDest string `json:\"dest\"`\n\t// DigestOriginal is the SHA256 of the uncompressed file (canonical ID).\n\tDigestOriginal string `json:\"digestOriginal\"`\n\t// Digest is the SHA256 of the compressed blob (OCI layer ID).\n\tDigest string `json:\"digest\"`\n\t// Size is the compressed size in bytes.\n\tSize int64 `json:\"size\"`\n\t// SizeUncompressed is the original size in bytes.\n\tSizeUncompressed int64 `json:\"sizeUncompressed\"`\n\t// MediaType is the OCI layer media type (e.g., application/vnd.cog.weight.layer.v1+gzip).\n\tMediaType string `json:\"mediaType\"`\n\t// ContentType is the file's MIME type (e.g., application/octet-stream).\n\tContentType string `json:\"contentType,omitempty\"`\n}\n"
  },
  {
    "path": "pkg/model/weights_lock.go",
    "content": "// pkg/model/weights_lock.go\npackage model\n\nimport (\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"time\"\n)\n\n// WeightsLockFilename is the default filename for the weights lock file.\nconst WeightsLockFilename = \"weights.lock\"\n\n// WeightsLock represents a weights.lock file that pins weight file metadata.\n// This is a placeholder format that will be replaced by the declarative weights implementation.\ntype WeightsLock struct {\n\t// Version is the lockfile format version.\n\tVersion string `json:\"version\"`\n\t// Created is when the lockfile was generated.\n\tCreated time.Time `json:\"created\"`\n\t// Files are the weight file entries.\n\tFiles []WeightFile `json:\"files\"`\n}\n\n// ParseWeightsLock parses a weights.lock JSON document.\nfunc ParseWeightsLock(data []byte) (*WeightsLock, error) {\n\tvar lock WeightsLock\n\tif err := json.Unmarshal(data, &lock); err != nil {\n\t\treturn nil, fmt.Errorf(\"parse weights.lock: %w\", err)\n\t}\n\treturn &lock, nil\n}\n\n// LoadWeightsLock loads a weights.lock file from disk.\nfunc LoadWeightsLock(path string) (*WeightsLock, error) {\n\tdata, err := os.ReadFile(path)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"read weights.lock: %w\", err)\n\t}\n\treturn ParseWeightsLock(data)\n}\n\n// Save writes the weights.lock to disk.\nfunc (wl *WeightsLock) Save(path string) error {\n\tdata, err := json.MarshalIndent(wl, \"\", \"  \")\n\tif err != nil {\n\t\treturn fmt.Errorf(\"marshal weights.lock: %w\", err)\n\t}\n\tif err := os.WriteFile(path, data, 0o644); err != nil {\n\t\treturn fmt.Errorf(\"write weights.lock: %w\", err)\n\t}\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/model/weights_lock_test.go",
    "content": "// pkg/model/weights_lock_test.go\npackage model\n\nimport (\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestWeightsLock(t *testing.T) {\n\tt.Run(\"parse valid lockfile\", func(t *testing.T) {\n\t\tjson := `{\n\t\t\t\"version\": \"1\",\n\t\t\t\"created\": \"2026-01-30T12:00:00Z\",\n\t\t\t\"files\": [\n\t\t\t\t{\n\t\t\t\t\t\"name\": \"model.safetensors\",\n\t\t\t\t\t\"dest\": \"/cache/model.safetensors\",\n\t\t\t\t\t\"digestOriginal\": \"sha256:abc123\",\n\t\t\t\t\t\"digest\": \"sha256:def456\",\n\t\t\t\t\t\"size\": 1000,\n\t\t\t\t\t\"sizeUncompressed\": 2000,\n\t\t\t\t\t\"mediaType\": \"application/vnd.cog.weights.layer.v1+gzip\"\n\t\t\t\t}\n\t\t\t]\n\t\t}`\n\n\t\tlock, err := ParseWeightsLock([]byte(json))\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"1\", lock.Version)\n\t\trequire.Len(t, lock.Files, 1)\n\t\trequire.Equal(t, \"model.safetensors\", lock.Files[0].Name)\n\t\trequire.Equal(t, \"/cache/model.safetensors\", lock.Files[0].Dest)\n\t\trequire.Equal(t, \"sha256:abc123\", lock.Files[0].DigestOriginal)\n\t\trequire.Equal(t, \"sha256:def456\", lock.Files[0].Digest)\n\t\trequire.Equal(t, int64(1000), lock.Files[0].Size)\n\t})\n\n\tt.Run(\"load from file\", func(t *testing.T) {\n\t\tdir := t.TempDir()\n\t\tlockPath := filepath.Join(dir, \"weights.lock\")\n\t\tcontent := `{\"version\": \"1\", \"created\": \"2026-01-30T12:00:00Z\", \"files\": []}`\n\t\trequire.NoError(t, os.WriteFile(lockPath, []byte(content), 0o644))\n\n\t\tlock, err := LoadWeightsLock(lockPath)\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"1\", lock.Version)\n\t})\n\n\tt.Run(\"save to file\", func(t *testing.T) {\n\t\tdir := t.TempDir()\n\t\tlockPath := filepath.Join(dir, \"weights.lock\")\n\n\t\tlock := &WeightsLock{\n\t\t\tVersion: \"1\",\n\t\t\tCreated: time.Date(2026, 1, 30, 12, 0, 0, 0, time.UTC),\n\t\t\tFiles: []WeightFile{\n\t\t\t\t{Name: \"test.bin\", Dest: \"/cache/test.bin\"},\n\t\t\t},\n\t\t}\n\n\t\trequire.NoError(t, lock.Save(lockPath))\n\n\t\tloaded, err := LoadWeightsLock(lockPath)\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, lock.Version, loaded.Version)\n\t\trequire.Len(t, loaded.Files, 1)\n\t})\n\n}\n"
  },
  {
    "path": "pkg/model/weights_test.go",
    "content": "package model\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestWeightFile(t *testing.T) {\n\tt.Run(\"media type constants\", func(t *testing.T) {\n\t\trequire.Equal(t, \"application/vnd.cog.weight.layer.v1+gzip\", MediaTypeWeightLayerGzip)\n\t\trequire.Equal(t, \"application/vnd.cog.weight.v1\", MediaTypeWeightArtifact)\n\t\trequire.Equal(t, \"application/vnd.cog.weight.layer.v1\", MediaTypeWeightLayer)\n\t})\n}\n"
  },
  {
    "path": "pkg/path/path.go",
    "content": "package path\n\nimport (\n\tgo_path \"path\"\n\t\"strconv\"\n\t\"strings\"\n)\n\nfunc TrimExt(s string) string {\n\treturn strings.TrimSuffix(s, go_path.Ext(s))\n}\n\nfunc IsExtInteger(ext string) bool {\n\text = strings.TrimPrefix(ext, \".\")\n\t_, err := strconv.Atoi(ext)\n\treturn err == nil\n}\n"
  },
  {
    "path": "pkg/path/path_test.go",
    "content": "package path\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestTrimExt(t *testing.T) {\n\tpath := TrimExt(\"/mydir/myoutput.bmp\")\n\trequire.Equal(t, path, \"/mydir/myoutput\")\n}\n\nfunc TestIsExtInteger(t *testing.T) {\n\trequire.True(t, IsExtInteger(\".0\"))\n}\n"
  },
  {
    "path": "pkg/predict/api.go",
    "content": "package predict\n\nimport \"github.com/replicate/cog/pkg/config\"\n\ntype HelpResponse struct {\n\tArguments map[string]*config.RunArgument `json:\"arguments\"`\n}\n"
  },
  {
    "path": "pkg/predict/input.go",
    "content": "package predict\n\nimport (\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"reflect\"\n\t\"strconv\"\n\t\"strings\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\t\"github.com/mitchellh/go-homedir\"\n\t\"github.com/vincent-petithory/dataurl\"\n\n\t\"github.com/replicate/cog/pkg/util/mime\"\n)\n\ntype Input struct {\n\tString *string\n\tFile   *string\n\tArray  *[]any\n\tJson   *json.RawMessage\n\tFloat  *float32\n\tInt    *int32\n}\n\ntype Inputs map[string]Input\n\nfunc NewInputs(keyVals map[string][]string, schema *openapi3.T) (Inputs, error) {\n\treturn NewInputsForMode(keyVals, schema, false)\n}\n\nfunc NewInputsForMode(keyVals map[string][]string, schema *openapi3.T, isTrain bool) (Inputs, error) {\n\tschemaKey := \"Input\"\n\tif isTrain {\n\t\tschemaKey = \"TrainingInput\"\n\t}\n\tvar inputComponent *openapi3.SchemaRef\n\tfor name, component := range schema.Components.Schemas {\n\t\tif name == schemaKey {\n\t\t\tinputComponent = component\n\t\t\tbreak\n\t\t}\n\t}\n\t// Fallback: if TrainingInput not found, try Input (legacy schemas)\n\tif inputComponent == nil && isTrain {\n\t\tfor name, component := range schema.Components.Schemas {\n\t\t\tif name == \"Input\" {\n\t\t\t\tinputComponent = component\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t}\n\n\tinput := Inputs{}\n\tfor key, vals := range keyVals {\n\t\tif len(vals) == 1 {\n\t\t\tval := vals[0]\n\t\t\tif strings.HasPrefix(val, \"@\") {\n\t\t\t\tval = val[1:]\n\t\t\t\tinput[key] = Input{File: &val}\n\t\t\t} else {\n\t\t\t\t// Check if we should explicitly parse the JSON based on a known schema\n\t\t\t\tif inputComponent != nil {\n\t\t\t\t\tproperties, err := inputComponent.JSONLookup(\"properties\")\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn input, err\n\t\t\t\t\t}\n\t\t\t\t\tpropertiesSchemas := properties.(openapi3.Schemas)\n\t\t\t\t\tproperty, err := propertiesSchemas.JSONLookup(key)\n\t\t\t\t\tif err == nil {\n\t\t\t\t\t\tpropertySchema := property.(*openapi3.Schema)\n\t\t\t\t\t\t// Resolve allOf/$ref to find the actual type.\n\t\t\t\t\t\t// cog-schema-gen emits allOf:[{$ref: ...}] for choices/enums,\n\t\t\t\t\t\t// where the referenced schema has the concrete type.\n\t\t\t\t\t\tpropertySchema = resolveSchemaType(propertySchema)\n\t\t\t\t\t\tswitch {\n\t\t\t\t\t\tcase propertySchema.Type.Is(\"object\"):\n\t\t\t\t\t\t\tencodedVal := json.RawMessage(val)\n\t\t\t\t\t\t\tinput[key] = Input{Json: &encodedVal}\n\t\t\t\t\t\t\tcontinue\n\t\t\t\t\t\tcase propertySchema.Type.Is(\"array\"):\n\t\t\t\t\t\t\tvar parsed any\n\t\t\t\t\t\t\terr := json.Unmarshal([]byte(val), &parsed)\n\t\t\t\t\t\t\tif err == nil {\n\t\t\t\t\t\t\t\tt := reflect.TypeOf(parsed)\n\t\t\t\t\t\t\t\tif t.Kind() == reflect.Slice || t.Kind() == reflect.Array {\n\t\t\t\t\t\t\t\t\tencodedVal := json.RawMessage(val)\n\t\t\t\t\t\t\t\t\tinput[key] = Input{Json: &encodedVal}\n\t\t\t\t\t\t\t\t\tcontinue\n\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\tvar arr = []any{val}\n\t\t\t\t\t\t\tinput[key] = Input{Array: &arr}\n\t\t\t\t\t\t\tcontinue\n\t\t\t\t\t\tcase propertySchema.Type.Is(\"number\"):\n\t\t\t\t\t\t\tvalue, err := strconv.ParseInt(val, 10, 32)\n\t\t\t\t\t\t\tif err == nil {\n\t\t\t\t\t\t\t\tvalueInt := int32(value)\n\t\t\t\t\t\t\t\tinput[key] = Input{Int: &valueInt}\n\t\t\t\t\t\t\t\tcontinue\n\t\t\t\t\t\t\t} else {\n\t\t\t\t\t\t\t\tvalue, err := strconv.ParseFloat(val, 32)\n\t\t\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\t\t\treturn input, err\n\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t\tfloat := float32(value)\n\t\t\t\t\t\t\t\tinput[key] = Input{Float: &float}\n\t\t\t\t\t\t\t\tcontinue\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\tcase propertySchema.Type.Is(\"integer\"):\n\t\t\t\t\t\t\tvalue, err := strconv.ParseInt(val, 10, 32)\n\t\t\t\t\t\t\tif err != nil {\n\t\t\t\t\t\t\t\treturn input, err\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\tvalueInt := int32(value)\n\t\t\t\t\t\t\tinput[key] = Input{Int: &valueInt}\n\t\t\t\t\t\t\tcontinue\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\tinput[key] = Input{String: &val}\n\t\t\t}\n\t\t} else if len(vals) > 1 {\n\t\t\tvar anyVals = make([]any, len(vals))\n\t\t\tfor i, v := range vals {\n\t\t\t\tanyVals[i] = v\n\t\t\t}\n\t\t\tinput[key] = Input{Array: &anyVals}\n\t\t}\n\t}\n\treturn input, nil\n}\n\nfunc NewInputsWithBaseDir(keyVals map[string]string, baseDir string) Inputs {\n\tinput := Inputs{}\n\tfor key, val := range keyVals {\n\t\tif strings.HasPrefix(val, \"@\") {\n\t\t\tval = filepath.Join(baseDir, val[1:])\n\t\t\tinput[key] = Input{File: &val}\n\t\t} else {\n\t\t\tinput[key] = Input{String: &val}\n\t\t}\n\t}\n\treturn input\n}\n\nfunc (inputs *Inputs) toMap() (map[string]any, error) {\n\tkeyVals := map[string]any{}\n\tfor key, input := range *inputs {\n\t\tswitch {\n\t\tcase input.String != nil:\n\t\t\t// Directly assign the string value\n\t\t\tkeyVals[key] = *input.String\n\t\tcase input.File != nil:\n\t\t\t// Single file handling: read content and convert to a data URL\n\t\t\tdataURL, err := fileToDataURL(*input.File)\n\t\t\tif err != nil {\n\t\t\t\treturn keyVals, err\n\t\t\t}\n\t\t\tkeyVals[key] = dataURL\n\t\tcase input.Array != nil:\n\t\t\t// Handle array, potentially containing file paths\n\t\t\tdataURLs := make([]string, len(*input.Array))\n\t\t\tfor i, elem := range *input.Array {\n\t\t\t\tif str, ok := elem.(string); ok && strings.HasPrefix(str, \"@\") {\n\t\t\t\t\tfilePath := str[1:] // Remove '@' prefix\n\t\t\t\t\tdataURL, err := fileToDataURL(filePath)\n\t\t\t\t\tif err != nil {\n\t\t\t\t\t\treturn keyVals, err\n\t\t\t\t\t}\n\t\t\t\t\tdataURLs[i] = dataURL\n\t\t\t\t} else if ok {\n\t\t\t\t\t// Directly use the string if it's not a file path\n\t\t\t\t\tdataURLs[i] = str\n\t\t\t\t}\n\t\t\t}\n\t\t\tkeyVals[key] = dataURLs\n\t\tcase input.Json != nil:\n\t\t\tkeyVals[key] = *input.Json\n\t\tcase input.Float != nil:\n\t\t\tkeyVals[key] = *input.Float\n\t\tcase input.Int != nil:\n\t\t\tkeyVals[key] = *input.Int\n\t\t}\n\t}\n\treturn keyVals, nil\n}\n\n// Helper function to read file content and convert to a data URL\nfunc fileToDataURL(filePath string) (string, error) {\n\t// Expand home directory if necessary\n\texpandedVal, err := homedir.Expand(filePath)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"error expanding homedir for '%s': %w\", filePath, err)\n\t}\n\n\tcontent, err := os.ReadFile(expandedVal)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tmimeType := mime.TypeByExtension(filepath.Ext(expandedVal))\n\tdataURL := dataurl.New(content, mimeType).String()\n\treturn dataURL, nil\n}\n\n// resolveSchemaType walks through allOf/anyOf/$ref wrappers to find a schema\n// that has a concrete Type set. This is needed because the static schema gen\n// emits allOf:[{$ref: \"#/components/schemas/Foo\"}] for enum/choices fields,\n// where the referenced schema carries the type (e.g. \"integer\") but the wrapper does not.\nfunc resolveSchemaType(s *openapi3.Schema) *openapi3.Schema {\n\tif s.Type != nil && s.Type.Slice() != nil {\n\t\treturn s\n\t}\n\t// Check allOf entries\n\tfor _, ref := range s.AllOf {\n\t\tif ref.Value != nil && ref.Value.Type != nil && ref.Value.Type.Slice() != nil {\n\t\t\treturn ref.Value\n\t\t}\n\t}\n\t// Check anyOf entries\n\tfor _, ref := range s.AnyOf {\n\t\tif ref.Value != nil && ref.Value.Type != nil && ref.Value.Type.Slice() != nil {\n\t\t\treturn ref.Value\n\t\t}\n\t}\n\treturn s\n}\n"
  },
  {
    "path": "pkg/predict/predictor.go",
    "content": "package predict\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"io\"\n\t\"net/http\"\n\t\"strings\"\n\t\"time\"\n\n\t\"github.com/getkin/kin-openapi/openapi3\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\ntype status string\n\ntype HealthcheckResponse struct {\n\tStatus string `json:\"status\"`\n}\n\ntype RequestContext struct {\n\tReplicateAPIToken string `json:\"replicate_api_token,omitempty\"`\n}\n\ntype Request struct {\n\t// TODO: could this be Inputs?\n\tInput   map[string]any `json:\"input\"`\n\tContext RequestContext `json:\"context\"`\n}\n\ntype Response struct {\n\tStatus status `json:\"status\"`\n\tOutput *any   `json:\"output\"`\n\tError  string `json:\"error\"`\n}\n\ntype ValidationErrorResponse struct {\n\tDetail []struct {\n\t\tLocation []string `json:\"loc\"`\n\t\tMessage  string   `json:\"msg\"`\n\t\tType     string   `json:\"type\"`\n\t} `json:\"detail\"`\n}\n\ntype Predictor struct {\n\trunOptions   command.RunOptions\n\tisTrain      bool\n\tdockerClient command.Command\n\n\t// Running state\n\tcontainerID string\n\tport        int\n}\n\nfunc NewPredictor(ctx context.Context, runOptions command.RunOptions, isTrain bool, dockerCommand command.Command) (*Predictor, error) {\n\tif global.Debug {\n\t\trunOptions.Env = append(runOptions.Env, \"COG_LOG_LEVEL=debug\")\n\t} else {\n\t\trunOptions.Env = append(runOptions.Env, \"COG_LOG_LEVEL=warning\")\n\t}\n\n\treturn &Predictor{\n\t\trunOptions:   runOptions,\n\t\tisTrain:      isTrain,\n\t\tdockerClient: dockerCommand,\n\t}, nil\n}\n\nfunc (p *Predictor) Start(ctx context.Context, logsWriter io.Writer, timeout time.Duration) error {\n\tvar err error\n\tcontainerPort := 5000\n\n\tp.runOptions.Ports = append(p.runOptions.Ports, command.Port{HostPort: 0, ContainerPort: containerPort})\n\n\tp.containerID, err = docker.RunDaemon(ctx, p.dockerClient, p.runOptions, logsWriter)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Failed to start container: %w\", err)\n\t}\n\n\tp.port, err = docker.GetHostPortForContainer(ctx, p.dockerClient, p.containerID, containerPort)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Failed to determine container port: %w\", err)\n\t}\n\n\tgo func() {\n\t\tif err := p.dockerClient.ContainerLogs(ctx, p.containerID, logsWriter); err != nil {\n\t\t\t// if user hits ctrl-c we expect an error signal\n\t\t\tif !strings.Contains(err.Error(), \"signal: interrupt\") {\n\t\t\t\tconsole.Warnf(\"Error getting container logs: %s\", err)\n\t\t\t}\n\t\t}\n\t}()\n\n\treturn p.waitForContainerReady(ctx, timeout)\n}\n\nfunc (p *Predictor) waitForContainerReady(ctx context.Context, timeout time.Duration) error {\n\turl := fmt.Sprintf(\"http://localhost:%d/health-check\", p.port)\n\n\tstart := time.Now()\n\tfor {\n\t\tif time.Since(start) > timeout {\n\t\t\treturn fmt.Errorf(\"Timed out\")\n\t\t}\n\n\t\ttime.Sleep(100 * time.Millisecond)\n\n\t\tcont, err := p.dockerClient.ContainerInspect(ctx, p.containerID)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"Failed to get container status: %w\", err)\n\t\t}\n\t\tif cont.State != nil && (cont.State.Status == \"exited\" || cont.State.Status == \"dead\") {\n\t\t\treturn fmt.Errorf(\"Container exited unexpectedly\")\n\t\t}\n\n\t\thealthcheck, err := func() (*HealthcheckResponse, error) {\n\t\t\tctx, cancel := context.WithTimeout(ctx, 250*time.Millisecond)\n\t\t\tdefer cancel()\n\t\t\treq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"Failed to create HTTP request to %s: %w\", url, err)\n\t\t\t}\n\n\t\t\tresp, err := http.DefaultClient.Do(req) //nolint:gosec // G704: URL from localhost health check\n\t\t\tif err != nil {\n\t\t\t\treturn nil, nil\n\t\t\t}\n\t\t\tdefer resp.Body.Close()\n\t\t\tif resp.StatusCode != http.StatusOK {\n\t\t\t\treturn nil, nil\n\t\t\t}\n\t\t\thealthcheck := &HealthcheckResponse{}\n\t\t\tif err := json.NewDecoder(resp.Body).Decode(healthcheck); err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"Container healthcheck returned invalid response: %w\", err)\n\t\t\t}\n\t\t\treturn healthcheck, nil\n\t\t}()\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tif healthcheck == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\t// These status values are defined in python/cog/server/http.py\n\t\tswitch healthcheck.Status {\n\t\tcase \"STARTING\":\n\t\t\tcontinue\n\t\tcase \"SETUP_FAILED\":\n\t\t\treturn fmt.Errorf(\"Model setup failed\")\n\t\tcase \"READY\":\n\t\t\treturn nil\n\t\tdefault:\n\t\t\treturn fmt.Errorf(\"Container healthcheck returned unexpected status: %s\", healthcheck.Status)\n\t\t}\n\t}\n}\n\nfunc (p *Predictor) Stop(ctx context.Context) error {\n\treturn p.dockerClient.ContainerStop(ctx, p.containerID)\n}\n\nfunc (p *Predictor) Predict(inputs Inputs, context RequestContext) (*Response, error) {\n\tinputMap, err := inputs.toMap()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\trequest := Request{\n\t\tInput:   inputMap,\n\t\tContext: context,\n\t}\n\trequestBody, err := json.Marshal(request)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\turl := p.url()\n\treq, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(requestBody))\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Failed to create HTTP request to %s: %w\", url, err)\n\t}\n\treq.Header.Set(\"Content-Type\", \"application/json\")\n\treq.Close = true\n\n\thttpClient := &http.Client{}\n\tresp, err := httpClient.Do(req) //nolint:gosec // G704: URL from localhost prediction endpoint\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Failed to POST HTTP request to %s: %w\", url, err)\n\t}\n\tdefer resp.Body.Close()\n\n\tif resp.StatusCode == http.StatusUnprocessableEntity {\n\t\terrorResponse := &ValidationErrorResponse{}\n\t\tif err := json.NewDecoder(resp.Body).Decode(errorResponse); err != nil {\n\t\t\treturn nil, fmt.Errorf(\"/%s call returned status 422, and the response body failed to decode: %w\", p.endpoint(), err)\n\t\t}\n\n\t\treturn nil, p.buildInputValidationErrorMessage(errorResponse)\n\t}\n\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn nil, fmt.Errorf(\"/%s call returned status %d\", p.endpoint(), resp.StatusCode)\n\t}\n\n\tprediction := &Response{}\n\tif err = json.NewDecoder(resp.Body).Decode(prediction); err != nil {\n\t\treturn nil, fmt.Errorf(\"Failed to decode prediction response: %w\", err)\n\t}\n\treturn prediction, nil\n}\n\nfunc (p *Predictor) GetSchema() (*openapi3.T, error) {\n\tresp, err := http.Get(fmt.Sprintf(\"http://localhost:%d/openapi.json\", p.port))\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn nil, fmt.Errorf(\"Failed to get OpenAPI schema: %d\", resp.StatusCode)\n\t}\n\n\tbody, err := io.ReadAll(resp.Body)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn openapi3.NewLoader().LoadFromData(body)\n}\n\nfunc (p *Predictor) endpoint() string {\n\tif p.isTrain {\n\t\treturn \"trainings\"\n\t}\n\treturn \"predictions\"\n}\n\nfunc (p *Predictor) url() string {\n\treturn fmt.Sprintf(\"http://localhost:%d/%s\", p.port, p.endpoint())\n}\n\nfunc (p *Predictor) buildInputValidationErrorMessage(errorResponse *ValidationErrorResponse) error {\n\terrorMessages := []string{}\n\n\tfor _, validationError := range errorResponse.Detail {\n\t\tif len(validationError.Location) != 3 || validationError.Location[0] != \"body\" || validationError.Location[1] != \"input\" {\n\t\t\tresponseBody, _ := json.MarshalIndent(errorResponse, \"\", \"\\t\")\n\t\t\treturn fmt.Errorf(\"/%s call returned status 422, and there was an unexpected message in response:\\n\\n%s\", p.endpoint(), responseBody)\n\t\t}\n\n\t\terrorMessages = append(errorMessages, fmt.Sprintf(\"- %s: %s\", validationError.Location[2], validationError.Message))\n\t}\n\n\tcommand := \"predict\"\n\tif p.isTrain {\n\t\tcommand = \"train\"\n\t}\n\n\treturn fmt.Errorf(\n\t\t`The inputs you passed could not be validated:\n\n%[2]s\n\nYou can provide an input with -i. For example:\n\n    cog %[1]s -i blur=3.5\n\nIf your input is a local file, you need to prefix the path with @ to tell Cog to read the file contents. For example:\n\n    cog %[1]s -i path=@image.jpg`,\n\t\tcommand,\n\t\tstrings.Join(errorMessages, \"\\n\"),\n\t)\n}\n"
  },
  {
    "path": "pkg/provider/generic/generic.go",
    "content": "package generic\n\nimport (\n\t\"bufio\"\n\t\"context\"\n\t\"fmt\"\n\t\"os\"\n\t\"strings\"\n\n\t\"golang.org/x/term\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/provider\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// GenericProvider works with any OCI-compliant registry\ntype GenericProvider struct{}\n\n// New creates a new GenericProvider\nfunc New() *GenericProvider {\n\treturn &GenericProvider{}\n}\n\nfunc (p *GenericProvider) Name() string {\n\treturn \"generic\"\n}\n\nfunc (p *GenericProvider) MatchesRegistry(host string) bool {\n\treturn true // Fallback - matches everything\n}\n\nfunc (p *GenericProvider) Login(ctx context.Context, opts provider.LoginOptions) error {\n\tconsole.InfoUnformattedf(\"Logging in to %s\", opts.Host)\n\tconsole.InfoUnformatted(\"\")\n\n\t// TODO: support non-interactive login with token stdin for generic registries\n\t// Prompt for username\n\tfmt.Print(\"Username: \")\n\treader := bufio.NewReader(os.Stdin)\n\tusername, err := reader.ReadString('\\n')\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to read username: %w\", err)\n\t}\n\tusername = strings.TrimSpace(username)\n\n\tif username == \"\" {\n\t\treturn fmt.Errorf(\"username cannot be empty\")\n\t}\n\n\t// Prompt for password (hidden input)\n\tfmt.Print(\"Password: \")\n\tpasswordBytes, err := term.ReadPassword(int(os.Stdin.Fd())) //nolint:gosec // G115: Fd() fits in int on all supported platforms\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to read password: %w\", err)\n\t}\n\tfmt.Println() // newline after hidden input\n\tpassword := string(passwordBytes)\n\n\tif password == \"\" {\n\t\treturn fmt.Errorf(\"password cannot be empty\")\n\t}\n\n\t// Save credentials using Docker's credential system\n\tif err := docker.SaveLoginToken(ctx, opts.Host, username, password); err != nil {\n\t\treturn fmt.Errorf(\"failed to save credentials: %w\", err)\n\t}\n\n\tconsole.Successf(\"Login succeeded for %s\", console.Bold(opts.Host))\n\treturn nil\n}\n\nfunc (p *GenericProvider) PostPush(ctx context.Context, opts provider.PushOptions, pushErr error) error {\n\t// No special post-push handling for generic registries\n\t// Just show a simple success message if push succeeded\n\tif pushErr == nil {\n\t\tconsole.Successf(\"Image %s pushed\", console.Bold(opts.Image))\n\t}\n\treturn nil\n}\n\n// Verify interface compliance at compile time\nvar _ provider.Provider = (*GenericProvider)(nil)\n"
  },
  {
    "path": "pkg/provider/generic/generic_test.go",
    "content": "package generic\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/provider\"\n)\n\nfunc TestGenericProvider_Name(t *testing.T) {\n\tp := New()\n\trequire.Equal(t, \"generic\", p.Name())\n}\n\nfunc TestGenericProvider_MatchesRegistry(t *testing.T) {\n\tp := New()\n\t// Generic provider matches everything (it's the fallback)\n\trequire.True(t, p.MatchesRegistry(\"ghcr.io\"))\n\trequire.True(t, p.MatchesRegistry(\"docker.io\"))\n\trequire.True(t, p.MatchesRegistry(\"ecr.aws\"))\n\trequire.True(t, p.MatchesRegistry(\"anything.example.com\"))\n}\n\nfunc TestGenericProvider_Login(t *testing.T) {\n\t// Login() prompts for username/password interactively and saves credentials\n\t// via Docker's credential system. This cannot be easily tested without mocking\n\t// stdin and the docker credential helpers.\n\t//\n\t// The Login method:\n\t// 1. Prompts for username (from stdin)\n\t// 2. Prompts for password (hidden input via terminal)\n\t// 3. Calls docker.SaveLoginToken() to store credentials\n\t//\n\t// For integration testing, use manual testing with 'cog login --registry '\n\tt.Skip(\"Login requires interactive input - test manually\")\n}\n\nfunc TestGenericProvider_PostPush(t *testing.T) {\n\tp := New()\n\n\tt.Run(\"success\", func(t *testing.T) {\n\t\topts := provider.PushOptions{\n\t\t\tImage: \"ghcr.io/org/model\",\n\t\t}\n\t\terr := p.PostPush(context.Background(), opts, nil)\n\t\trequire.NoError(t, err)\n\t})\n\n\tt.Run(\"with error\", func(t *testing.T) {\n\t\topts := provider.PushOptions{\n\t\t\tImage: \"ghcr.io/org/model\",\n\t\t}\n\t\tpushErr := errors.New(\"push failed\")\n\t\terr := p.PostPush(context.Background(), opts, pushErr)\n\t\trequire.NoError(t, err) // PostPush itself doesn't error\n\t})\n}\n"
  },
  {
    "path": "pkg/provider/provider.go",
    "content": "package provider\n\nimport (\n\t\"context\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\n// PushOptions contains all options for a push operation\ntype PushOptions struct {\n\tImage      string\n\tConfig     *config.Config\n\tProjectDir string\n}\n\ntype LoginOptions struct {\n\tTokenStdin bool\n\tHost       string\n}\n\n// Provider encapsulates registry-specific behavior\ntype Provider interface {\n\t// Name returns the provider identifier (e.g., \"replicate\", \"generic\")\n\tName() string\n\n\t// MatchesRegistry returns true if this provider handles the given registry host\n\tMatchesRegistry(host string) bool\n\n\t// Login performs provider-specific authentication\n\tLogin(ctx context.Context, opts LoginOptions) error\n\n\t// PostPush is called after push attempt (success or failure)\n\t// - Shows success message (e.g., Replicate model URL)\n\t// - May transform errors into provider-specific messages\n\t// - pushErr is nil on success, contains the push error on failure\n\tPostPush(ctx context.Context, opts PushOptions, pushErr error) error\n}\n"
  },
  {
    "path": "pkg/provider/registry.go",
    "content": "package provider\n\nimport (\n\t\"strings\"\n\t\"sync\"\n)\n\n// defaultRegistry is the global singleton registry\nvar defaultRegistry *Registry\n\n// DefaultRegistry returns the global provider registry, initializing it on first call\n// The registry is pre-populated with Replicate and Generic providers\nfunc DefaultRegistry() *Registry {\n\tif defaultRegistry == nil {\n\t\tdefaultRegistry = NewRegistry()\n\t\t// Note: providers are registered by init() functions in their respective packages\n\t\t// via RegisterProvider(), or can be set up explicitly\n\t}\n\treturn defaultRegistry\n}\n\n// RegisterProvider adds a provider to the default registry\n// This should be called from init() functions in provider packages\nfunc RegisterProvider(p Provider) {\n\tDefaultRegistry().Register(p)\n}\n\n// Registry manages provider lookup and registration\ntype Registry struct {\n\tproviders []Provider\n\tmu        sync.RWMutex\n}\n\n// NewRegistry creates a new Registry with no providers registered\nfunc NewRegistry() *Registry {\n\treturn &Registry{\n\t\tproviders: make([]Provider, 0),\n\t}\n}\n\n// Register adds a provider to the registry\n// Providers are checked in registration order, so register more specific\n// providers before generic fallback providers\nfunc (r *Registry) Register(p Provider) {\n\tr.mu.Lock()\n\tdefer r.mu.Unlock()\n\tr.providers = append(r.providers, p)\n}\n\n// ForImage returns the provider for a given image name\n// It extracts the registry host from the image and delegates to ForHost\nfunc (r *Registry) ForImage(image string) Provider {\n\thost := ExtractHost(image)\n\treturn r.ForHost(host)\n}\n\n// ForHost returns the provider for a given registry host\n// Returns the first provider that matches, or nil if none match\nfunc (r *Registry) ForHost(host string) Provider {\n\tr.mu.RLock()\n\tdefer r.mu.RUnlock()\n\n\tfor _, p := range r.providers {\n\t\tif p.MatchesRegistry(host) {\n\t\t\treturn p\n\t\t}\n\t}\n\treturn nil\n}\n\n// ExtractHost extracts the registry host from an image name\n// Examples:\n//   - \"r8.im/user/model\" -> \"r8.im\"\n//   - \"ghcr.io/owner/repo:tag\" -> \"ghcr.io\"\n//   - \"gcr.io/project/image\" -> \"gcr.io\"\n//   - \"docker.io/library/nginx\" -> \"docker.io\"\n//   - \"nginx\" -> \"docker.io\" (Docker Hub default)\n//   - \"myregistry.com:5000/image\" -> \"myregistry.com:5000\"\n//   - \"localhost:5000/image\" -> \"localhost:5000\"\nfunc ExtractHost(image string) string {\n\t// Handle empty image\n\tif image == \"\" {\n\t\treturn \"docker.io\"\n\t}\n\n\t// Remove digest first (@sha256:...)\n\tif idx := strings.Index(image, \"@\"); idx != -1 {\n\t\timage = image[:idx]\n\t}\n\n\t// Get the first component (everything before the first slash)\n\t// If there's no slash, it's a Docker Hub image (e.g., \"nginx\" or \"nginx:latest\")\n\tfirstComponent, _, found := strings.Cut(image, \"/\")\n\tif !found {\n\t\treturn \"docker.io\"\n\t}\n\n\t// Check if it looks like a registry host:\n\t// - Contains a dot (e.g., gcr.io, ghcr.io, r8.im, myregistry.com)\n\t// - Contains a colon (e.g., localhost:5000, myregistry.com:5000)\n\t// - Is \"localhost\"\n\tif strings.Contains(firstComponent, \".\") ||\n\t\tstrings.Contains(firstComponent, \":\") ||\n\t\tfirstComponent == \"localhost\" {\n\t\treturn firstComponent\n\t}\n\n\t// Otherwise it's a Docker Hub user/image (e.g., \"user/image\")\n\treturn \"docker.io\"\n}\n"
  },
  {
    "path": "pkg/provider/registry_test.go",
    "content": "package provider\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\n// mockProvider implements Provider for testing\ntype mockProvider struct {\n\tname    string\n\tmatches func(host string) bool\n}\n\nfunc (m *mockProvider) Name() string {\n\treturn m.name\n}\n\nfunc (m *mockProvider) MatchesRegistry(host string) bool {\n\treturn m.matches(host)\n}\n\nfunc (m *mockProvider) Login(ctx context.Context, opts LoginOptions) error {\n\treturn nil\n}\n\nfunc (m *mockProvider) PrePush(ctx context.Context, opts PushOptions) error {\n\treturn nil\n}\n\nfunc (m *mockProvider) PostPush(ctx context.Context, opts PushOptions, pushErr error) error {\n\treturn nil\n}\n\nfunc TestRegistry_ForHost(t *testing.T) {\n\tr := NewRegistry()\n\n\treplicateProvider := &mockProvider{\n\t\tname:    \"replicate\",\n\t\tmatches: func(host string) bool { return host == \"r8.im\" },\n\t}\n\tgenericProvider := &mockProvider{\n\t\tname:    \"generic\",\n\t\tmatches: func(host string) bool { return true },\n\t}\n\n\t// Register replicate first (more specific), then generic (fallback)\n\tr.Register(replicateProvider)\n\tr.Register(genericProvider)\n\n\tt.Run(\"matches replicate\", func(t *testing.T) {\n\t\tp := r.ForHost(\"r8.im\")\n\t\trequire.NotNil(t, p)\n\t\trequire.Equal(t, \"replicate\", p.Name())\n\t})\n\n\tt.Run(\"falls back to generic\", func(t *testing.T) {\n\t\tp := r.ForHost(\"ghcr.io\")\n\t\trequire.NotNil(t, p)\n\t\trequire.Equal(t, \"generic\", p.Name())\n\t})\n\n\tt.Run(\"empty host falls back to generic\", func(t *testing.T) {\n\t\tp := r.ForHost(\"\")\n\t\trequire.NotNil(t, p)\n\t\trequire.Equal(t, \"generic\", p.Name())\n\t})\n}\n\nfunc TestRegistry_ForImage(t *testing.T) {\n\tr := NewRegistry()\n\n\treplicateProvider := &mockProvider{\n\t\tname:    \"replicate\",\n\t\tmatches: func(host string) bool { return host == \"r8.im\" },\n\t}\n\tgenericProvider := &mockProvider{\n\t\tname:    \"generic\",\n\t\tmatches: func(host string) bool { return true },\n\t}\n\n\tr.Register(replicateProvider)\n\tr.Register(genericProvider)\n\n\ttests := []struct {\n\t\timage        string\n\t\texpectedName string\n\t}{\n\t\t{\"r8.im/user/model\", \"replicate\"},\n\t\t{\"r8.im/user/model:v1\", \"replicate\"},\n\t\t{\"ghcr.io/owner/repo\", \"generic\"},\n\t\t{\"gcr.io/project/image:tag\", \"generic\"},\n\t\t{\"docker.io/library/nginx\", \"generic\"},\n\t\t{\"nginx\", \"generic\"},\n\t\t{\"myregistry.com/image\", \"generic\"},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.image, func(t *testing.T) {\n\t\t\tp := r.ForImage(tt.image)\n\t\t\trequire.NotNil(t, p)\n\t\t\trequire.Equal(t, tt.expectedName, p.Name())\n\t\t})\n\t}\n}\n\nfunc TestRegistry_NoProviders(t *testing.T) {\n\tr := NewRegistry()\n\n\tp := r.ForHost(\"any.registry.io\")\n\trequire.Nil(t, p)\n}\n\nfunc TestExtractHost(t *testing.T) {\n\ttests := []struct {\n\t\timage    string\n\t\texpected string\n\t}{\n\t\t// Replicate\n\t\t{\"r8.im/user/model\", \"r8.im\"},\n\t\t{\"r8.im/user/model:v1\", \"r8.im\"},\n\t\t{\"r8.im/user/model:latest\", \"r8.im\"},\n\n\t\t// GitHub Container Registry\n\t\t{\"ghcr.io/owner/repo\", \"ghcr.io\"},\n\t\t{\"ghcr.io/owner/repo:tag\", \"ghcr.io\"},\n\n\t\t// Google Container Registry\n\t\t{\"gcr.io/project/image\", \"gcr.io\"},\n\t\t{\"gcr.io/project/image:tag\", \"gcr.io\"},\n\t\t{\"us.gcr.io/project/image\", \"us.gcr.io\"},\n\n\t\t// Docker Hub explicit\n\t\t{\"docker.io/library/nginx\", \"docker.io\"},\n\t\t{\"docker.io/user/image\", \"docker.io\"},\n\n\t\t// Docker Hub implicit (no registry specified)\n\t\t{\"nginx\", \"docker.io\"},\n\t\t{\"nginx:latest\", \"docker.io\"},\n\t\t{\"user/image\", \"docker.io\"},\n\t\t{\"user/image:tag\", \"docker.io\"},\n\n\t\t// Custom registries\n\t\t{\"myregistry.com/image\", \"myregistry.com\"},\n\t\t{\"myregistry.example.com/path/to/image\", \"myregistry.example.com\"},\n\n\t\t// Registries with ports\n\t\t{\"localhost:5000/image\", \"localhost:5000\"},\n\t\t{\"myregistry.com:5000/image\", \"myregistry.com:5000\"},\n\t\t{\"myregistry.com:5000/image:tag\", \"myregistry.com:5000\"},\n\n\t\t// With digest\n\t\t{\"ghcr.io/owner/repo@sha256:abc123\", \"ghcr.io\"},\n\n\t\t// Edge cases\n\t\t{\"\", \"docker.io\"},\n\t\t{\"localhost/image\", \"localhost\"},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.image, func(t *testing.T) {\n\t\t\tresult := ExtractHost(tt.image)\n\t\t\trequire.Equal(t, tt.expected, result, \"ExtractHost(%q)\", tt.image)\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "pkg/provider/replicate/replicate.go",
    "content": "package replicate\n\nimport (\n\t\"bufio\"\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"os\"\n\t\"os/exec\"\n\t\"runtime\"\n\t\"strings\"\n\n\t\"golang.org/x/term\"\n\n\t\"github.com/replicate/cog/pkg/docker\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/provider\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// ReplicateProvider handles Replicate's r8.im registry\ntype ReplicateProvider struct{}\n\n// New creates a new ReplicateProvider\nfunc New() *ReplicateProvider {\n\treturn &ReplicateProvider{}\n}\n\nfunc (p *ReplicateProvider) Name() string {\n\treturn \"replicate\"\n}\n\nfunc (p *ReplicateProvider) MatchesRegistry(host string) bool {\n\t// Only match the default Replicate registry host (r8.im)\n\t// Note: We don't use global.ReplicateRegistryHost here because that variable\n\t// gets updated by the --registry flag, which would cause us to incorrectly\n\t// match any registry the user specifies.\n\treturn host == global.DefaultReplicateRegistryHost\n}\n\n// Login performs login to the registry with options\nfunc (p *ReplicateProvider) Login(ctx context.Context, opts provider.LoginOptions) error {\n\tvar (\n\t\ttoken string\n\t\terr   error\n\t)\n\n\tif opts.TokenStdin {\n\t\ttoken, err = readTokenFromStdin()\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t} else {\n\t\ttoken, err = readTokenInteractively(opts.Host)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t}\n\n\ttoken = strings.TrimSpace(token)\n\n\tif err := checkTokenFormat(token); err != nil {\n\t\treturn err\n\t}\n\n\tusername, err := verifyToken(opts.Host, token)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tif err := docker.SaveLoginToken(ctx, opts.Host, username, token); err != nil {\n\t\treturn err\n\t}\n\n\tconsole.Successf(\"You've successfully authenticated as %s! You can now use the %s registry.\", console.Bold(username), console.Bold(opts.Host))\n\treturn nil\n}\n\nfunc (p *ReplicateProvider) PostPush(ctx context.Context, opts provider.PushOptions, pushErr error) error {\n\tif pushErr != nil {\n\t\t// Return Replicate-specific error message for repository not found errors\n\t\tif command.IsNotFoundError(pushErr) {\n\t\t\treturn fmt.Errorf(`Unable to find existing Replicate model for %s. Go to replicate.com and create a new model before pushing.\n\nIf the model already exists, you may be getting this error because you're not logged in as owner of the model. This can happen if you did 'sudo cog login' instead of 'cog login' or 'sudo cog push' instead of 'cog push', which causes Docker to use the wrong Docker credentials.`, opts.Image)\n\t\t}\n\t\treturn pushErr\n\t}\n\n\t// Success - show Replicate model URL\n\tconsole.Successf(\"Image %s pushed\", console.Bold(opts.Image))\n\treplicatePage := fmt.Sprintf(\"https://%s\", strings.Replace(opts.Image, global.ReplicateRegistryHost, global.ReplicateWebsiteHost, 1))\n\tconsole.Infof(\"\\nRun your model on Replicate:\\n    %s\", console.Bold(replicatePage))\n\n\treturn nil\n}\n\n// readTokenFromStdin reads the authentication token from stdin\nfunc readTokenFromStdin() (string, error) {\n\ttokenBytes, err := os.ReadFile(\"/dev/stdin\")\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to read token from stdin: %w\", err)\n\t}\n\treturn string(tokenBytes), nil\n}\n\n// readTokenInteractively guides user through browser-based token flow\nfunc readTokenInteractively(registryHost string) (string, error) {\n\ttokenURL, err := getDisplayTokenURL(registryHost)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tconsole.InfoUnformattedf(\"This command will authenticate Docker with Replicate's '%s' Docker registry. You will need a Replicate account.\", registryHost)\n\tconsole.InfoUnformatted(\"\")\n\tconsole.InfoUnformatted(\"Hit enter to get started. A browser will open with an authentication token that you need to paste here.\")\n\n\tinputReader := os.Stdin\n\tinputFd := int(os.Stdin.Fd()) //nolint:gosec // G115: Fd() fits in int on all supported platforms\n\n\treader := bufio.NewReader(inputReader)\n\tif _, err := reader.ReadString('\\n'); err != nil {\n\t\treturn \"\", err\n\t}\n\n\tconsole.InfoUnformatted(\"If it didn't open automatically, open this URL in a web browser:\")\n\tconsole.InfoUnformatted(tokenURL)\n\tmaybeOpenBrowser(tokenURL)\n\n\tconsole.InfoUnformatted(\"\")\n\tconsole.InfoUnformatted(\"Once you've signed in, copy the token from that web page, paste it here, then hit enter:\")\n\tconsole.InfoUnformatted(\"\")\n\n\tfmt.Print(\"Token: \")\n\t// Read the token securely, masking the input\n\ttokenBytes, err := term.ReadPassword(inputFd)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to read token: %w\", err)\n\t}\n\n\t// Print a newline after the hidden input\n\tfmt.Println()\n\tconsole.InfoUnformatted(\"\")\n\n\treturn string(tokenBytes), nil\n}\n\n// getDisplayTokenURL fetches the token URL from Replicate's API\nfunc getDisplayTokenURL(registryHost string) (string, error) {\n\tresp, err := http.Get(addressWithScheme(registryHost) + \"/cog/v1/display-token-url\")\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to log in to %s: %w\", registryHost, err)\n\t}\n\tdefer resp.Body.Close()\n\n\tif resp.StatusCode == http.StatusNotFound {\n\t\treturn \"\", fmt.Errorf(\"%s is not the Replicate registry\\nPlease log in using 'docker login'\", registryHost)\n\t}\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn \"\", fmt.Errorf(\"%s returned HTTP status %d\", registryHost, resp.StatusCode)\n\t}\n\n\tbody := &struct {\n\t\tURL string `json:\"url\"`\n\t}{}\n\tif err := json.NewDecoder(resp.Body).Decode(body); err != nil {\n\t\treturn \"\", err\n\t}\n\treturn body.URL, nil\n}\n\n// addressWithScheme ensures the address has an https:// scheme\nfunc addressWithScheme(address string) string {\n\tif strings.Contains(address, \"://\") {\n\t\treturn address\n\t}\n\treturn \"https://\" + address\n}\n\n// maybeOpenBrowser attempts to open the URL in the default browser\nfunc maybeOpenBrowser(urlToOpen string) {\n\tswitch runtime.GOOS {\n\tcase \"linux\":\n\t\t_ = exec.Command(\"xdg-open\", urlToOpen).Start()\n\tcase \"windows\":\n\t\t_ = exec.Command(\"rundll32\", \"url.dll,FileProtocolHandler\", urlToOpen).Start()\n\tcase \"darwin\":\n\t\t_ = exec.Command(\"open\", urlToOpen).Start()\n\t}\n}\n\n// checkTokenFormat validates the token isn't an API token\nfunc checkTokenFormat(token string) error {\n\tif strings.HasPrefix(token, \"r8_\") {\n\t\treturn fmt.Errorf(\"that looks like a Replicate API token, not a CLI auth token. Please fetch a token from https://replicate.com/auth/token to log in\")\n\t}\n\treturn nil\n}\n\n// verifyToken validates the token with Replicate and returns the username\nfunc verifyToken(registryHost string, token string) (username string, err error) {\n\tif token == \"\" {\n\t\treturn \"\", fmt.Errorf(\"token is empty\")\n\t}\n\n\tresp, err := http.PostForm(addressWithScheme(registryHost)+\"/cog/v1/verify-token\", url.Values{\n\t\t\"token\": []string{token},\n\t})\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to verify token: %w\", err)\n\t}\n\tdefer resp.Body.Close()\n\n\tif resp.StatusCode == http.StatusNotFound {\n\t\treturn \"\", fmt.Errorf(\"user does not exist\")\n\t}\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn \"\", fmt.Errorf(\"failed to verify token, got status %d\", resp.StatusCode)\n\t}\n\n\tbody := &struct {\n\t\tUsername string `json:\"username\"`\n\t}{}\n\tif err := json.NewDecoder(resp.Body).Decode(body); err != nil {\n\t\treturn \"\", err\n\t}\n\treturn body.Username, nil\n}\n\n// Verify interface compliance at compile time\nvar _ provider.Provider = (*ReplicateProvider)(nil)\n"
  },
  {
    "path": "pkg/provider/replicate/replicate_test.go",
    "content": "package replicate\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"net/http\"\n\t\"net/http/httptest\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/provider\"\n)\n\nfunc TestReplicateProvider_Name(t *testing.T) {\n\tp := New()\n\trequire.Equal(t, \"replicate\", p.Name())\n}\n\nfunc TestReplicateProvider_MatchesRegistry(t *testing.T) {\n\tp := New()\n\n\t// Should match default r8.im\n\trequire.True(t, p.MatchesRegistry(\"r8.im\"))\n\n\t// Should match the current global registry host (in case it was overridden)\n\trequire.True(t, p.MatchesRegistry(global.ReplicateRegistryHost))\n\n\t// Should not match other registries\n\trequire.False(t, p.MatchesRegistry(\"ghcr.io\"))\n\trequire.False(t, p.MatchesRegistry(\"docker.io\"))\n\trequire.False(t, p.MatchesRegistry(\"gcr.io\"))\n\trequire.False(t, p.MatchesRegistry(\"myregistry.example.com\"))\n}\n\nfunc TestReplicateProvider_PostPush(t *testing.T) {\n\tp := New()\n\topts := provider.PushOptions{\n\t\tImage: \"r8.im/user/model\",\n\t}\n\n\tt.Run(\"success\", func(t *testing.T) {\n\t\terr := p.PostPush(context.Background(), opts, nil)\n\t\trequire.NoError(t, err)\n\t})\n\n\tt.Run(\"repository not found error\", func(t *testing.T) {\n\t\t// Simulate a NotFoundError from docker push (repository doesn't exist)\n\t\tpushErr := &command.NotFoundError{Ref: \"r8.im/user/model\", Object: \"repository\"}\n\t\terr := p.PostPush(context.Background(), opts, pushErr)\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"Unable to find existing Replicate model\")\n\t\trequire.Contains(t, err.Error(), \"replicate.com and create a new model\")\n\t})\n\n\tt.Run(\"tag not found error\", func(t *testing.T) {\n\t\t// Tag not found errors should also trigger the helpful message\n\t\tpushErr := &command.NotFoundError{Ref: \"r8.im/user/model:v1\", Object: \"tag\"}\n\t\terr := p.PostPush(context.Background(), opts, pushErr)\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"Unable to find existing Replicate model\")\n\t})\n}\n\nfunc TestCheckTokenFormat(t *testing.T) {\n\ttests := []struct {\n\t\tname    string\n\t\ttoken   string\n\t\twantErr bool\n\t}{\n\t\t{\n\t\t\tname:    \"valid CLI token\",\n\t\t\ttoken:   \"abc123def456\",\n\t\t\twantErr: false,\n\t\t},\n\t\t{\n\t\t\tname:    \"API token rejected\",\n\t\t\ttoken:   \"r8_abc123\",\n\t\t\twantErr: true,\n\t\t},\n\t\t{\n\t\t\tname:    \"empty token allowed (separate validation)\",\n\t\t\ttoken:   \"\",\n\t\t\twantErr: false,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\terr := checkTokenFormat(tt.token)\n\t\t\tif tt.wantErr {\n\t\t\t\trequire.Error(t, err)\n\t\t\t\trequire.Contains(t, err.Error(), \"API token\")\n\t\t\t} else {\n\t\t\t\trequire.NoError(t, err)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestVerifyToken(t *testing.T) {\n\tt.Run(\"successful verification\", func(t *testing.T) {\n\t\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\t\trequire.Equal(t, \"/cog/v1/verify-token\", r.URL.Path)\n\t\t\trequire.Equal(t, \"POST\", r.Method)\n\n\t\t\tw.WriteHeader(http.StatusOK)\n\t\t\tjson.NewEncoder(w).Encode(map[string]string{\"username\": \"testuser\"})\n\t\t}))\n\t\tdefer server.Close()\n\n\t\tusername, err := verifyToken(server.URL, \"valid-token\")\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"testuser\", username)\n\t})\n\n\tt.Run(\"empty token\", func(t *testing.T) {\n\t\t_, err := verifyToken(\"http://localhost\", \"\")\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"empty\")\n\t})\n\n\tt.Run(\"user not found\", func(t *testing.T) {\n\t\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\t\tw.WriteHeader(http.StatusNotFound)\n\t\t}))\n\t\tdefer server.Close()\n\n\t\t_, err := verifyToken(server.URL, \"unknown-token\")\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"does not exist\")\n\t})\n\n\tt.Run(\"server error\", func(t *testing.T) {\n\t\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\t\tw.WriteHeader(http.StatusInternalServerError)\n\t\t}))\n\t\tdefer server.Close()\n\n\t\t_, err := verifyToken(server.URL, \"some-token\")\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"500\")\n\t})\n}\n\nfunc TestGetDisplayTokenURL(t *testing.T) {\n\tt.Run(\"successful fetch\", func(t *testing.T) {\n\t\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\t\trequire.Equal(t, \"/cog/v1/display-token-url\", r.URL.Path)\n\t\t\trequire.Equal(t, \"GET\", r.Method)\n\n\t\t\tw.WriteHeader(http.StatusOK)\n\t\t\tjson.NewEncoder(w).Encode(map[string]string{\"url\": \"https://replicate.com/auth/token\"})\n\t\t}))\n\t\tdefer server.Close()\n\n\t\turl, err := getDisplayTokenURL(server.URL)\n\t\trequire.NoError(t, err)\n\t\trequire.Equal(t, \"https://replicate.com/auth/token\", url)\n\t})\n\n\tt.Run(\"not replicate registry\", func(t *testing.T) {\n\t\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\t\tw.WriteHeader(http.StatusNotFound)\n\t\t}))\n\t\tdefer server.Close()\n\n\t\t_, err := getDisplayTokenURL(server.URL)\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"not the Replicate registry\")\n\t})\n}\n\nfunc TestAddressWithScheme(t *testing.T) {\n\ttests := []struct {\n\t\tinput    string\n\t\texpected string\n\t}{\n\t\t{\"r8.im\", \"https://r8.im\"},\n\t\t{\"https://r8.im\", \"https://r8.im\"},\n\t\t{\"http://localhost:8080\", \"http://localhost:8080\"},\n\t\t{\"myregistry.com\", \"https://myregistry.com\"},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.input, func(t *testing.T) {\n\t\t\tresult := addressWithScheme(tt.input)\n\t\t\trequire.Equal(t, tt.expected, result)\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "pkg/provider/setup/setup.go",
    "content": "// Package setup initializes the default provider registry\npackage setup\n\nimport (\n\t\"sync\"\n\n\t\"github.com/replicate/cog/pkg/provider\"\n\t\"github.com/replicate/cog/pkg/provider/generic\"\n\t\"github.com/replicate/cog/pkg/provider/replicate\"\n)\n\nvar once sync.Once\n\n// Init initializes the default provider registry with all built-in providers\n// This function is idempotent - it only runs once even if called multiple times\nfunc Init() {\n\tonce.Do(func() {\n\t\tregistry := provider.DefaultRegistry()\n\n\t\t// Register Replicate provider first (more specific)\n\t\tregistry.Register(replicate.New())\n\n\t\t// Register Generic provider last (fallback for any OCI registry)\n\t\tregistry.Register(generic.New())\n\t})\n}\n"
  },
  {
    "path": "pkg/provider/setup/setup_test.go",
    "content": "package setup\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/provider\"\n)\n\nfunc TestInit(t *testing.T) {\n\t// Call Init multiple times - should be idempotent\n\tInit()\n\tInit()\n\n\tregistry := provider.DefaultRegistry()\n\n\t// Replicate images should get the Replicate provider\n\tp := registry.ForImage(\"r8.im/user/model\")\n\trequire.NotNil(t, p)\n\trequire.Equal(t, \"replicate\", p.Name())\n\n\t// Other images should get the Generic provider\n\tp = registry.ForImage(\"ghcr.io/owner/repo\")\n\trequire.NotNil(t, p)\n\trequire.Equal(t, \"generic\", p.Name())\n\n\t// Docker Hub images should get Generic provider\n\tp = registry.ForImage(\"nginx\")\n\trequire.NotNil(t, p)\n\trequire.Equal(t, \"generic\", p.Name())\n}\n"
  },
  {
    "path": "pkg/registry/client.go",
    "content": "package registry\n\nimport (\n\t\"context\"\n\t\"time\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote\"\n)\n\ntype Platform struct {\n\tOS           string\n\tArchitecture string\n\tVariant      string\n}\n\ntype PlatformManifest struct {\n\tDigest       string\n\tMediaType    string\n\tSize         int64\n\tOS           string\n\tArchitecture string\n\tVariant      string\n\tAnnotations  map[string]string\n}\n\n// RetryEvent contains information about a retry attempt.\ntype RetryEvent struct {\n\t// Attempt is the current retry attempt number (1-indexed).\n\tAttempt int\n\t// MaxAttempts is the maximum number of retry attempts.\n\tMaxAttempts int\n\t// Err is the error that caused the retry.\n\tErr error\n\t// NextRetryIn is the duration until the next retry attempt.\n\tNextRetryIn time.Duration\n}\n\n// RetryCallback is called when a retry occurs. Return false to abort retrying.\ntype RetryCallback func(event RetryEvent) bool\n\n// RetryConfig configures retry behavior for registry operations.\ntype RetryConfig struct {\n\t// Backoff configures the exponential backoff for retries.\n\t// If nil, the default backoff from go-containerregistry is used (3 attempts, 1s initial, 3x factor).\n\tBackoff *remote.Backoff\n\t// OnRetry is called when a retry occurs. If nil, no callback is invoked.\n\t// The callback receives information about the retry attempt.\n\tOnRetry RetryCallback\n}\n\n// WriteLayerOptions configures the WriteLayer operation.\ntype WriteLayerOptions struct {\n\t// Repo is the repository to push to.\n\tRepo string\n\t// Layer is the layer to push.\n\tLayer v1.Layer\n\t// ProgressCh receives progress updates. Use a buffered channel to avoid deadlocks.\n\t// If nil, no progress updates are sent.\n\tProgressCh chan<- v1.Update\n\t// Retry configures retry behavior. If nil, default retry behavior is used\n\t// (5 attempts with exponential backoff starting at 2 seconds).\n\tRetry *RetryConfig\n}\n\ntype Client interface {\n\t// Read methods\n\tInspect(ctx context.Context, imageRef string, platform *Platform) (*ManifestResult, error)\n\tGetImage(ctx context.Context, imageRef string, platform *Platform) (v1.Image, error)\n\tExists(ctx context.Context, imageRef string) (bool, error)\n\n\t// GetDescriptor returns the OCI descriptor for an image reference without downloading\n\t// the full image. This is a lightweight HEAD request useful for building OCI indexes\n\t// from already-pushed manifests.\n\tGetDescriptor(ctx context.Context, imageRef string) (v1.Descriptor, error)\n\n\t// Write methods for OCI index support\n\tPushImage(ctx context.Context, ref string, img v1.Image) error\n\tPushIndex(ctx context.Context, ref string, idx v1.ImageIndex) error\n\n\t// WriteLayer pushes a single layer (blob) to a repository with retry and optional progress reporting.\n\t// This method handles transient failures automatically with exponential backoff.\n\t// Use WriteLayerOptions to configure progress reporting and retry callbacks.\n\tWriteLayer(ctx context.Context, opts WriteLayerOptions) error\n}\n"
  },
  {
    "path": "pkg/registry/client_test.go",
    "content": "package registry\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"io\"\n\t\"net\"\n\t\"os\"\n\t\"syscall\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/registry_testhelpers\"\n)\n\nfunc TestInspect(t *testing.T) {\n\tif testing.Short() {\n\t\t// TODO[md]: this is a hack to skip the test in GitHub Actions because\n\t\t// because macos runners don't have rootless docker. this should get added back\n\t\t// and be part of a normal integration suite we run on all target platforms\n\t\tt.Skip(\"skipping integration tests\")\n\t}\n\n\tregistry := registry_testhelpers.StartTestRegistry(t)\n\n\tt.Run(\"it returns an index for multi-platform images when a platform isn't provided\", func(t *testing.T) {\n\t\timageRef := registry.ImageRef(\"alpine:latest\")\n\n\t\tclient := NewRegistryClient()\n\t\tresp, err := client.Inspect(t.Context(), imageRef, nil)\n\t\trequire.NoError(t, err)\n\t\trequire.NotNil(t, resp)\n\t\tassert.True(t, resp.IsIndex(), \"expected index\")\n\t\tjson.NewEncoder(os.Stdout).Encode(resp)\n\t})\n\n\tt.Run(\"it returns a single platform image when a platform is provided\", func(t *testing.T) {\n\t\timageRef := registry.ImageRef(\"alpine:latest\")\n\t\tclient := NewRegistryClient()\n\t\tresp, err := client.Inspect(t.Context(), imageRef, &Platform{OS: \"linux\", Architecture: \"amd64\"})\n\t\trequire.NoError(t, err)\n\t\trequire.NotNil(t, resp)\n\t\tassert.False(t, resp.IsIndex(), \"expected single platform image\")\n\t\tassert.True(t, resp.IsSinglePlatform(), \"expected single platform image\")\n\t\tjson.NewEncoder(os.Stdout).Encode(resp)\n\t})\n\n\tt.Run(\"when a repo does not exist\", func(t *testing.T) {\n\t\timageRef := registry.ImageRef(\"i-do-not-exist:latest\")\n\t\tclient := NewRegistryClient()\n\t\tresp, err := client.Inspect(t.Context(), imageRef, nil)\n\t\tassert.ErrorIs(t, err, NotFoundError, \"expected not found error\")\n\t\tassert.Nil(t, resp)\n\t})\n\n\tt.Run(\"when a repo with a slashdoes not exist\", func(t *testing.T) {\n\t\timageRef := registry.ImageRef(\"i-do-not-exist/with-a-slash:latest\")\n\t\tclient := NewRegistryClient()\n\t\tresp, err := client.Inspect(t.Context(), imageRef, nil)\n\t\tassert.ErrorIs(t, err, NotFoundError, \"expected not found error\")\n\t\tassert.Nil(t, resp)\n\t})\n\n\tt.Run(\"when the repo exists but the tag does not\", func(t *testing.T) {\n\t\timageRef := registry.ImageRef(\"alpine:not-found\")\n\t\tclient := NewRegistryClient()\n\t\tresp, err := client.Inspect(t.Context(), imageRef, nil)\n\t\tassert.ErrorIs(t, err, NotFoundError, \"expected not found error\")\n\t\tassert.Nil(t, resp)\n\t})\n\n\tt.Run(\"when the repo and tag exist but platform does not\", func(t *testing.T) {\n\t\timageRef := registry.ImageRef(\"alpine:latest\")\n\t\tclient := NewRegistryClient()\n\t\tresp, err := client.Inspect(t.Context(), imageRef, &Platform{OS: \"windows\", Architecture: \"i386\"})\n\t\tassert.ErrorContains(t, err, \"platform not found\")\n\t\tassert.Nil(t, resp)\n\t})\n}\n\nfunc TestIsRetryableError(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\terr      error\n\t\texpected bool\n\t}{\n\t\t{\n\t\t\tname:     \"nil error\",\n\t\t\terr:      nil,\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"io.EOF\",\n\t\t\terr:      io.EOF,\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"io.ErrUnexpectedEOF\",\n\t\t\terr:      io.ErrUnexpectedEOF,\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"syscall.EPIPE (broken pipe)\",\n\t\t\terr:      syscall.EPIPE,\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"syscall.ECONNRESET\",\n\t\t\terr:      syscall.ECONNRESET,\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"net.ErrClosed\",\n\t\t\terr:      net.ErrClosed,\n\t\t\texpected: true,\n\t\t},\n\t\t{\n\t\t\tname:     \"context.Canceled\",\n\t\t\terr:      context.Canceled,\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"context.DeadlineExceeded\",\n\t\t\terr:      context.DeadlineExceeded,\n\t\t\texpected: false,\n\t\t},\n\t\t{\n\t\t\tname:     \"generic error (not retryable)\",\n\t\t\terr:      errors.New(\"something completely different\"),\n\t\t\texpected: false,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := isRetryableError(tt.err)\n\t\t\tassert.Equal(t, tt.expected, result, \"isRetryableError(%v) = %v, want %v\", tt.err, result, tt.expected)\n\t\t})\n\t}\n}\n"
  },
  {
    "path": "pkg/registry/config.go",
    "content": "package registry\n\nimport (\n\t\"os\"\n\t\"strconv\"\n)\n\nconst (\n\t// DefaultChunkSize is the size (in bytes) of each chunk in a multipart upload.\n\t// This is used as a fallback when the registry does not advertise chunk size\n\t// limits via OCI-Chunk-Min-Length / OCI-Chunk-Max-Length headers.\n\t// 96 MB stays under common CDN/proxy request body limits while still being\n\t// large enough to reduce HTTP round-trips for multi-GB files.\n\tDefaultChunkSize = 96 * 1024 * 1024 // 96 MB\n\n\t// DefaultMultipartThreshold is the minimum blob size (in bytes) before using multipart upload.\n\t// Blobs smaller than this are uploaded in a single request to avoid multipart overhead.\n\t// Set higher than DefaultChunkSize so that blobs that would fit in a single chunk\n\t// are uploaded in one request, avoiding unnecessary multipart overhead.\n\tDefaultMultipartThreshold = 128 * 1024 * 1024 // 128 MB\n\n\t// chunkSizeMargin is subtracted from the server's OCI-Chunk-Max-Length to stay\n\t// safely under the limit (e.g. for HTTP framing overhead).\n\tchunkSizeMargin = 64 * 1024 // 64 KB\n\n\t// envPushDefaultChunkSize sets the default chunk size for multipart uploads.\n\t// This is only used when the registry does not advertise OCI-Chunk-Max-Length.\n\t// When the registry does advertise a maximum, the server's limit takes precedence.\n\tenvPushDefaultChunkSize = \"COG_PUSH_DEFAULT_CHUNK_SIZE\"\n\n\t// envMultipartThreshold overrides the minimum blob size for multipart uploads.\n\tenvMultipartThreshold = \"COG_PUSH_MULTIPART_THRESHOLD\"\n)\n\n// getDefaultChunkSize returns the client-configured default chunk size for multipart uploads.\n// This is used as a fallback when the registry does not advertise chunk size limits.\nfunc getDefaultChunkSize() int64 {\n\tif v := os.Getenv(envPushDefaultChunkSize); v != \"\" {\n\t\tif n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {\n\t\t\treturn n\n\t\t}\n\t}\n\treturn DefaultChunkSize\n}\n\n// getMultipartThreshold returns the minimum blob size for multipart uploads.\nfunc getMultipartThreshold() int64 {\n\tif v := os.Getenv(envMultipartThreshold); v != \"\" {\n\t\tif n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {\n\t\t\treturn n\n\t\t}\n\t}\n\treturn DefaultMultipartThreshold\n}\n"
  },
  {
    "path": "pkg/registry/config_test.go",
    "content": "package registry\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n)\n\nfunc TestGetDefaultChunkSize(t *testing.T) {\n\tt.Run(\"returns default when env not set\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"\")\n\t\tassert.Equal(t, int64(DefaultChunkSize), getDefaultChunkSize())\n\t})\n\n\tt.Run(\"returns env var value\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"134217728\") // 128 MB\n\t\tassert.Equal(t, int64(134217728), getDefaultChunkSize())\n\t})\n\n\tt.Run(\"returns default for invalid value\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"xyz\")\n\t\tassert.Equal(t, int64(DefaultChunkSize), getDefaultChunkSize())\n\t})\n\n\tt.Run(\"returns default for zero\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"0\")\n\t\tassert.Equal(t, int64(DefaultChunkSize), getDefaultChunkSize())\n\t})\n\n\tt.Run(\"returns default for negative\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"-100\")\n\t\tassert.Equal(t, int64(DefaultChunkSize), getDefaultChunkSize())\n\t})\n}\n\nfunc TestEffectiveChunkSize(t *testing.T) {\n\tt.Run(\"uses client default when server provides no limits\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"\")\n\t\ts := uploadSession{}\n\t\tassert.Equal(t, int64(DefaultChunkSize), s.effectiveChunkSize())\n\t})\n\n\tt.Run(\"uses env var default when server provides no limits\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"50000000\") // 50 MB\n\t\ts := uploadSession{}\n\t\tassert.Equal(t, int64(50000000), s.effectiveChunkSize())\n\t})\n\n\tt.Run(\"server max takes precedence over client default\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"\")\n\t\tserverMax := int64(90 * 1024 * 1024)\n\t\ts := uploadSession{ChunkMaxBytes: serverMax}\n\t\texpected := serverMax - chunkSizeMargin\n\t\tassert.Equal(t, expected, s.effectiveChunkSize())\n\t})\n\n\tt.Run(\"server max takes precedence even when larger than client default\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"\")\n\t\t// Server max of 200 MB -- server still dictates, not the client default\n\t\tserverMax := int64(200 * 1024 * 1024)\n\t\ts := uploadSession{ChunkMaxBytes: serverMax}\n\t\texpected := serverMax - chunkSizeMargin\n\t\tassert.Equal(t, expected, s.effectiveChunkSize())\n\t})\n\n\tt.Run(\"server max takes precedence over env var\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"50000000\") // 50 MB -- ignored when server provides max\n\t\tserverMax := int64(100 * 1000 * 1000)         // 100 MB\n\t\ts := uploadSession{ChunkMaxBytes: serverMax}\n\t\texpected := serverMax - chunkSizeMargin\n\t\tassert.Equal(t, expected, s.effectiveChunkSize())\n\t})\n\n\tt.Run(\"handles very small server max gracefully\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"\")\n\t\ts := uploadSession{ChunkMaxBytes: 1000} // 1000 bytes, smaller than margin\n\t\t// Margin is bigger than max, so we use the max directly\n\t\tassert.Equal(t, int64(1000), s.effectiveChunkSize())\n\t})\n\n\tt.Run(\"server min does not raise chunk size when already above it\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"\")\n\t\t// Server says min=5MiB max=90MiB; max-margin is well above min, so min has no effect\n\t\tserverMax := int64(90 * 1024 * 1024)\n\t\ts := uploadSession{ChunkMinBytes: 5 * 1024 * 1024, ChunkMaxBytes: serverMax}\n\t\texpected := serverMax - chunkSizeMargin\n\t\tassert.Equal(t, expected, s.effectiveChunkSize())\n\t})\n\n\tt.Run(\"server min clamps up a too-small client default\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"1000\") // 1 KB, below server min\n\t\tserverMin := int64(5 * 1024 * 1024)       // 5 MiB\n\t\ts := uploadSession{ChunkMinBytes: serverMin}\n\t\tassert.Equal(t, serverMin, s.effectiveChunkSize())\n\t})\n\n\tt.Run(\"server min clamps up when max minus margin falls below min\", func(t *testing.T) {\n\t\tt.Setenv(envPushDefaultChunkSize, \"\")\n\t\t// Contrived: max is just above min, so max-margin < min. Min should win.\n\t\tserverMin := int64(5 * 1024 * 1024)\n\t\tserverMax := serverMin + chunkSizeMargin/2 // max - margin < min\n\t\ts := uploadSession{ChunkMinBytes: serverMin, ChunkMaxBytes: serverMax}\n\t\tassert.Equal(t, serverMin, s.effectiveChunkSize())\n\t})\n}\n\nfunc TestGetMultipartThreshold(t *testing.T) {\n\tt.Run(\"returns default when env not set\", func(t *testing.T) {\n\t\tt.Setenv(envMultipartThreshold, \"\")\n\t\tassert.Equal(t, int64(DefaultMultipartThreshold), getMultipartThreshold())\n\t})\n\n\tt.Run(\"returns env var value\", func(t *testing.T) {\n\t\tt.Setenv(envMultipartThreshold, \"104857600\") // 100 MB\n\t\tassert.Equal(t, int64(104857600), getMultipartThreshold())\n\t})\n\n\tt.Run(\"returns default for invalid value\", func(t *testing.T) {\n\t\tt.Setenv(envMultipartThreshold, \"abc\")\n\t\tassert.Equal(t, int64(DefaultMultipartThreshold), getMultipartThreshold())\n\t})\n\n\tt.Run(\"returns default for zero\", func(t *testing.T) {\n\t\tt.Setenv(envMultipartThreshold, \"0\")\n\t\tassert.Equal(t, int64(DefaultMultipartThreshold), getMultipartThreshold())\n\t})\n\n\tt.Run(\"returns default for negative\", func(t *testing.T) {\n\t\tt.Setenv(envMultipartThreshold, \"-50\")\n\t\tassert.Equal(t, int64(DefaultMultipartThreshold), getMultipartThreshold())\n\t})\n}\n"
  },
  {
    "path": "pkg/registry/manifest_result.go",
    "content": "package registry\n\nimport \"github.com/google/go-containerregistry/pkg/v1/types\"\n\ntype ManifestResult struct {\n\tSchemaVersion int64\n\tMediaType     string\n\t// Digest is the content-addressable digest of the manifest (sha256:...).\n\tDigest string\n\n\tManifests []PlatformManifest\n\tLayers    []string\n\tConfig    string\n\tLabels    map[string]string\n}\n\nfunc (m *ManifestResult) IsIndex() bool {\n\treturn m.MediaType == string(types.OCIImageIndex) || m.MediaType == string(types.DockerManifestList)\n}\n\nfunc (m *ManifestResult) IsSinglePlatform() bool {\n\treturn !m.IsIndex()\n}\n"
  },
  {
    "path": "pkg/registry/push_test.go",
    "content": "package registry\n\nimport (\n\t\"context\"\n\t\"testing\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/empty\"\n\t\"github.com/google/go-containerregistry/pkg/v1/mutate\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/registry_testhelpers\"\n)\n\nfunc TestPushOperations(t *testing.T) {\n\tif testing.Short() {\n\t\tt.Skip(\"skipping integration tests in short mode\")\n\t}\n\n\t// Start a test registry using testcontainers\n\tregistry := registry_testhelpers.StartTestRegistry(t)\n\tregistryAddr := registry.RegistryHost()\n\n\tctx := context.Background()\n\tclient := NewRegistryClient()\n\n\tt.Run(\"push image\", func(t *testing.T) {\n\t\timg := empty.Image\n\t\timg, _ = mutate.Config(img, v1.Config{})\n\n\t\terr := client.PushImage(ctx, registryAddr+\"/test/push-test:v1\", img)\n\t\trequire.NoError(t, err)\n\n\t\t// Verify it exists\n\t\texists, err := client.Exists(ctx, registryAddr+\"/test/push-test:v1\")\n\t\trequire.NoError(t, err)\n\t\trequire.True(t, exists)\n\t})\n\n\tt.Run(\"push index\", func(t *testing.T) {\n\t\timg := empty.Image\n\t\timg, _ = mutate.Config(img, v1.Config{})\n\n\t\t// Push the child image first — PushIndex only writes the index manifest,\n\t\t// it does not recursively push child manifests/blobs.\n\t\terr := client.PushImage(ctx, registryAddr+\"/test/push-idx:child\", img)\n\t\trequire.NoError(t, err)\n\n\t\tidx := mutate.IndexMediaType(empty.Index, types.OCIImageIndex)\n\t\tidx = mutate.AppendManifests(idx, mutate.IndexAddendum{\n\t\t\tAdd:        img,\n\t\t\tDescriptor: v1.Descriptor{Platform: &v1.Platform{OS: \"linux\", Architecture: \"amd64\"}},\n\t\t})\n\n\t\terr = client.PushIndex(ctx, registryAddr+\"/test/push-idx:v1\", idx)\n\t\trequire.NoError(t, err)\n\n\t\t// Verify it exists\n\t\texists, err := client.Exists(ctx, registryAddr+\"/test/push-idx:v1\")\n\t\trequire.NoError(t, err)\n\t\trequire.True(t, exists)\n\t})\n}\n"
  },
  {
    "path": "pkg/registry/registry_client.go",
    "content": "package registry\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"crypto/tls\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"math/rand\"\n\t\"net\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"slices\"\n\t\"strconv\"\n\t\"sync\"\n\t\"syscall\"\n\t\"time\"\n\n\t\"github.com/google/go-containerregistry/pkg/authn\"\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote/transport\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n)\n\n//nolint:staticcheck // ST1012: exported API, renaming would be breaking change\nvar NotFoundError = errors.New(\"image reference not found\")\n\n// chunkBufPool pools byte slices used for multipart chunk reads to avoid\n// re-allocating large buffers (up to ~96 MB each) on every layer upload.\nvar chunkBufPool = sync.Pool{} //nolint:gochecknoglobals\n\ntype RegistryClient struct {\n\t// transport is the HTTP transport used for all registry operations.\n\t// Uses HTTP/1.1 only to avoid HTTP/2 head-of-line blocking and stream\n\t// errors (RST_STREAM INTERNAL_ERROR) that occur when uploading large\n\t// blobs through CDN/proxy edges.\n\ttransport http.RoundTripper\n}\n\nfunc NewRegistryClient() Client {\n\treturn &RegistryClient{\n\t\ttransport: http1OnlyTransport(),\n\t}\n}\n\n// remoteOptions returns the common remote.Option set for go-containerregistry calls,\n// including authentication, context, and HTTP/1.1 transport.\nfunc (c *RegistryClient) remoteOptions(ctx context.Context) []remote.Option {\n\treturn []remote.Option{\n\t\tremote.WithContext(ctx),\n\t\tremote.WithAuthFromKeychain(authn.DefaultKeychain),\n\t\tremote.WithTransport(c.transport),\n\t}\n}\n\nfunc (c *RegistryClient) Inspect(ctx context.Context, imageRef string, platform *Platform) (*ManifestResult, error) {\n\tref, err := name.ParseReference(imageRef, name.Insecure)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"parsing reference: %w\", err)\n\t}\n\n\tdesc, err := remote.Get(ref, c.remoteOptions(ctx)...)\n\tif err != nil {\n\t\tif checkError(err, transport.ManifestUnknownErrorCode, transport.NameUnknownErrorCode) {\n\t\t\treturn nil, NotFoundError\n\t\t}\n\n\t\treturn nil, fmt.Errorf(\"fetching descriptor: %w\", err)\n\t}\n\n\tmediaType := desc.MediaType\n\n\tif platform == nil {\n\t\tswitch mediaType {\n\t\tcase types.OCIImageIndex, types.DockerManifestList:\n\t\t\tidx, err := desc.ImageIndex()\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"loading image index: %w\", err)\n\t\t\t}\n\t\t\tindexManifest, err := idx.IndexManifest()\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"getting index manifest: %w\", err)\n\t\t\t}\n\t\t\tresult := &ManifestResult{\n\t\t\t\tSchemaVersion: indexManifest.SchemaVersion,\n\t\t\t\tMediaType:     string(mediaType),\n\t\t\t\tDigest:        desc.Digest.String(),\n\t\t\t}\n\t\t\tfor _, m := range indexManifest.Manifests {\n\t\t\t\tresult.Manifests = append(result.Manifests, PlatformManifest{\n\t\t\t\t\tDigest:       m.Digest.String(),\n\t\t\t\t\tMediaType:    string(m.MediaType),\n\t\t\t\t\tSize:         m.Size,\n\t\t\t\t\tOS:           m.Platform.OS,\n\t\t\t\t\tArchitecture: m.Platform.Architecture,\n\t\t\t\t\tVariant:      m.Platform.Variant,\n\t\t\t\t\tAnnotations:  m.Annotations,\n\t\t\t\t})\n\t\t\t}\n\t\t\t// For indexes, pick a default image to get labels from.\n\t\t\t// Prefer linux/amd64, otherwise use the first manifest.\n\t\t\tdefaultImg, err := pickDefaultImage(ref, indexManifest, c.remoteOptions(ctx)...)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"failed to read image config from index: %w\", err)\n\t\t\t}\n\t\t\tconfigFile, err := defaultImg.ConfigFile()\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"failed to get image config: %w\", err)\n\t\t\t}\n\t\t\tresult.Labels = configFile.Config.Labels\n\t\t\treturn result, nil\n\n\t\tcase types.OCIManifestSchema1, types.DockerManifestSchema2:\n\t\t\timg, err := desc.Image()\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"loading image: %w\", err)\n\t\t\t}\n\t\t\tmanifest, err := img.Manifest()\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"getting manifest: %w\", err)\n\t\t\t}\n\t\t\tconfigFile, err := img.ConfigFile()\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"getting config file: %w\", err)\n\t\t\t}\n\t\t\tresult := &ManifestResult{\n\t\t\t\tSchemaVersion: manifest.SchemaVersion,\n\t\t\t\tMediaType:     string(mediaType),\n\t\t\t\tDigest:        desc.Digest.String(),\n\t\t\t\tConfig:        manifest.Config.Digest.String(),\n\t\t\t\tLabels:        configFile.Config.Labels,\n\t\t\t}\n\t\t\tfor _, layer := range manifest.Layers {\n\t\t\t\tresult.Layers = append(result.Layers, layer.Digest.String())\n\t\t\t}\n\t\t\treturn result, nil\n\t\tdefault:\n\t\t\treturn nil, fmt.Errorf(\"unsupported media type: %s\", mediaType)\n\t\t}\n\t}\n\n\t// platform is set, we expect a manifest list or error\n\tif mediaType != types.OCIImageIndex && mediaType != types.DockerManifestList {\n\t\treturn nil, fmt.Errorf(\"image is not a manifest list but platform was specified\")\n\t}\n\n\tidx, err := desc.ImageIndex()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"loading image index: %w\", err)\n\t}\n\tindexManifest, err := idx.IndexManifest()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"getting index manifest: %w\", err)\n\t}\n\n\tvar matchedDigest string\n\tfor _, m := range indexManifest.Manifests {\n\t\tif m.Platform.OS == platform.OS &&\n\t\t\tm.Platform.Architecture == platform.Architecture &&\n\t\t\tm.Platform.Variant == platform.Variant {\n\t\t\tmatchedDigest = m.Digest.String()\n\t\t\tbreak\n\t\t}\n\t}\n\n\tif matchedDigest == \"\" {\n\t\treturn nil, fmt.Errorf(\"platform not found in manifest list\")\n\t}\n\n\tdigestRef, err := name.NewDigest(ref.Context().Name() + \"@\" + matchedDigest)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"creating digest ref: %w\", err)\n\t}\n\tmanifestDesc, err := remote.Get(digestRef, c.remoteOptions(ctx)...)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"fetching platform manifest: %w\", err)\n\t}\n\timg, err := manifestDesc.Image()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"loading platform image: %w\", err)\n\t}\n\tmanifest, err := img.Manifest()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"getting manifest: %w\", err)\n\t}\n\tconfigFile, err := img.ConfigFile()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"getting config file: %w\", err)\n\t}\n\tresult := &ManifestResult{\n\t\tSchemaVersion: manifest.SchemaVersion,\n\t\tMediaType:     string(manifestDesc.MediaType),\n\t\tDigest:        manifestDesc.Digest.String(),\n\t\tConfig:        manifest.Config.Digest.String(),\n\t\tLabels:        configFile.Config.Labels,\n\t}\n\tfor _, layer := range manifest.Layers {\n\t\tresult.Layers = append(result.Layers, layer.Digest.String())\n\t}\n\treturn result, nil\n}\n\nfunc (c *RegistryClient) GetImage(ctx context.Context, imageRef string, platform *Platform) (v1.Image, error) {\n\tref, err := name.ParseReference(imageRef, name.Insecure)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"parsing reference: %w\", err)\n\t}\n\n\tdesc, err := remote.Get(ref, c.remoteOptions(ctx)...)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"fetching descriptor: %w\", err)\n\t}\n\n\tmediaType := desc.MediaType\n\n\t// If no platform is specified and it's a single image, return it directly\n\tif platform == nil {\n\t\tswitch mediaType {\n\t\tcase types.OCIManifestSchema1, types.DockerManifestSchema2:\n\t\t\treturn desc.Image()\n\t\tcase types.OCIImageIndex, types.DockerManifestList:\n\t\t\treturn nil, fmt.Errorf(\"platform must be specified for multi-platform image\")\n\t\tdefault:\n\t\t\treturn nil, fmt.Errorf(\"unsupported media type: %s\", mediaType)\n\t\t}\n\t}\n\n\t// For platform-specific requests, we need to handle manifest lists\n\tif mediaType != types.OCIImageIndex && mediaType != types.DockerManifestList {\n\t\treturn nil, fmt.Errorf(\"image is not a manifest list but platform was specified\")\n\t}\n\n\tidx, err := desc.ImageIndex()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"loading image index: %w\", err)\n\t}\n\n\tindexManifest, err := idx.IndexManifest()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"getting index manifest: %w\", err)\n\t}\n\n\t// Find the matching platform manifest\n\tvar matchedDigest string\n\tfor _, m := range indexManifest.Manifests {\n\t\tif m.Platform.OS == platform.OS &&\n\t\t\tm.Platform.Architecture == platform.Architecture &&\n\t\t\tm.Platform.Variant == platform.Variant {\n\t\t\tmatchedDigest = m.Digest.String()\n\t\t\tbreak\n\t\t}\n\t}\n\n\tif matchedDigest == \"\" {\n\t\treturn nil, fmt.Errorf(\"platform not found in manifest list\")\n\t}\n\n\t// Get the image for the matched digest\n\tdigestRef, err := name.NewDigest(ref.Context().Name() + \"@\" + matchedDigest)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"creating digest ref: %w\", err)\n\t}\n\n\tmanifestDesc, err := remote.Get(digestRef, c.remoteOptions(ctx)...)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"fetching platform manifest: %w\", err)\n\t}\n\n\treturn manifestDesc.Image()\n}\n\n// GetDescriptor returns the OCI descriptor for an image reference using a HEAD request.\n// This is lightweight — it does not download the full manifest or image layers.\nfunc (c *RegistryClient) GetDescriptor(ctx context.Context, imageRef string) (v1.Descriptor, error) {\n\tref, err := name.ParseReference(imageRef, name.Insecure)\n\tif err != nil {\n\t\treturn v1.Descriptor{}, fmt.Errorf(\"parsing reference: %w\", err)\n\t}\n\n\tdesc, err := remote.Head(ref, c.remoteOptions(ctx)...)\n\tif err != nil {\n\t\treturn v1.Descriptor{}, fmt.Errorf(\"head request for %s: %w\", imageRef, err)\n\t}\n\n\treturn *desc, nil\n}\n\nfunc (c *RegistryClient) Exists(ctx context.Context, imageRef string) (bool, error) {\n\tif _, err := c.Inspect(ctx, imageRef, nil); err != nil {\n\t\tif errors.Is(err, NotFoundError) {\n\t\t\treturn false, nil\n\t\t}\n\t\treturn false, err\n\t}\n\treturn true, nil\n}\n\nfunc checkError(err error, codes ...transport.ErrorCode) bool {\n\tif err == nil {\n\t\treturn false\n\t}\n\n\tvar e *transport.Error\n\tif errors.As(err, &e) {\n\t\tfor _, diagnosticErr := range e.Errors {\n\t\t\tif slices.Contains(codes, diagnosticErr.Code) {\n\t\t\t\treturn true\n\t\t\t}\n\t\t}\n\t}\n\treturn false\n}\n\n// PushImage pushes a single image to a registry.\nfunc (c *RegistryClient) PushImage(ctx context.Context, ref string, img v1.Image) error {\n\tparsedRef, err := name.ParseReference(ref, name.Insecure)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"parsing reference: %w\", err)\n\t}\n\n\tif err := remote.Write(parsedRef, img, c.remoteOptions(ctx)...); err != nil {\n\t\treturn fmt.Errorf(\"pushing image %s: %w\", ref, err)\n\t}\n\n\treturn nil\n}\n\n// PushIndex pushes an OCI Image Index to a registry.\nfunc (c *RegistryClient) PushIndex(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\tparsedRef, err := name.ParseReference(ref, name.Insecure)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"parsing reference: %w\", err)\n\t}\n\n\t// Use remote.Put instead of remote.WriteIndex because all child manifests\n\t// (image + weights) are already pushed to the registry. WriteIndex would\n\t// try to recursively resolve and push children via idx.Image(), which fails\n\t// for our descriptor-only index. Put just writes the index manifest.\n\tif err := remote.Put(parsedRef, idx, c.remoteOptions(ctx)...); err != nil {\n\t\treturn fmt.Errorf(\"pushing index %s: %w\", ref, err)\n\t}\n\n\treturn nil\n}\n\n// http1OnlyTransport returns an http.Transport that only speaks HTTP/1.1.\n// HTTP/2 is avoided for all registry operations because high-throughput uploads\n// suffer from head-of-line blocking and stream errors (RST_STREAM INTERNAL_ERROR)\n// when pushed through CDN/proxy edges. Multiple concurrent HTTP/1.1 connections\n// outperform a single multiplexed HTTP/2 connection for large blob uploads.\nfunc http1OnlyTransport() *http.Transport {\n\tt := http.DefaultTransport.(*http.Transport).Clone()\n\tt.TLSClientConfig = tlsConfigHTTP1Only(t.TLSClientConfig)\n\t// ForceAttemptHTTP2 is true by default on cloned transports; disable it.\n\tt.ForceAttemptHTTP2 = false\n\treturn t\n}\n\n// tlsConfigHTTP1Only returns a TLS config that only advertises HTTP/1.1 via ALPN.\nfunc tlsConfigHTTP1Only(base *tls.Config) *tls.Config {\n\tif base == nil {\n\t\tbase = &tls.Config{MinVersion: tls.VersionTLS12}\n\t}\n\tcfg := base.Clone()\n\tcfg.NextProtos = []string{\"http/1.1\"}\n\treturn cfg\n}\n\n// DefaultRetryBackoff returns the default retry backoff configuration for weight pushes.\n// It retries 5 times with exponential backoff starting at 2 seconds.\nfunc DefaultRetryBackoff() remote.Backoff {\n\treturn remote.Backoff{\n\t\tDuration: 2 * time.Second,\n\t\tFactor:   2.0,\n\t\tJitter:   0.1,\n\t\tSteps:    5,\n\t}\n}\n\n// isRetryableError determines if an error should trigger a retry.\n// This matches the go-containerregistry default retry predicate plus additional cases.\nfunc isRetryableError(err error) bool {\n\tif err == nil {\n\t\treturn false\n\t}\n\n\t// Check for context cancellation - don't retry these\n\tif errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {\n\t\treturn false\n\t}\n\n\t// Check for temporary errors (network issues, etc.)\n\tvar tempErr interface{ Temporary() bool }\n\tif errors.As(err, &tempErr) && tempErr.Temporary() {\n\t\treturn true\n\t}\n\n\t// Check for common transient errors\n\tif errors.Is(err, io.ErrUnexpectedEOF) ||\n\t\terrors.Is(err, io.EOF) ||\n\t\terrors.Is(err, syscall.EPIPE) ||\n\t\terrors.Is(err, syscall.ECONNRESET) ||\n\t\terrors.Is(err, net.ErrClosed) {\n\t\treturn true\n\t}\n\n\t// Check for retryable HTTP status codes in transport errors\n\tvar transportErr *transport.Error\n\tif errors.As(err, &transportErr) {\n\t\tswitch transportErr.StatusCode {\n\t\tcase http.StatusRequestTimeout,\n\t\t\thttp.StatusInternalServerError,\n\t\t\thttp.StatusBadGateway,\n\t\t\thttp.StatusServiceUnavailable,\n\t\t\thttp.StatusGatewayTimeout,\n\t\t\t499, // nginx-specific, client closed request\n\t\t\t522: // Cloudflare-specific, connection timeout\n\t\t\treturn true\n\t\t}\n\t}\n\n\t// Check for network operation errors (connection refused, timeout, etc.)\n\tvar netErr *net.OpError\n\tif errors.As(err, &netErr) {\n\t\treturn true\n\t}\n\n\t// Check for DNS errors\n\tvar dnsErr *net.DNSError\n\tif errors.As(err, &dnsErr) {\n\t\treturn dnsErr.Temporary()\n\t}\n\n\treturn false\n}\n\n// WriteLayer pushes a single layer with retry and optional progress reporting.\n// This implements retry at the application level with callbacks for CLI feedback.\n// Unlike the standard remote.WriteLayer, this implementation performs multipart uploads\n// using Content-Range headers to upload the blob in chunks.\nfunc (c *RegistryClient) WriteLayer(ctx context.Context, opts WriteLayerOptions) error {\n\tparsedRepo, err := name.NewRepository(opts.Repo, name.Insecure)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"parsing repository: %w\", err)\n\t}\n\n\t// Determine retry configuration\n\tbackoff := DefaultRetryBackoff()\n\tif opts.Retry != nil && opts.Retry.Backoff != nil {\n\t\tbackoff = *opts.Retry.Backoff\n\t}\n\n\tvar lastErr error\n\tcurrentDelay := backoff.Duration\n\n\tfor attempt := 1; attempt <= backoff.Steps; attempt++ {\n\t\t// Check for context cancellation\n\t\tselect {\n\t\tcase <-ctx.Done():\n\t\t\treturn ctx.Err()\n\t\tdefault:\n\t\t}\n\n\t\t// Attempt the push using custom multipart upload\n\t\terr := c.writeLayerMultipart(ctx, parsedRepo, opts)\n\t\tif err == nil {\n\t\t\treturn nil // Success\n\t\t}\n\n\t\tlastErr = err\n\n\t\t// Check if this error is retryable\n\t\tif !isRetryableError(err) {\n\t\t\treturn fmt.Errorf(\"pushing layer to %s: %w\", opts.Repo, err)\n\t\t}\n\n\t\t// Don't retry if this was the last attempt\n\t\tif attempt >= backoff.Steps {\n\t\t\tbreak\n\t\t}\n\n\t\t// Calculate next delay with randomized jitter to avoid thundering herd\n\t\tnextDelay := currentDelay\n\t\tif backoff.Jitter > 0 {\n\t\t\tjitterAmount := time.Duration(float64(currentDelay) * backoff.Jitter * rand.Float64()) //nolint:gosec\n\t\t\tnextDelay = currentDelay + jitterAmount\n\t\t}\n\n\t\t// Invoke retry callback if configured\n\t\tif opts.Retry != nil && opts.Retry.OnRetry != nil {\n\t\t\tevent := RetryEvent{\n\t\t\t\tAttempt:     attempt,\n\t\t\t\tMaxAttempts: backoff.Steps,\n\t\t\t\tErr:         err,\n\t\t\t\tNextRetryIn: nextDelay,\n\t\t\t}\n\t\t\tif !opts.Retry.OnRetry(event) {\n\t\t\t\t// Callback returned false, abort retrying\n\t\t\t\treturn fmt.Errorf(\"pushing layer to %s (retry aborted): %w\", opts.Repo, err)\n\t\t\t}\n\t\t}\n\n\t\t// Wait before retrying\n\t\tselect {\n\t\tcase <-ctx.Done():\n\t\t\treturn ctx.Err()\n\t\tcase <-time.After(nextDelay):\n\t\t}\n\n\t\t// Update delay for next iteration\n\t\tcurrentDelay = time.Duration(float64(currentDelay) * backoff.Factor)\n\t}\n\n\treturn fmt.Errorf(\"pushing layer to %s (after %d attempts): %w\", opts.Repo, backoff.Steps, lastErr)\n}\n\n// writeLayerMultipart uploads a layer using multipart uploads with Content-Range headers.\n// This is a custom implementation that supports chunked uploads compatible with the\n// server-side code provided.\nfunc (c *RegistryClient) writeLayerMultipart(ctx context.Context, repo name.Repository, opts WriteLayerOptions) error {\n\t// Get layer metadata\n\tdigest, err := opts.Layer.Digest()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"getting layer digest: %w\", err)\n\t}\n\n\tsize, err := opts.Layer.Size()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"getting layer size: %w\", err)\n\t}\n\n\t// Create authenticated HTTP client\n\tauth, err := authn.Resolve(ctx, authn.DefaultKeychain, repo)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"resolving auth: %w\", err)\n\t}\n\n\tscopes := []string{repo.Scope(transport.PushScope)}\n\ttr, err := transport.NewWithContext(ctx, repo.Registry, auth, c.transport, scopes)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"creating transport: %w\", err)\n\t}\n\n\tclient := &http.Client{Transport: tr}\n\n\t// Check if blob already exists\n\texists, err := c.checkBlobExists(ctx, client, repo, digest)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"checking blob existence: %w\", err)\n\t}\n\tif exists {\n\t\tif opts.ProgressCh != nil {\n\t\t\topts.ProgressCh <- v1.Update{Complete: size, Total: size}\n\t\t}\n\t\treturn nil\n\t}\n\n\t// Initiate upload\n\tsession, err := c.initiateUpload(ctx, client, repo)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"initiating upload: %w\", err)\n\t}\n\n\t// Upload the blob in chunks\n\tfinalLocation, err := c.uploadBlobChunks(ctx, client, repo, opts.Layer, session, size, opts.ProgressCh)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"uploading blob chunks: %w\", err)\n\t}\n\n\t// Commit the upload using the final location (which contains updated state hash)\n\terr = c.commitUpload(ctx, client, finalLocation, digest)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"committing upload: %w\", err)\n\t}\n\n\treturn nil\n}\n\n// checkBlobExists checks if a blob already exists in the repository.\nfunc (c *RegistryClient) checkBlobExists(ctx context.Context, client *http.Client, repo name.Repository, digest v1.Hash) (bool, error) {\n\tu := url.URL{\n\t\tScheme: repo.Scheme(),\n\t\tHost:   repo.RegistryStr(),\n\t\tPath:   fmt.Sprintf(\"/v2/%s/blobs/%s\", repo.RepositoryStr(), digest.String()),\n\t}\n\n\treq, err := http.NewRequestWithContext(ctx, http.MethodHead, u.String(), nil)\n\tif err != nil {\n\t\treturn false, err\n\t}\n\n\tresp, err := client.Do(req) //nolint:gosec // G704: URL from registry reference, not user input\n\tif err != nil {\n\t\treturn false, err\n\t}\n\tdefer resp.Body.Close()\n\n\tif err := transport.CheckError(resp, http.StatusOK, http.StatusNotFound); err != nil {\n\t\treturn false, err\n\t}\n\n\treturn resp.StatusCode == http.StatusOK, nil\n}\n\n// uploadSession holds the result of initiating a blob upload, including the\n// upload location URL and any server-advertised chunk size constraints.\ntype uploadSession struct {\n\t// Location is the URL to which blob data should be uploaded.\n\tLocation string\n\t// ChunkMinBytes is the minimum chunk size the server accepts (from OCI-Chunk-Min-Length).\n\t// Zero means the server did not advertise a minimum.\n\tChunkMinBytes int64\n\t// ChunkMaxBytes is the maximum chunk size the server accepts (from OCI-Chunk-Max-Length).\n\t// Zero means the server did not advertise a maximum.\n\tChunkMaxBytes int64\n}\n\n// effectiveChunkSize returns the chunk size to use for uploads.\n// The server-advertised OCI-Chunk-Max-Length always takes precedence: when\n// present, we use (max - margin) to stay safely under the limit regardless\n// of any client-side configuration. The result is also clamped to be at least\n// OCI-Chunk-Min-Length when the server advertises one. The client default\n// (COG_PUSH_DEFAULT_CHUNK_SIZE env var or DefaultChunkSize) is only used when\n// the server does not advertise a maximum.\nfunc (s uploadSession) effectiveChunkSize() int64 {\n\tvar chunkSize = getDefaultChunkSize() // Start with client default as baseline\n\n\tif s.ChunkMaxBytes > 0 {\n\t\t// Server advertised a maximum — use it minus a small margin.\n\t\tchunkSize = s.ChunkMaxBytes - chunkSizeMargin\n\t\tif chunkSize <= 0 {\n\t\t\t// Degenerate case: margin bigger than max. Use max directly.\n\t\t\tchunkSize = s.ChunkMaxBytes\n\t\t}\n\t}\n\n\t// Enforce the server-advertised minimum.\n\tif s.ChunkMinBytes > 0 && chunkSize < s.ChunkMinBytes {\n\t\tchunkSize = s.ChunkMinBytes\n\t}\n\n\treturn chunkSize\n}\n\n// initiateUpload initiates a blob upload and returns an uploadSession containing\n// the upload location URL and server-advertised chunk size limits.\nfunc (c *RegistryClient) initiateUpload(ctx context.Context, client *http.Client, repo name.Repository) (uploadSession, error) {\n\tu := url.URL{\n\t\tScheme: repo.Scheme(),\n\t\tHost:   repo.RegistryStr(),\n\t\tPath:   fmt.Sprintf(\"/v2/%s/blobs/uploads/\", repo.RepositoryStr()),\n\t}\n\n\treq, err := http.NewRequestWithContext(ctx, http.MethodPost, u.String(), nil)\n\tif err != nil {\n\t\treturn uploadSession{}, err\n\t}\n\treq.Header.Set(\"Content-Type\", \"application/json\")\n\n\tresp, err := client.Do(req) //nolint:gosec // G704: URL from registry reference, not user input\n\tif err != nil {\n\t\treturn uploadSession{}, err\n\t}\n\tdefer resp.Body.Close()\n\n\tif err := transport.CheckError(resp, http.StatusAccepted); err != nil {\n\t\treturn uploadSession{}, err\n\t}\n\n\tloc := resp.Header.Get(\"Location\")\n\tif loc == \"\" {\n\t\treturn uploadSession{}, errors.New(\"missing Location header in initiate upload response\")\n\t}\n\n\t// Resolve relative URLs\n\tlocURL, err := url.Parse(loc)\n\tif err != nil {\n\t\treturn uploadSession{}, fmt.Errorf(\"parsing location URL: %w\", err)\n\t}\n\n\tbaseURL := url.URL{\n\t\tScheme: repo.Scheme(),\n\t\tHost:   repo.RegistryStr(),\n\t}\n\n\tsession := uploadSession{\n\t\tLocation: baseURL.ResolveReference(locURL).String(),\n\t}\n\n\t// Parse OCI chunk size headers if the registry advertises them.\n\tif v := resp.Header.Get(\"OCI-Chunk-Min-Length\"); v != \"\" {\n\t\tif n, parseErr := strconv.ParseInt(v, 10, 64); parseErr == nil && n > 0 {\n\t\t\tsession.ChunkMinBytes = n\n\t\t}\n\t}\n\tif v := resp.Header.Get(\"OCI-Chunk-Max-Length\"); v != \"\" {\n\t\tif n, parseErr := strconv.ParseInt(v, 10, 64); parseErr == nil && n > 0 {\n\t\t\tsession.ChunkMaxBytes = n\n\t\t}\n\t}\n\n\treturn session, nil\n}\n\n// uploadBlobChunks uploads a blob using either multipart or single-part upload depending on server support.\n// The repo parameter is needed to restart the upload session if multipart fails.\n// The session carries the upload location and any server-advertised chunk size limits\n// (OCI-Chunk-Min-Length / OCI-Chunk-Max-Length).\n// Returns the final upload location URL which must be used for committing the upload.\nfunc (c *RegistryClient) uploadBlobChunks(ctx context.Context, client *http.Client, repo name.Repository, layer v1.Layer, session uploadSession, totalSize int64, progressCh chan<- v1.Update) (string, error) {\n\t// The chunk size is determined by the server's OCI-Chunk-Max-Length header\n\t// (minus a small margin). When the server does not advertise a maximum,\n\t// the client falls back to COG_PUSH_DEFAULT_CHUNK_SIZE or DefaultChunkSize (96 MiB).\n\t// COG_PUSH_MULTIPART_THRESHOLD controls the minimum blob size for multipart upload (default: 128 MiB).\n\tvar (\n\t\tmultipartThreshold = getMultipartThreshold()\n\t\tchunkSize          = session.effectiveChunkSize()\n\t\tlocation           = session.Location\n\t)\n\n\tif totalSize > multipartThreshold {\n\t\tfinalLocation, newLocation, fallback, err := c.tryMultipartWithFallback(ctx, client, repo, layer, location, totalSize, chunkSize, progressCh)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\tif !fallback {\n\t\t\treturn finalLocation, nil\n\t\t}\n\t\t// Multipart not supported, continue with single-part using the new location\n\t\tlocation = newLocation\n\t}\n\n\t// Single-part upload for small blobs or servers that don't support multipart\n\tblob, err := layer.Compressed()\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"getting compressed blob: %w\", err)\n\t}\n\tdefer blob.Close()\n\n\tfinalLocation, err := c.uploadBlobSingle(ctx, client, location, blob, totalSize, progressCh)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn finalLocation, nil\n}\n\n// tryMultipartWithFallback attempts multipart upload and handles fallback if not supported.\n// Returns (finalLocation, newLocation, fallback, error):\n//   - If multipart succeeds: (finalLocation, \"\", false, nil)\n//   - If multipart not supported: (\"\", newLocation, true, nil) - caller should use single-part with newLocation\n//   - If error: (\"\", \"\", false, error)\nfunc (c *RegistryClient) tryMultipartWithFallback(ctx context.Context, client *http.Client, repo name.Repository, layer v1.Layer, location string, totalSize int64, chunkSize int64, progressCh chan<- v1.Update) (finalLocation string, newLocation string, fallback bool, err error) {\n\tblob, err := layer.Compressed()\n\tif err != nil {\n\t\treturn \"\", \"\", false, fmt.Errorf(\"getting compressed blob: %w\", err)\n\t}\n\tdefer blob.Close()\n\n\tfinalLocation, err = c.tryMultipartUpload(ctx, client, location, blob, totalSize, chunkSize, progressCh)\n\tif err == nil {\n\t\treturn finalLocation, \"\", false, nil\n\t}\n\n\t// Check if error indicates multipart not supported\n\tvar transportErr *transport.Error\n\tif errors.As(err, &transportErr) && (transportErr.StatusCode == http.StatusRequestedRangeNotSatisfiable ||\n\t\ttransportErr.StatusCode == http.StatusBadRequest) {\n\t\t// Multipart not supported - restart upload session for single-part fallback\n\t\tnewSession, err := c.initiateUpload(ctx, client, repo)\n\t\tif err != nil {\n\t\t\treturn \"\", \"\", false, fmt.Errorf(\"restarting upload after multipart failure: %w\", err)\n\t\t}\n\t\treturn \"\", newSession.Location, true, nil\n\t}\n\n\treturn \"\", \"\", false, err\n}\n\n// tryMultipartUpload attempts to upload using Content-Range headers.\n// Returns the final location or an error.\nfunc (c *RegistryClient) tryMultipartUpload(ctx context.Context, client *http.Client, location string, blob io.Reader, totalSize int64, chunkSize int64, progressCh chan<- v1.Update) (string, error) {\n\tvar uploaded int64\n\n\t// Reuse chunk buffers via pool to reduce memory pressure when pushing\n\t// multiple layers concurrently (default concurrency 5 × up to 96 MB each).\n\t// No need to zero the buffer before reuse: io.ReadFull overwrites from\n\t// index 0, and we slice to buffer[:n] so stale bytes are never sent.\n\tvar buffer []byte\n\tif v, ok := chunkBufPool.Get().(*[]byte); ok && int64(len(*v)) == chunkSize {\n\t\tbuffer = *v\n\t} else {\n\t\tbuffer = make([]byte, chunkSize)\n\t}\n\tdefer func() { chunkBufPool.Put(&buffer) }()\n\n\tcurrentLocation := location\n\n\tfor uploaded < totalSize {\n\t\t// Read the next chunk\n\t\tn, err := io.ReadFull(blob, buffer)\n\t\tif err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) {\n\t\t\treturn \"\", fmt.Errorf(\"reading blob: %w\", err)\n\t\t}\n\t\tif n == 0 {\n\t\t\tbreak\n\t\t}\n\n\t\tchunk := buffer[:n]\n\t\tstart := uploaded\n\t\tend := uploaded + int64(n) - 1 // Range is inclusive\n\n\t\t// Upload the chunk with Content-Range (progress is reported within uploadChunk)\n\t\tnewLocation, err := c.uploadChunk(ctx, client, currentLocation, chunk, start, end, totalSize, progressCh)\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\n\t\t// Update location for next chunk (server may change it)\n\t\tif newLocation != \"\" {\n\t\t\tcurrentLocation = newLocation\n\t\t}\n\n\t\tuploaded += int64(n)\n\n\t\t// Check for context cancellation\n\t\tselect {\n\t\tcase <-ctx.Done():\n\t\t\treturn \"\", ctx.Err()\n\t\tdefault:\n\t\t}\n\t}\n\n\treturn currentLocation, nil\n}\n\n// uploadBlobSingle uploads the entire blob in one request without Content-Range headers.\nfunc (c *RegistryClient) uploadBlobSingle(ctx context.Context, client *http.Client, location string, blob io.Reader, totalSize int64, progressCh chan<- v1.Update) (string, error) {\n\t// Wrap the reader to report progress\n\tvar uploaded int64\n\treader := &progressReader{\n\t\treader: blob,\n\t\tonRead: func(n int) {\n\t\t\tuploaded += int64(n)\n\t\t\tif progressCh != nil {\n\t\t\t\t// Cap at totalSize defensively\n\t\t\t\tcomplete := min(uploaded, totalSize)\n\t\t\t\tselect {\n\t\t\t\tcase progressCh <- v1.Update{Complete: complete, Total: totalSize}:\n\t\t\t\tdefault:\n\t\t\t\t\t// Don't block if channel is full\n\t\t\t\t}\n\t\t\t}\n\t\t},\n\t}\n\n\treq, err := http.NewRequestWithContext(ctx, http.MethodPatch, location, reader)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treq.Header.Set(\"Content-Type\", \"application/octet-stream\")\n\treq.ContentLength = totalSize\n\n\tresp, err := client.Do(req) //nolint:gosec // G704: URL from registry upload session, not user input\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tdefer resp.Body.Close()\n\n\tif err := transport.CheckError(resp, http.StatusAccepted, http.StatusNoContent, http.StatusCreated); err != nil {\n\t\treturn \"\", err\n\t}\n\n\t// Return the updated Location header — the registry includes upload state\n\t// that commitUpload needs for the final PUT.\n\tif loc := resp.Header.Get(\"Location\"); loc != \"\" {\n\t\tlocURL, parseErr := url.Parse(loc)\n\t\tif parseErr == nil {\n\t\t\tbaseURL := url.URL{Scheme: \"http\", Host: req.URL.Host}\n\t\t\tif req.URL.Scheme != \"\" {\n\t\t\t\tbaseURL.Scheme = req.URL.Scheme\n\t\t\t}\n\t\t\treturn baseURL.ResolveReference(locURL).String(), nil\n\t\t}\n\t}\n\treturn location, nil\n}\n\n// uploadChunk uploads a single chunk of a blob with Content-Range header.\n// Returns the new location URL if the server returns one.\n// If progressCh is provided, progress updates are sent as bytes are uploaded.\n// Progress updates occur approximately every 32-64KB based on HTTP client buffer size.\nfunc (c *RegistryClient) uploadChunk(ctx context.Context, client *http.Client, location string, chunk []byte, start, end int64, totalSize int64, progressCh chan<- v1.Update) (string, error) {\n\t// Wrap the chunk reader to report progress as bytes are written\n\tvar reader io.Reader\n\tif progressCh != nil {\n\t\tvar chunkUploaded int64\n\t\treader = &progressReader{\n\t\t\treader: bytes.NewReader(chunk),\n\t\t\tonRead: func(n int) {\n\t\t\t\tchunkUploaded += int64(n)\n\t\t\t\t// Cap at totalSize defensively\n\t\t\t\tcomplete := min(start+chunkUploaded, totalSize)\n\t\t\t\tselect {\n\t\t\t\tcase progressCh <- v1.Update{Complete: complete, Total: totalSize}:\n\t\t\t\tdefault:\n\t\t\t\t\t// Don't block if channel is full\n\t\t\t\t}\n\t\t\t},\n\t\t}\n\t} else {\n\t\treader = bytes.NewReader(chunk)\n\t}\n\n\treq, err := http.NewRequestWithContext(ctx, http.MethodPatch, location, reader)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treq.Header.Set(\"Content-Type\", \"application/octet-stream\")\n\treq.Header.Set(\"Content-Length\", strconv.FormatInt(int64(len(chunk)), 10))\n\treq.Header.Set(\"Content-Range\", fmt.Sprintf(\"%d-%d\", start, end))\n\n\tresp, err := client.Do(req) //nolint:gosec // G704: URL from registry upload session, not user input\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tdefer resp.Body.Close()\n\n\tif err := transport.CheckError(resp, http.StatusAccepted, http.StatusNoContent, http.StatusCreated); err != nil {\n\t\treturn \"\", err\n\t}\n\n\t// Get the new location for the next chunk\n\tnewLocation := resp.Header.Get(\"Location\")\n\tif newLocation != \"\" {\n\t\t// Resolve relative URLs\n\t\tlocURL, err := url.Parse(newLocation)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"parsing location URL: %w\", err)\n\t\t}\n\n\t\t// Parse the original location to get the base URL\n\t\torigURL, err := url.Parse(location)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"parsing original location URL: %w\", err)\n\t\t}\n\n\t\treturn origURL.ResolveReference(locURL).String(), nil\n\t}\n\n\treturn \"\", nil\n}\n\n// progressReader wraps an io.Reader to report progress.\ntype progressReader struct {\n\treader io.Reader\n\tonRead func(int)\n}\n\nfunc (pr *progressReader) Read(p []byte) (int, error) {\n\tn, err := pr.reader.Read(p)\n\tif n > 0 {\n\t\tpr.onRead(n)\n\t}\n\treturn n, err\n}\n\n// commitUpload finalizes the upload by sending a PUT request with the digest.\nfunc (c *RegistryClient) commitUpload(ctx context.Context, client *http.Client, location string, digest v1.Hash) error {\n\tu, err := url.Parse(location)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"parsing location URL: %w\", err)\n\t}\n\n\t// Add digest query parameter\n\tq := u.Query()\n\tq.Set(\"digest\", digest.String())\n\tu.RawQuery = q.Encode()\n\n\treq, err := http.NewRequestWithContext(ctx, http.MethodPut, u.String(), nil)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treq.Header.Set(\"Content-Type\", \"application/octet-stream\")\n\n\tresp, err := client.Do(req) //nolint:gosec // G704: URL from registry upload session, not user input\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer resp.Body.Close()\n\n\treturn transport.CheckError(resp, http.StatusCreated)\n}\n\n// pickDefaultImage selects an image from a manifest index to use for fetching labels.\n// Prefers linux/amd64, otherwise returns the first image manifest.\n// Returns an error if no suitable image is found or if fetching fails.\nfunc pickDefaultImage(ref name.Reference, idx *v1.IndexManifest, opts ...remote.Option) (v1.Image, error) {\n\tvar targetDigest string\n\n\t// First, look for linux/amd64\n\tfor _, m := range idx.Manifests {\n\t\tif m.Platform != nil && m.Platform.OS == \"linux\" && m.Platform.Architecture == \"amd64\" {\n\t\t\ttargetDigest = m.Digest.String()\n\t\t\tbreak\n\t\t}\n\t}\n\n\t// Fall back to first manifest\n\tif targetDigest == \"\" && len(idx.Manifests) > 0 {\n\t\ttargetDigest = idx.Manifests[0].Digest.String()\n\t}\n\n\tif targetDigest == \"\" {\n\t\treturn nil, fmt.Errorf(\"index for %s contains no manifests\", ref.String())\n\t}\n\n\tdigestRef, err := name.NewDigest(ref.Context().Name()+\"@\"+targetDigest, name.Insecure)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to create digest reference: %w\", err)\n\t}\n\n\tdesc, err := remote.Get(digestRef, opts...)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to fetch image %s: %w\", digestRef.String(), err)\n\t}\n\n\timg, err := desc.Image()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to load image %s: %w\", digestRef.String(), err)\n\t}\n\n\treturn img, nil\n}\n"
  },
  {
    "path": "pkg/registry/registrytest/mock_client.go",
    "content": "package registrytest\n\nimport (\n\t\"context\"\n\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\n\t\"github.com/replicate/cog/pkg/registry\"\n)\n\ntype MockRegistryClient struct {\n\tmockImages map[string]bool\n}\n\nfunc NewMockRegistryClient() *MockRegistryClient {\n\treturn &MockRegistryClient{\n\t\tmockImages: map[string]bool{},\n\t}\n}\n\nfunc (c *MockRegistryClient) Exists(ctx context.Context, imageRef string) (bool, error) {\n\t_, exists := c.mockImages[imageRef]\n\treturn exists, nil\n}\n\nfunc (c *MockRegistryClient) GetImage(ctx context.Context, imageRef string, platform *registry.Platform) (v1.Image, error) {\n\treturn nil, nil\n}\n\nfunc (c *MockRegistryClient) Inspect(ctx context.Context, imageRef string, platform *registry.Platform) (*registry.ManifestResult, error) {\n\treturn nil, nil\n}\n\nfunc (c *MockRegistryClient) AddMockImage(imageRef string) {\n\tc.mockImages[imageRef] = true\n}\n\nfunc (c *MockRegistryClient) PushImage(ctx context.Context, ref string, img v1.Image) error {\n\tc.mockImages[ref] = true\n\treturn nil\n}\n\nfunc (c *MockRegistryClient) PushIndex(ctx context.Context, ref string, idx v1.ImageIndex) error {\n\tc.mockImages[ref] = true\n\treturn nil\n}\n\nfunc (c *MockRegistryClient) GetDescriptor(ctx context.Context, imageRef string) (v1.Descriptor, error) {\n\treturn v1.Descriptor{}, nil\n}\n\nfunc (c *MockRegistryClient) WriteLayer(ctx context.Context, opts registry.WriteLayerOptions) error {\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/registry_testhelpers/registry_container.go",
    "content": "package registry_testhelpers\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"path\"\n\t\"path/filepath\"\n\t\"runtime\"\n\t\"strconv\"\n\t\"strings\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/go-connections/nat\"\n\t\"github.com/google/go-containerregistry/pkg/authn\"\n\t\"github.com/google/go-containerregistry/pkg/crane\"\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote\"\n\t\"github.com/stretchr/testify/require\"\n\t\"github.com/testcontainers/testcontainers-go\"\n\t\"github.com/testcontainers/testcontainers-go/modules/registry\"\n\t\"github.com/testcontainers/testcontainers-go/wait\"\n\t\"golang.org/x/crypto/bcrypt\"\n\n\tdockerregistry \"github.com/docker/docker/api/types/registry\"\n\n\t\"github.com/replicate/cog/pkg/util\"\n)\n\n// StartTestRegistry starts a test registry container on a random local port populated\n// with image data from the testdata/docker directory. It returns a RegistryContainer\n// that can be used to inspect the registry and generate absolute image references. It will\n// automatically be cleaned when the test finishes.\n// This is safe to run concurrently across multiple tests.\nfunc StartTestRegistry(t *testing.T, opts ...Option) *RegistryContainer {\n\tt.Helper()\n\n\tcontainer, cleanup, err := StartTestRegistryWithCleanup(t.Context(), opts...)\n\trequire.NoError(t, err, \"Failed to start registry container\")\n\n\t// Register cleanup with testing.T\n\tt.Cleanup(cleanup)\n\n\treturn container\n}\n\n// StartTestRegistryWithCleanup starts a test registry and returns a cleanup function.\n// Use this when you don't have a *testing.T (e.g., in testscript harness).\n// The caller is responsible for calling the cleanup function when done.\nfunc StartTestRegistryWithCleanup(ctx context.Context, opts ...Option) (*RegistryContainer, func(), error) {\n\toptions := &options{}\n\tfor _, opt := range opts {\n\t\topt(options)\n\t}\n\n\t_, filename, _, _ := runtime.Caller(0)\n\ttestdataDir := filepath.Join(filepath.Dir(filename), \"testdata\", \"docker\")\n\n\t// Pick a port in the insecure range (Docker considers localhost:1-9999 as insecure)\n\tport, err := util.PickFreePort(1024, 9999)\n\tif err != nil {\n\t\treturn nil, nil, fmt.Errorf(\"pick free port: %w\", err)\n\t}\n\n\tcontainerCustomizers := []testcontainers.ContainerCustomizer{\n\t\ttestcontainers.WithFiles(testcontainers.ContainerFile{\n\t\t\tHostFilePath:      testdataDir,\n\t\t\tContainerFilePath: \"/var/lib/registry/\",\n\t\t\tFileMode:          0o755,\n\t\t}),\n\t\ttestcontainers.WithWaitStrategy(\n\t\t\twait.ForHTTP(\"/\").WithPort(\"5000/tcp\").\n\t\t\t\tWithStartupTimeout(10 * time.Second),\n\t\t),\n\t\ttestcontainers.WithHostConfigModifier(func(hostConfig *container.HostConfig) {\n\t\t\thostConfig.PortBindings = map[nat.Port][]nat.PortBinding{\n\t\t\t\tnat.Port(\"5000/tcp\"): {{HostIP: \"0.0.0.0\", HostPort: strconv.Itoa(port)}},\n\t\t\t}\n\t\t}),\n\t}\n\n\tif options.auth != nil {\n\t\thtpasswd, err := generateHtpasswd(options.auth.Username, options.auth.Password)\n\t\tif err != nil {\n\t\t\treturn nil, nil, fmt.Errorf(\"generate htpasswd: %w\", err)\n\t\t}\n\t\tcontainerCustomizers = append(containerCustomizers,\n\t\t\tregistry.WithHtpasswd(htpasswd),\n\t\t)\n\t}\n\n\tregistryContainer, err := registry.Run(\n\t\tctx,\n\t\t\"registry:3\",\n\t\tcontainerCustomizers...,\n\t)\n\tif err != nil {\n\t\treturn nil, nil, fmt.Errorf(\"start registry container: %w\", err)\n\t}\n\n\tcleanup := func() {\n\t\tif registryContainer != nil {\n\t\t\t_ = registryContainer.Terminate(context.Background())\n\t\t}\n\t}\n\n\treturn &RegistryContainer{\n\t\tContainer: registryContainer,\n\t\toptions:   options,\n\t}, cleanup, nil\n}\n\ntype RegistryContainer struct {\n\tContainer *registry.RegistryContainer\n\toptions   *options\n}\n\nfunc (c *RegistryContainer) ImageRef(ref string) string {\n\treturn path.Join(c.Container.RegistryName, ref)\n}\n\nfunc (c *RegistryContainer) ImageRefForTest(t *testing.T, label string) string {\n\tif label == \"\" {\n\t\tlabel = fmt.Sprintf(\"test-%d\", time.Now().Unix())\n\t}\n\trepo := strings.ToLower(t.Name())\n\treturn c.ImageRef(fmt.Sprintf(\"%s:%s\", repo, label))\n}\n\nfunc (c *RegistryContainer) CloneRepo(t *testing.T, existingRepo, newRepo string) string {\n\texistingRepo = c.ImageRef(existingRepo)\n\tnewRepo = c.ImageRef(newRepo)\n\n\terr := crane.CopyRepository(existingRepo, newRepo)\n\trequire.NoError(t, err, \"Failed to clone repo %q to %q\", existingRepo, newRepo)\n\treturn newRepo\n}\n\nfunc (c *RegistryContainer) CloneRepoForTest(t *testing.T, repo string) string {\n\treturn c.CloneRepo(t, repo, strings.ToLower(t.Name()))\n}\n\nfunc (c *RegistryContainer) ImageExists(t *testing.T, ref string) error {\n\tparsedRef, err := name.ParseReference(ref, name.WithDefaultRegistry(c.RegistryHost()))\n\trequire.NoError(t, err)\n\n\tvar opts []remote.Option\n\n\tif c.options.auth != nil {\n\t\topts = append(opts, remote.WithAuth(authn.FromConfig(authn.AuthConfig{\n\t\t\tUsername: c.options.auth.Username,\n\t\t\tPassword: c.options.auth.Password,\n\t\t})))\n\t}\n\t_, err = remote.Head(parsedRef, opts...)\n\treturn err\n}\n\nfunc (c *RegistryContainer) RegistryHost() string {\n\treturn c.Container.RegistryName\n}\n\ntype Option func(*options)\n\nfunc WithAuth(username, password string) func(*options) {\n\treturn func(o *options) {\n\t\to.auth = &dockerregistry.AuthConfig{\n\t\t\tUsername: username,\n\t\t\tPassword: password,\n\t\t}\n\t}\n}\n\ntype options struct {\n\tauth *dockerregistry.AuthConfig\n}\n\nfunc generateHtpasswd(username, password string) (string, error) {\n\thash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn fmt.Sprintf(\"%s:%s\", username, string(hash)), nil\n}\n"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/blobs/sha256/1c/1c4eef651f65e2f7daee7ee785882ac164b02b78fb74503052a26dc061c90474/data",
    "content": "{\n  \"schemaVersion\": 2,\n  \"mediaType\": \"application/vnd.oci.image.manifest.v1+json\",\n  \"config\": {\n    \"mediaType\": \"application/vnd.oci.image.config.v1+json\",\n    \"digest\": \"sha256:aded1e1a5b3705116fa0a92ba074a5e0b0031647d9c315983ccba2ee5428ec8b\",\n    \"size\": 581\n  },\n  \"layers\": [\n    {\n      \"mediaType\": \"application/vnd.oci.image.layer.v1.tar+gzip\",\n      \"digest\": \"sha256:f18232174bc91741fdf3da96d85011092101a032a93a388b79e99e69c2d5c870\",\n      \"size\": 3642247\n    }\n  ],\n  \"annotations\": {\n    \"com.docker.official-images.bashbrew.arch\": \"amd64\",\n    \"org.opencontainers.image.base.name\": \"scratch\",\n    \"org.opencontainers.image.created\": \"2025-02-14T03:28:36Z\",\n    \"org.opencontainers.image.revision\": \"17fe3d1e2d2cbf54d745139eab749c252e35b883\",\n    \"org.opencontainers.image.source\": \"https://github.com/alpinelinux/docker-alpine.git#17fe3d1e2d2cbf54d745139eab749c252e35b883:x86_64\",\n    \"org.opencontainers.image.url\": \"https://hub.docker.com/_/alpine\",\n    \"org.opencontainers.image.version\": \"3.21.3\"\n  }\n}"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/blobs/sha256/75/757d680068d77be46fd1ea20fb21db16f150468c5e7079a08a2e4705aec096ac/data",
    "content": "{\n  \"schemaVersion\": 2,\n  \"mediaType\": \"application/vnd.oci.image.manifest.v1+json\",\n  \"config\": {\n    \"mediaType\": \"application/vnd.oci.image.config.v1+json\",\n    \"digest\": \"sha256:8d591b0b7dea080ea3be9e12ae563eebf9869168ffced1cb25b2470a3d9fe15e\",\n    \"size\": 597\n  },\n  \"layers\": [\n    {\n      \"mediaType\": \"application/vnd.oci.image.layer.v1.tar+gzip\",\n      \"digest\": \"sha256:6e771e15690e2fabf2332d3a3b744495411d6e0b00b2aea64419b58b0066cf81\",\n      \"size\": 3993029\n    }\n  ],\n  \"annotations\": {\n    \"com.docker.official-images.bashbrew.arch\": \"arm64v8\",\n    \"org.opencontainers.image.base.name\": \"scratch\",\n    \"org.opencontainers.image.created\": \"2025-02-14T03:28:36Z\",\n    \"org.opencontainers.image.revision\": \"17fe3d1e2d2cbf54d745139eab749c252e35b883\",\n    \"org.opencontainers.image.source\": \"https://github.com/alpinelinux/docker-alpine.git#17fe3d1e2d2cbf54d745139eab749c252e35b883:aarch64\",\n    \"org.opencontainers.image.url\": \"https://hub.docker.com/_/alpine\",\n    \"org.opencontainers.image.version\": \"3.21.3\"\n  }\n}"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/blobs/sha256/8d/8d591b0b7dea080ea3be9e12ae563eebf9869168ffced1cb25b2470a3d9fe15e/data",
    "content": "{\"architecture\":\"arm64\",\"config\":{\"Env\":[\"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin\"],\"Cmd\":[\"/bin/sh\"],\"WorkingDir\":\"/\"},\"created\":\"2025-02-14T03:28:36Z\",\"history\":[{\"created\":\"2025-02-14T03:28:36Z\",\"created_by\":\"ADD alpine-minirootfs-3.21.3-aarch64.tar.gz / # buildkit\",\"comment\":\"buildkit.dockerfile.v0\"},{\"created\":\"2025-02-14T03:28:36Z\",\"created_by\":\"CMD [\\\"/bin/sh\\\"]\",\"comment\":\"buildkit.dockerfile.v0\",\"empty_layer\":true}],\"os\":\"linux\",\"rootfs\":{\"type\":\"layers\",\"diff_ids\":[\"sha256:a16e98724c05975ee8c40d8fe389c3481373d34ab20a1cf52ea2accc43f71f4c\"]},\"variant\":\"v8\"}"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/blobs/sha256/9a/9a0ff41dccad7a96f324a4655a715c623ed3511c7336361ffa9dadcecbdb99e5/data",
    "content": "{\"schemaVersion\":2,\"mediaType\":\"application/vnd.oci.image.index.v1+json\",\"manifests\":[{\"mediaType\":\"application/vnd.oci.image.manifest.v1+json\",\"size\":1022,\"digest\":\"sha256:1c4eef651f65e2f7daee7ee785882ac164b02b78fb74503052a26dc061c90474\",\"annotations\":{\"com.docker.official-images.bashbrew.arch\":\"amd64\",\"org.opencontainers.image.base.name\":\"scratch\",\"org.opencontainers.image.created\":\"2025-02-14T18:27:58Z\",\"org.opencontainers.image.revision\":\"17fe3d1e2d2cbf54d745139eab749c252e35b883\",\"org.opencontainers.image.source\":\"https://github.com/alpinelinux/docker-alpine.git#17fe3d1e2d2cbf54d745139eab749c252e35b883:x86_64\",\"org.opencontainers.image.url\":\"https://hub.docker.com/_/alpine\",\"org.opencontainers.image.version\":\"3.21.3\"},\"platform\":{\"architecture\":\"amd64\",\"os\":\"linux\"}},{\"mediaType\":\"application/vnd.oci.image.manifest.v1+json\",\"size\":1025,\"digest\":\"sha256:757d680068d77be46fd1ea20fb21db16f150468c5e7079a08a2e4705aec096ac\",\"annotations\":{\"com.docker.official-images.bashbrew.arch\":\"arm64v8\",\"org.opencontainers.image.base.name\":\"scratch\",\"org.opencontainers.image.created\":\"2025-02-14T18:27:49Z\",\"org.opencontainers.image.revision\":\"17fe3d1e2d2cbf54d745139eab749c252e35b883\",\"org.opencontainers.image.source\":\"https://github.com/alpinelinux/docker-alpine.git#17fe3d1e2d2cbf54d745139eab749c252e35b883:aarch64\",\"org.opencontainers.image.url\":\"https://hub.docker.com/_/alpine\",\"org.opencontainers.image.version\":\"3.21.3\"},\"platform\":{\"architecture\":\"arm64\",\"os\":\"linux\"}}]}"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/blobs/sha256/ad/aded1e1a5b3705116fa0a92ba074a5e0b0031647d9c315983ccba2ee5428ec8b/data",
    "content": "{\"architecture\":\"amd64\",\"config\":{\"Env\":[\"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin\"],\"Cmd\":[\"/bin/sh\"],\"WorkingDir\":\"/\"},\"created\":\"2025-02-14T03:28:36Z\",\"history\":[{\"created\":\"2025-02-14T03:28:36Z\",\"created_by\":\"ADD alpine-minirootfs-3.21.3-x86_64.tar.gz / # buildkit\",\"comment\":\"buildkit.dockerfile.v0\"},{\"created\":\"2025-02-14T03:28:36Z\",\"created_by\":\"CMD [\\\"/bin/sh\\\"]\",\"comment\":\"buildkit.dockerfile.v0\",\"empty_layer\":true}],\"os\":\"linux\",\"rootfs\":{\"type\":\"layers\",\"diff_ids\":[\"sha256:08000c18d16dadf9553d747a58cf44023423a9ab010aab96cf263d2216b8b350\"]}}"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_layers/sha256/6e771e15690e2fabf2332d3a3b744495411d6e0b00b2aea64419b58b0066cf81/link",
    "content": "sha256:6e771e15690e2fabf2332d3a3b744495411d6e0b00b2aea64419b58b0066cf81"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_layers/sha256/8d591b0b7dea080ea3be9e12ae563eebf9869168ffced1cb25b2470a3d9fe15e/link",
    "content": "sha256:8d591b0b7dea080ea3be9e12ae563eebf9869168ffced1cb25b2470a3d9fe15e"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_layers/sha256/aded1e1a5b3705116fa0a92ba074a5e0b0031647d9c315983ccba2ee5428ec8b/link",
    "content": "sha256:aded1e1a5b3705116fa0a92ba074a5e0b0031647d9c315983ccba2ee5428ec8b"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_layers/sha256/f18232174bc91741fdf3da96d85011092101a032a93a388b79e99e69c2d5c870/link",
    "content": "sha256:f18232174bc91741fdf3da96d85011092101a032a93a388b79e99e69c2d5c870"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_manifests/revisions/sha256/1c4eef651f65e2f7daee7ee785882ac164b02b78fb74503052a26dc061c90474/link",
    "content": "sha256:1c4eef651f65e2f7daee7ee785882ac164b02b78fb74503052a26dc061c90474"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_manifests/revisions/sha256/757d680068d77be46fd1ea20fb21db16f150468c5e7079a08a2e4705aec096ac/link",
    "content": "sha256:757d680068d77be46fd1ea20fb21db16f150468c5e7079a08a2e4705aec096ac"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_manifests/revisions/sha256/9a0ff41dccad7a96f324a4655a715c623ed3511c7336361ffa9dadcecbdb99e5/link",
    "content": "sha256:9a0ff41dccad7a96f324a4655a715c623ed3511c7336361ffa9dadcecbdb99e5"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_manifests/tags/latest/current/link",
    "content": "sha256:9a0ff41dccad7a96f324a4655a715c623ed3511c7336361ffa9dadcecbdb99e5"
  },
  {
    "path": "pkg/registry_testhelpers/testdata/docker/registry/v2/repositories/alpine/_manifests/tags/latest/index/sha256/9a0ff41dccad7a96f324a4655a715c623ed3511c7336361ffa9dadcecbdb99e5/link",
    "content": "sha256:9a0ff41dccad7a96f324a4655a715c623ed3511c7336361ffa9dadcecbdb99e5"
  },
  {
    "path": "pkg/requirements/requirements.go",
    "content": "package requirements\n\nimport (\n\t\"bufio\"\n\t\"errors\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"regexp\"\n\t\"strings\"\n\n\t\"github.com/replicate/cog/pkg/util/files\"\n)\n\nconst RequirementsFile = \"requirements.txt\"\nconst OverridesFile = \"overrides.txt\"\n\nfunc GenerateRequirements(tmpDir string, path string, fileName string) (string, error) {\n\tbs, err := os.ReadFile(path)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\trequirements := string(bs)\n\n\t// Check against the old requirements\n\trequirementsFile := filepath.Join(tmpDir, fileName)\n\tif err := files.WriteIfDifferent(requirementsFile, requirements); err != nil {\n\t\treturn \"\", err\n\t}\n\treturn requirementsFile, err\n}\n\nfunc CurrentRequirements(tmpDir string) (string, error) {\n\trequirementsFile := filepath.Join(tmpDir, RequirementsFile)\n\t_, err := os.Stat(requirementsFile)\n\tif err != nil {\n\t\tif errors.Is(err, os.ErrNotExist) {\n\t\t\treturn \"\", nil\n\t\t}\n\t\treturn \"\", err\n\t}\n\treturn requirementsFile, nil\n}\n\nfunc ReadRequirements(path string) ([]string, error) {\n\tre := regexp.MustCompile(`(?m)^\\s*-e\\s+\\.\\s*$`)\n\n\tfh, err := os.Open(path)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdefer fh.Close()\n\n\t// Use scanner to handle CRLF endings\n\tscanner := bufio.NewScanner(fh)\n\tscanner.Split(scanLinesWithContinuations)\n\n\tvar requirements []string\n\n\tfor scanner.Scan() {\n\t\tline := strings.TrimSpace(scanner.Text())\n\n\t\t// Skip empty lines and comment lines\n\t\tif strings.HasPrefix(line, \"#\") {\n\t\t\tcontinue\n\t\t}\n\n\t\tif re.MatchString(line) {\n\t\t\tcontinue\n\t\t}\n\n\t\t// Remove any trailing comments\n\t\tif idx := strings.Index(line, \"#\"); idx >= 0 {\n\t\t\tline = line[:idx]\n\t\t}\n\n\t\tif line != \"\" {\n\t\t\trequirements = append(requirements, line)\n\t\t}\n\t}\n\n\treturn requirements, scanner.Err()\n}\n\n// scanLinesWithContinuations is a modified version of bufio.ScanLines that\n// also handles line continuations (lines ending with a backslash).\nfunc scanLinesWithContinuations(data []byte, atEOF bool) (advance int, token []byte, err error) {\n\t// If we're at EOF and there's no data, return nil\n\tif atEOF && len(data) == 0 {\n\t\treturn 0, nil, nil\n\t}\n\n\tvar line []byte\n\tstart := 0\n\tfor i := range data {\n\t\tif data[i] == '\\n' {\n\t\t\tend := i\n\t\t\tif end > 0 && data[end-1] == '\\r' {\n\t\t\t\tend--\n\t\t\t}\n\t\t\t// Add this segment to our accumulated line\n\t\t\tline = append(line, data[start:end]...)\n\n\t\t\tif len(line) > 0 && line[len(line)-1] == '\\\\' {\n\t\t\t\t// This is a continuation - remove the backslash and continue\n\t\t\t\tline = line[:len(line)-1]\n\t\t\t\tstart = i + 1\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\tif len(line) == 0 {\n\t\t\t\tcontinue\n\t\t\t}\n\n\t\t\t// Not a continuation, return the accumulated line\n\t\t\treturn i + 1, line, nil\n\t\t}\n\t}\n\n\t// If we're at EOF, we have a final, non-terminated line\n\tif atEOF {\n\t\tif len(data) > start {\n\t\t\tline = append(line, data[start:]...)\n\t\t\tif len(line) > 0 && line[len(line)-1] == '\\r' {\n\t\t\t\tline = line[:len(line)-1]\n\t\t\t}\n\t\t}\n\t\treturn len(data), line, nil\n\t}\n\n\t// Need more data\n\treturn 0, nil, nil\n}\n\n// SplitPinnedPythonRequirement returns the name, version, findLinks, and extraIndexURLs from a requirements.txt line\n// in the form name==version [--find-links=] [-f ] [--extra-index-url=]\nfunc SplitPinnedPythonRequirement(requirement string) (name string, version string, findLinks []string, extraIndexURLs []string, err error) {\n\tpinnedPackageRe := regexp.MustCompile(`(?:([a-zA-Z0-9\\-_]+)==([^ ]+)|--find-links=([^\\s]+)|-f\\s+([^\\s]+)|--extra-index-url=([^\\s]+))`)\n\n\tmatches := pinnedPackageRe.FindAllStringSubmatch(requirement, -1)\n\tif matches == nil {\n\t\treturn \"\", \"\", nil, nil, fmt.Errorf(\"Package %s is not in the expected format\", requirement)\n\t}\n\n\tnameFound := false\n\tversionFound := false\n\n\tfor _, match := range matches {\n\t\tif match[1] != \"\" {\n\t\t\tname = match[1]\n\t\t\tnameFound = true\n\t\t}\n\n\t\tif match[2] != \"\" {\n\t\t\tversion = match[2]\n\t\t\tversionFound = true\n\t\t}\n\n\t\tif match[3] != \"\" {\n\t\t\tfindLinks = append(findLinks, match[3])\n\t\t}\n\n\t\tif match[4] != \"\" {\n\t\t\tfindLinks = append(findLinks, match[4])\n\t\t}\n\n\t\tif match[5] != \"\" {\n\t\t\textraIndexURLs = append(extraIndexURLs, match[5])\n\t\t}\n\t}\n\n\tif !nameFound || !versionFound {\n\t\treturn \"\", \"\", nil, nil, fmt.Errorf(\"Package name or version is missing in %s\", requirement)\n\t}\n\n\treturn name, version, findLinks, extraIndexURLs, nil\n}\n\nfunc PackageName(pipRequirement string) string {\n\tre := regexp.MustCompile(`^([a-zA-Z0-9_\\-\\.]+(?:\\[[^\\]]+\\])?)`)\n\tmatch := re.FindStringSubmatch(pipRequirement)\n\tif len(match) > 1 {\n\t\treturn match[1]\n\t}\n\treturn \"\"\n}\n\nfunc VersionSpecifier(pipRequirement string) string {\n\tre := regexp.MustCompile(`^[a-zA-Z0-9_\\-\\.]+(?:\\[[^\\]]+\\])?\\s*([<>=!~]=?\\s*[^;,#\\s]+(?:\\s*,\\s*[<>=!~]=?\\s*[^;,#\\s]+)*(?:\\s*\\|\\|\\s*[<>=!~]=?\\s*[^;,#\\s]+(?:\\s*,\\s*[<>=!~]=?\\s*[^;,#\\s]+)*)*)?`)\n\tmatch := re.FindStringSubmatch(pipRequirement)\n\tif len(match) > 1 {\n\t\t// Optional: strip spaces for uniform output\n\t\treturn strings.ReplaceAll(match[1], \" \", \"\")\n\t}\n\treturn \"\"\n}\n\nfunc Versions(pipRequirement string) []string {\n\tvar versions []string\n\n\t// Match standard specifier versions\n\treVersion := regexp.MustCompile(`[<>=!~]=?\\s*([^\\s,;|]+)`)\n\tmatches := reVersion.FindAllStringSubmatch(pipRequirement, -1)\n\tfor _, match := range matches {\n\t\tif len(match) > 1 {\n\t\t\tversions = append(versions, match[1])\n\t\t}\n\t}\n\n\t// Match @ file/url version\n\treURL := regexp.MustCompile(`@\\s*([^\\s]+)`)\n\tif match := reURL.FindStringSubmatch(pipRequirement); len(match) > 1 {\n\t\tversions = append(versions, match[1])\n\t}\n\n\treturn versions\n}\n"
  },
  {
    "path": "pkg/requirements/requirements_test.go",
    "content": "package requirements\n\nimport (\n\t\"os\"\n\t\"path\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestPythonRequirements(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \"requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(\"torch==2.5.1\"), 0o644)\n\trequire.NoError(t, err)\n\n\ttmpDir := t.TempDir()\n\trequirementsFile, err := GenerateRequirements(tmpDir, reqFile, RequirementsFile)\n\trequire.NoError(t, err)\n\trequire.Equal(t, filepath.Join(tmpDir, \"requirements.txt\"), requirementsFile)\n}\n\nfunc TestReadRequirements(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \"requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(\"torch==2.5.1\"), 0o644)\n\trequire.NoError(t, err)\n\n\trequirements, err := ReadRequirements(reqFile)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"torch==2.5.1\"}, requirements)\n}\n\nfunc TestReadRequirementsLineContinuations(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \"requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(\"torch==\\\\\\n2.5.1\\ntorchvision==\\\\\\r\\n2.5.1\"), 0o644)\n\trequire.NoError(t, err)\n\n\trequirements, err := ReadRequirements(reqFile)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"torch==2.5.1\", \"torchvision==2.5.1\"}, requirements)\n}\n\nfunc TestReadRequirementsStripComments(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \"requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(\"torch==\\\\\\n2.5.1# Heres my comment\\ntorchvision==2.5.1\\n# Heres a beginning of line comment\"), 0o644)\n\trequire.NoError(t, err)\n\n\trequirements, err := ReadRequirements(reqFile)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"torch==2.5.1\", \"torchvision==2.5.1\"}, requirements)\n}\n\nfunc TestReadRequirementsComplex(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \"requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(`foo==1.0.0\n# complex requirements\nfastapi>=0.6,<1\nflask>0.4\n# comments!\n# blank lines!\n\n# arguments\n-f http://example.com`), 0o644)\n\trequire.NoError(t, err)\n\n\trequirements, err := ReadRequirements(reqFile)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"foo==1.0.0\", \"fastapi>=0.6,<1\", \"flask>0.4\", \"-f http://example.com\"}, requirements)\n}\n\nfunc TestReadRequirementsLongLine(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \"requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(`\nantlr4-python3-runtime==4.9.3 \\\n    --hash=sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b\ncolorama==0.4.6 ; sys_platform == 'win32' \\\n    --hash=sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44 \\\n    --hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6\ncontourpy==1.3.2 \\\n    --hash=sha256:0475b1f6604896bc7c53bb070e355e9321e1bc0d381735421a2d2068ec56531f \\\n    --hash=sha256:106fab697af11456fcba3e352ad50effe493a90f893fca6c2ca5c033820cea92 \\\n    --hash=sha256:15ce6ab60957ca74cff444fe66d9045c1fd3e92c8936894ebd1f3eef2fff075f \\\n    --hash=sha256:1c48188778d4d2f3d48e4643fb15d8608b1d01e4b4d6b0548d9b336c28fc9b6f \\\n    --hash=sha256:3859783aefa2b8355697f16642695a5b9792e7a46ab86da1118a4a23a51a33d7 \\\n    --hash=sha256:3d80b2c0300583228ac98d0a927a1ba6a2ba6b8a742463c564f1d419ee5b211e \\\n    --hash=sha256:3f9e896f447c5c8618f1edb2bafa9a4030f22a575ec418ad70611450720b5b08 \\\n    --hash=sha256:434f0adf84911c924519d2b08fc10491dd282b20bdd3fa8f60fd816ea0b48841 \\\n    --hash=sha256:49b65a95d642d4efa8f64ba12558fcb83407e58a2dfba9d796d77b63ccfcaff5 \\\n    --hash=sha256:4caf2bcd2969402bf77edc4cb6034c7dd7c0803213b3523f111eb7460a51b8d2 \\\n    --hash=sha256:532fd26e715560721bb0d5fc7610fce279b3699b018600ab999d1be895b09415 \\\n    --hash=sha256:5ebac872ba09cb8f2131c46b8739a7ff71de28a24c869bcad554477eb089a878 \\\n    --hash=sha256:5f5964cdad279256c084b69c3f412b7801e15356b16efa9d78aa974041903da0 \\\n    --hash=sha256:65a887a6e8c4cd0897507d814b14c54a8c2e2aa4ac9f7686292f9769fcf9a6ab \\\n    --hash=sha256:6a37a2fb93d4df3fc4c0e363ea4d16f83195fc09c891bc8ce072b9d084853445 \\\n    --hash=sha256:70771a461aaeb335df14deb6c97439973d253ae70660ca085eec25241137ef43 \\\n    --hash=sha256:71e2bd4a1c4188f5c2b8d274da78faab884b59df20df63c34f74aa1813c4427c \\\n    --hash=sha256:745b57db7758f3ffc05a10254edd3182a2a83402a89c00957a8e8a22f5582823 \\\n    --hash=sha256:78e9253c3de756b3f6a5174d024c4835acd59eb3f8e2ca13e775dbffe1558f69 \\\n    --hash=sha256:82199cb78276249796419fe36b7386bd8d2cc3f28b3bc19fe2454fe2e26c4c15 \\\n    --hash=sha256:8b7fc0cd78ba2f4695fd0a6ad81a19e7e3ab825c31b577f384aa9d7817dc3bef \\\n    --hash=sha256:8c5acb8dddb0752bf252e01a3035b21443158910ac16a3b0d20e7fed7d534ce5 \\\n    --hash=sha256:8c942a01d9163e2e5cfb05cb66110121b8d07ad438a17f9e766317bcb62abf73 \\\n    --hash=sha256:90df94c89a91b7362e1142cbee7568f86514412ab8a2c0d0fca72d7e91b62912 \\\n    --hash=sha256:970e9173dbd7eba9b4e01aab19215a48ee5dd3f43cef736eebde064a171f89a5 \\\n    --hash=sha256:977e98a0e0480d3fe292246417239d2d45435904afd6d7332d8455981c408b85 \\\n    --hash=sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54 \\\n    --hash=sha256:b7cd50c38f500bbcc9b6a46643a40e0913673f869315d8e70de0438817cb7773 \\\n    --hash=sha256:c49f73e61f1f774650a55d221803b101d966ca0c5a2d6d5e4320ec3997489441 \\\n    --hash=sha256:c66c4906cdbc50e9cba65978823e6e00b45682eb09adbb78c9775b74eb222422 \\\n    --hash=sha256:c6c4639a9c22230276b7bffb6a850dfc8258a2521305e1faefe804d006b2e532 \\\n    --hash=sha256:c85bb486e9be652314bb5b9e2e3b0d1b2e643d5eec4992c0fbe8ac71775da739 \\\n    --hash=sha256:cc829960f34ba36aad4302e78eabf3ef16a3a100863f0d4eeddf30e8a485a03b \\\n    --hash=sha256:d0e589ae0d55204991450bb5c23f571c64fe43adaa53f93fc902a84c96f52fe1 \\\n    --hash=sha256:d14f12932a8d620e307f715857107b1d1845cc44fdb5da2bc8e850f5ceba9f87 \\\n    --hash=sha256:d32530b534e986374fc19eaa77fcb87e8a99e5431499949b828312bdcd20ac52 \\\n    --hash=sha256:d6658ccc7251a4433eebd89ed2672c2ed96fba367fd25ca9512aa92a4b46c4f1 \\\n    --hash=sha256:d91a3ccc7fea94ca0acab82ceb77f396d50a1f67412efe4c526f5d20264e6ecd \\\n    --hash=sha256:de39db2604ae755316cb5967728f4bea92685884b1e767b7c24e983ef5f771cb \\\n    --hash=sha256:de425af81b6cea33101ae95ece1f696af39446db9682a0b56daaa48cfc29f38f \\\n    --hash=sha256:e1578f7eafce927b168752ed7e22646dad6cd9bca673c60bff55889fa236ebf9 \\\n    --hash=sha256:e298e7e70cf4eb179cc1077be1c725b5fd131ebc81181bf0c03525c8abc297fd \\\n    --hash=sha256:eab0f6db315fa4d70f1d8ab514e527f0366ec021ff853d7ed6a2d33605cf4b83 \\\n    --hash=sha256:f26b383144cf2d2c29f01a1e8170f50dacf0eac02d64139dcd709a8ac4eb3cfe\ncycler==0.12.1 \\\n    --hash=sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30 \\\n    --hash=sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c`), 0o644)\n\trequire.NoError(t, err)\n\trequirements, err := ReadRequirements(reqFile)\n\trequire.NoError(t, err)\n\tcheckRequirements(t, []string{\n\t\t\"antlr4-python3-runtime==4.9.3 --hash=sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b\",\n\t\t\"colorama==0.4.6 ; sys_platform == 'win32' --hash=sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44 --hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6\",\n\t\t\"contourpy==1.3.2 --hash=sha256:0475b1f6604896bc7c53bb070e355e9321e1bc0d381735421a2d2068ec56531f --hash=sha256:106fab697af11456fcba3e352ad50effe493a90f893fca6c2ca5c033820cea92 --hash=sha256:15ce6ab60957ca74cff444fe66d9045c1fd3e92c8936894ebd1f3eef2fff075f --hash=sha256:1c48188778d4d2f3d48e4643fb15d8608b1d01e4b4d6b0548d9b336c28fc9b6f --hash=sha256:3859783aefa2b8355697f16642695a5b9792e7a46ab86da1118a4a23a51a33d7 --hash=sha256:3d80b2c0300583228ac98d0a927a1ba6a2ba6b8a742463c564f1d419ee5b211e --hash=sha256:3f9e896f447c5c8618f1edb2bafa9a4030f22a575ec418ad70611450720b5b08 --hash=sha256:434f0adf84911c924519d2b08fc10491dd282b20bdd3fa8f60fd816ea0b48841 --hash=sha256:49b65a95d642d4efa8f64ba12558fcb83407e58a2dfba9d796d77b63ccfcaff5 --hash=sha256:4caf2bcd2969402bf77edc4cb6034c7dd7c0803213b3523f111eb7460a51b8d2 --hash=sha256:532fd26e715560721bb0d5fc7610fce279b3699b018600ab999d1be895b09415 --hash=sha256:5ebac872ba09cb8f2131c46b8739a7ff71de28a24c869bcad554477eb089a878 --hash=sha256:5f5964cdad279256c084b69c3f412b7801e15356b16efa9d78aa974041903da0 --hash=sha256:65a887a6e8c4cd0897507d814b14c54a8c2e2aa4ac9f7686292f9769fcf9a6ab --hash=sha256:6a37a2fb93d4df3fc4c0e363ea4d16f83195fc09c891bc8ce072b9d084853445 --hash=sha256:70771a461aaeb335df14deb6c97439973d253ae70660ca085eec25241137ef43 --hash=sha256:71e2bd4a1c4188f5c2b8d274da78faab884b59df20df63c34f74aa1813c4427c --hash=sha256:745b57db7758f3ffc05a10254edd3182a2a83402a89c00957a8e8a22f5582823 --hash=sha256:78e9253c3de756b3f6a5174d024c4835acd59eb3f8e2ca13e775dbffe1558f69 --hash=sha256:82199cb78276249796419fe36b7386bd8d2cc3f28b3bc19fe2454fe2e26c4c15 --hash=sha256:8b7fc0cd78ba2f4695fd0a6ad81a19e7e3ab825c31b577f384aa9d7817dc3bef --hash=sha256:8c5acb8dddb0752bf252e01a3035b21443158910ac16a3b0d20e7fed7d534ce5 --hash=sha256:8c942a01d9163e2e5cfb05cb66110121b8d07ad438a17f9e766317bcb62abf73 --hash=sha256:90df94c89a91b7362e1142cbee7568f86514412ab8a2c0d0fca72d7e91b62912 --hash=sha256:970e9173dbd7eba9b4e01aab19215a48ee5dd3f43cef736eebde064a171f89a5 --hash=sha256:977e98a0e0480d3fe292246417239d2d45435904afd6d7332d8455981c408b85 --hash=sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54 --hash=sha256:b7cd50c38f500bbcc9b6a46643a40e0913673f869315d8e70de0438817cb7773 --hash=sha256:c49f73e61f1f774650a55d221803b101d966ca0c5a2d6d5e4320ec3997489441 --hash=sha256:c66c4906cdbc50e9cba65978823e6e00b45682eb09adbb78c9775b74eb222422 --hash=sha256:c6c4639a9c22230276b7bffb6a850dfc8258a2521305e1faefe804d006b2e532 --hash=sha256:c85bb486e9be652314bb5b9e2e3b0d1b2e643d5eec4992c0fbe8ac71775da739 --hash=sha256:cc829960f34ba36aad4302e78eabf3ef16a3a100863f0d4eeddf30e8a485a03b --hash=sha256:d0e589ae0d55204991450bb5c23f571c64fe43adaa53f93fc902a84c96f52fe1 --hash=sha256:d14f12932a8d620e307f715857107b1d1845cc44fdb5da2bc8e850f5ceba9f87 --hash=sha256:d32530b534e986374fc19eaa77fcb87e8a99e5431499949b828312bdcd20ac52 --hash=sha256:d6658ccc7251a4433eebd89ed2672c2ed96fba367fd25ca9512aa92a4b46c4f1 --hash=sha256:d91a3ccc7fea94ca0acab82ceb77f396d50a1f67412efe4c526f5d20264e6ecd --hash=sha256:de39db2604ae755316cb5967728f4bea92685884b1e767b7c24e983ef5f771cb --hash=sha256:de425af81b6cea33101ae95ece1f696af39446db9682a0b56daaa48cfc29f38f --hash=sha256:e1578f7eafce927b168752ed7e22646dad6cd9bca673c60bff55889fa236ebf9 --hash=sha256:e298e7e70cf4eb179cc1077be1c725b5fd131ebc81181bf0c03525c8abc297fd --hash=sha256:eab0f6db315fa4d70f1d8ab514e527f0366ec021ff853d7ed6a2d33605cf4b83 --hash=sha256:f26b383144cf2d2c29f01a1e8170f50dacf0eac02d64139dcd709a8ac4eb3cfe\",\n\t\t\"cycler==0.12.1 --hash=sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30 --hash=sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c\",\n\t}, requirements)\n}\n\nfunc TestComfyUIRequirements(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \"requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(`torch\ntorchvision\ntorchaudio\ntorchsde\neinops\ntransformers>=4.49.0\ntokenizers>=0.13.3\nsentencepiece\nsafetensors>=0.3.0\naiohttp\naccelerate>=1.1.1\npyyaml\nPillow\nscipy\ntqdm\npsutil\nspandrel\nsoundfile\nkornia>=0.7.1\nwebsocket-client==1.6.3\ndiffusers>=0.31.0\nav>=14.1.0\ncomfyui-frontend-package==1.17.11\ncomfyui-workflow-templates==0.1.3\n\n# ComfyUI-AdvancedLivePortrait\ndill\n\n# Inspire\nwebcolors\n\nalbumentations==1.4.3\n\n# was-node-suite-comfyui\n# https://github.com/WASasquatch/was-node-suite-comfyui/blob/main/requirements.txt\ncmake\nimageio\njoblib\nmatplotlib\npilgram\nscikit-learn\nrembg\n\n# ComfyUI_essentials\nnumba\n\n# ComfyUI_FizzNodes\npandas\nnumexpr\n\n# comfyui-reactor-node\ninsightface\nonnx\n\n# ComfyUI-Impact-Pack\nsegment-anything\npiexif\n\n# ComfyUI-Impact-Subpack\nultralytics!=8.0.177\n\n# comfyui_segment_anything\ntimm\n\n# comfyui_controlnet_aux\n# https://github.com/Fannovel16/comfyui_controlnet_aux/blob/main/requirements.txt\nimportlib_metadata\nopencv-python-headless>=4.0.1.24\nfilelock\nnumpy\nscikit-image\npython-dateutil\nmediapipe\nsvglib\nfvcore\nyapf\nomegaconf\nftfy\naddict\nyacs\ntrimesh[easy]\n\n# ComfyUI-KJNodes\nlibrosa\ncolor-matcher\n\n# PuLID\nfacexlib\n\n# SUPIR\nopen-clip-torch>=2.24.0\npytorch-lightning>=2.2.1\n\n# For train.py and custom loras\nhuggingface_hub[hf-transfer]\n\n# ComfyUI-segment-anything-2\niopath`), 0o644)\n\trequire.NoError(t, err)\n\n\trequirements, err := ReadRequirements(reqFile)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\n\t\t\"torch\",\n\t\t\"torchvision\",\n\t\t\"torchaudio\",\n\t\t\"torchsde\",\n\t\t\"einops\",\n\t\t\"transformers>=4.49.0\",\n\t\t\"tokenizers>=0.13.3\",\n\t\t\"sentencepiece\",\n\t\t\"safetensors>=0.3.0\",\n\t\t\"aiohttp\",\n\t\t\"accelerate>=1.1.1\",\n\t\t\"pyyaml\",\n\t\t\"Pillow\",\n\t\t\"scipy\",\n\t\t\"tqdm\",\n\t\t\"psutil\",\n\t\t\"spandrel\",\n\t\t\"soundfile\",\n\t\t\"kornia>=0.7.1\",\n\t\t\"websocket-client==1.6.3\",\n\t\t\"diffusers>=0.31.0\",\n\t\t\"av>=14.1.0\",\n\t\t\"comfyui-frontend-package==1.17.11\",\n\t\t\"comfyui-workflow-templates==0.1.3\",\n\t\t\"dill\",\n\t\t\"webcolors\",\n\t\t\"albumentations==1.4.3\",\n\t\t\"cmake\",\n\t\t\"imageio\",\n\t\t\"joblib\",\n\t\t\"matplotlib\",\n\t\t\"pilgram\",\n\t\t\"scikit-learn\",\n\t\t\"rembg\",\n\t\t\"numba\",\n\t\t\"pandas\",\n\t\t\"numexpr\",\n\t\t\"insightface\",\n\t\t\"onnx\",\n\t\t\"segment-anything\",\n\t\t\"piexif\",\n\t\t\"ultralytics!=8.0.177\",\n\t\t\"timm\",\n\t\t\"importlib_metadata\",\n\t\t\"opencv-python-headless>=4.0.1.24\",\n\t\t\"filelock\",\n\t\t\"numpy\",\n\t\t\"scikit-image\",\n\t\t\"python-dateutil\",\n\t\t\"mediapipe\",\n\t\t\"svglib\",\n\t\t\"fvcore\",\n\t\t\"yapf\",\n\t\t\"omegaconf\",\n\t\t\"ftfy\",\n\t\t\"addict\",\n\t\t\"yacs\",\n\t\t\"trimesh[easy]\",\n\t\t\"librosa\",\n\t\t\"color-matcher\",\n\t\t\"facexlib\",\n\t\t\"open-clip-torch>=2.24.0\",\n\t\t\"pytorch-lightning>=2.2.1\",\n\t\t\"huggingface_hub[hf-transfer]\",\n\t\t\"iopath\",\n\t}, requirements)\n}\n\nfunc TestTensorflowRequirements(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \".requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(`compel==2.0.3\ndiffusers>=0.27.1\ngputil==1.4.0\nloguru==0.7.2\nopencv-python>=4.9.0.80\npillow>=10.2.0\npsutil==6.1.1\nreplicate>=1.0.4\nsentry-sdk[fastapi,loguru]>=2.16.0\nantialiased_cnns==0.3\nbeautifulsoup4==4.13.4\nimageio==2.37.0\nipdb==0.13.13\nkornia==0.8.1\nmatplotlib==3.10.3\nnumpy==1.23.5\nopencv_python==4.11.0.86\nPillow==11.2.1\npytorch_lightning==2.3.3\nPyYAML==6.0.2\nRequests==2.32.3\nscipy==1.15.3\nscikit-image==0.24.0\ntensorflow==2.10.0\ntensorlayer==2.2.5\ntf_slim==1.1.0\ntimm==1.0.15\ntorch==2.0.1\ntorchvision==0.15.2\ntqdm==4.67.1`), 0o644)\n\trequire.NoError(t, err)\n\trequirements, err := ReadRequirements(reqFile)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\n\t\t\"compel==2.0.3\",\n\t\t\"diffusers>=0.27.1\",\n\t\t\"gputil==1.4.0\",\n\t\t\"loguru==0.7.2\",\n\t\t\"opencv-python>=4.9.0.80\",\n\t\t\"pillow>=10.2.0\",\n\t\t\"psutil==6.1.1\",\n\t\t\"replicate>=1.0.4\",\n\t\t\"sentry-sdk[fastapi,loguru]>=2.16.0\",\n\t\t\"antialiased_cnns==0.3\",\n\t\t\"beautifulsoup4==4.13.4\",\n\t\t\"imageio==2.37.0\",\n\t\t\"ipdb==0.13.13\",\n\t\t\"kornia==0.8.1\",\n\t\t\"matplotlib==3.10.3\",\n\t\t\"numpy==1.23.5\",\n\t\t\"opencv_python==4.11.0.86\",\n\t\t\"Pillow==11.2.1\",\n\t\t\"pytorch_lightning==2.3.3\",\n\t\t\"PyYAML==6.0.2\",\n\t\t\"Requests==2.32.3\",\n\t\t\"scipy==1.15.3\",\n\t\t\"scikit-image==0.24.0\",\n\t\t\"tensorflow==2.10.0\",\n\t\t\"tensorlayer==2.2.5\",\n\t\t\"tf_slim==1.1.0\",\n\t\t\"timm==1.0.15\",\n\t\t\"torch==2.0.1\",\n\t\t\"torchvision==0.15.2\",\n\t\t\"tqdm==4.67.1\",\n\t}, requirements)\n}\n\nfunc TestSplitPinnedPythonRequirement(t *testing.T) {\n\ttestCases := []struct {\n\t\tinput                  string\n\t\texpectedName           string\n\t\texpectedVersion        string\n\t\texpectedFindLinks      []string\n\t\texpectedExtraIndexURLs []string\n\t\texpectedError          bool\n\t}{\n\t\t{\"package1==1.0.0\", \"package1\", \"1.0.0\", nil, nil, false},\n\t\t{\"package1==1.0.0+alpha\", \"package1\", \"1.0.0+alpha\", nil, nil, false},\n\t\t{\"--find-links=link1 --find-links=link2 package3==3.0.0\", \"package3\", \"3.0.0\", []string{\"link1\", \"link2\"}, nil, false},\n\t\t{\"package4==4.0.0 --extra-index-url=url1 --extra-index-url=url2\", \"package4\", \"4.0.0\", nil, []string{\"url1\", \"url2\"}, false},\n\t\t{\"-f link1 --find-links=link2 package5==5.0.0 --extra-index-url=url1 --extra-index-url=url2\", \"package5\", \"5.0.0\", []string{\"link1\", \"link2\"}, []string{\"url1\", \"url2\"}, false},\n\t\t{\"package6 --find-links=link1 --find-links=link2 --extra-index-url=url1 --extra-index-url=url2\", \"\", \"\", nil, nil, true},\n\t\t{\"invalid package\", \"\", \"\", nil, nil, true},\n\t\t{\"package8==\", \"\", \"\", nil, nil, true},\n\t\t{\"==8.0.0\", \"\", \"\", nil, nil, true},\n\t}\n\n\tfor _, tc := range testCases {\n\t\tname, version, findLinks, extraIndexURLs, err := SplitPinnedPythonRequirement(tc.input)\n\n\t\tif tc.expectedError {\n\t\t\trequire.Error(t, err)\n\t\t} else {\n\t\t\trequire.NoError(t, err)\n\t\t\trequire.Equal(t, tc.expectedName, name, \"input: \"+tc.input)\n\t\t\trequire.Equal(t, tc.expectedVersion, version, \"input: \"+tc.input)\n\t\t\trequire.Equal(t, tc.expectedFindLinks, findLinks, \"input: \"+tc.input)\n\t\t\trequire.Equal(t, tc.expectedExtraIndexURLs, extraIndexURLs, \"input: \"+tc.input)\n\t\t}\n\t}\n}\n\nfunc TestReadRequirementsWithEditable(t *testing.T) {\n\tsrcDir := t.TempDir()\n\treqFile := path.Join(srcDir, \"requirements.txt\")\n\terr := os.WriteFile(reqFile, []byte(\"-e .\\ntorch==2.5.1\"), 0o644)\n\trequire.NoError(t, err)\n\n\trequirements, err := ReadRequirements(reqFile)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"torch==2.5.1\"}, requirements)\n}\n\nfunc TestVersionSpecifier(t *testing.T) {\n\tspecifier := VersionSpecifier(\"mypackage>= 1.0, < 1.4 || > 2.0\")\n\trequire.Equal(t, specifier, \">=1.0,<1.4||>2.0\")\n}\n\nfunc TestPackageName(t *testing.T) {\n\tname := PackageName(\"mypackage>= 1.0, < 1.4 || > 2.0\")\n\trequire.Equal(t, name, \"mypackage\")\n}\n\nfunc TestVersions(t *testing.T) {\n\tversions := Versions(\"another @ https://some.domain/package.whl\")\n\trequire.Equal(t, versions, []string{\"https://some.domain/package.whl\"})\n}\n\nfunc checkRequirements(t *testing.T, expected []string, actual []string) {\n\tt.Helper()\n\tfor n, expectLine := range expected {\n\t\tactualLine := actual[n]\n\t\t// collapse any multiple-space runs with single spaces in the actual line - the generator may output these\n\t\t// but we don't care about them for comparison purposes\n\t\tactualLine = strings.Join(strings.Fields(actualLine), \" \")\n\t\trequire.Equal(t, expectLine, actualLine)\n\t}\n\trequire.Equal(t, len(expected), len(actual))\n}\n"
  },
  {
    "path": "pkg/schema/errors.go",
    "content": "package schema\n\nimport \"fmt\"\n\n// SchemaError represents errors during schema generation.\ntype SchemaError struct {\n\tKind    SchemaErrorKind\n\tMessage string\n}\n\nfunc (e *SchemaError) Error() string { return e.Message }\n\n// SchemaErrorKind classifies schema generation errors.\ntype SchemaErrorKind int\n\nconst (\n\tErrParse SchemaErrorKind = iota\n\tErrPredictorNotFound\n\tErrMethodNotFound\n\tErrMissingReturnType\n\tErrMissingTypeAnnotation\n\tErrUnsupportedType\n\tErrDefaultFactoryNotSupported\n\tErrInvalidConstraint\n\tErrInvalidPredictRef\n\tErrOptionalOutput\n\tErrConcatIteratorNotStr\n\tErrChoicesNotResolvable\n\tErrDefaultNotResolvable\n\tErrUnresolvableType\n\tErrOther\n)\n\n// NewError creates a SchemaError with the given kind and message.\nfunc NewError(kind SchemaErrorKind, msg string) *SchemaError {\n\treturn &SchemaError{Kind: kind, Message: msg}\n}\n\n// WrapError creates a SchemaError, appending the inner error's message if non-nil.\nfunc WrapError(kind SchemaErrorKind, msg string, inner error) *SchemaError {\n\tif inner != nil {\n\t\treturn &SchemaError{Kind: kind, Message: fmt.Sprintf(\"%s: %s\", msg, inner.Error())}\n\t}\n\treturn &SchemaError{Kind: kind, Message: msg}\n}\n\nfunc errParse(msg string) error { //nolint:unused // used by generator.go (not yet written)\n\treturn &SchemaError{Kind: ErrParse, Message: fmt.Sprintf(\"failed to parse Python source: %s\", msg)}\n}\n\nfunc errPredictorNotFound(name string) error { //nolint:unused // used by generator.go (not yet written)\n\treturn &SchemaError{Kind: ErrPredictorNotFound, Message: fmt.Sprintf(\"predictor not found: %s\", name)}\n}\n\nfunc errMethodNotFound(class, method string) error { //nolint:unused // used by generator.go (not yet written)\n\treturn &SchemaError{Kind: ErrMethodNotFound, Message: fmt.Sprintf(\"%s method not found on %s\", method, class)}\n}\n\nfunc errMissingReturnType(method string) error { //nolint:unused // used by generator.go (not yet written)\n\treturn &SchemaError{Kind: ErrMissingReturnType, Message: fmt.Sprintf(\"missing return type annotation on %s\", method)}\n}\n\nfunc errMissingTypeAnnotation(method, param string) error { //nolint:unused // used by generator.go (not yet written)\n\treturn &SchemaError{Kind: ErrMissingTypeAnnotation, Message: fmt.Sprintf(\"missing type annotation for parameter '%s' on %s\", param, method)}\n}\n\nfunc errUnsupportedType(msg string) error {\n\treturn &SchemaError{Kind: ErrUnsupportedType, Message: fmt.Sprintf(\"unsupported type: %s\", msg)}\n}\n\nfunc errDefaultFactoryNotSupported(param string) error { //nolint:unused // used by generator.go (not yet written)\n\treturn &SchemaError{\n\t\tKind:    ErrDefaultFactoryNotSupported,\n\t\tMessage: fmt.Sprintf(\"default_factory is not supported in Input() — use a literal default value instead (parameter '%s')\", param),\n\t}\n}\n\nfunc errInvalidPredictRef(ref string) error {\n\treturn &SchemaError{\n\t\tKind:    ErrInvalidPredictRef,\n\t\tMessage: fmt.Sprintf(\"invalid predict reference '%s' — expected format: file.py:ClassName or file.py:function_name\", ref),\n\t}\n}\n\nfunc errOptionalOutput() error {\n\treturn &SchemaError{Kind: ErrOptionalOutput, Message: \"unsupported output type: Optional is not allowed as a return type\"}\n}\n\nfunc errConcatIteratorNotStr(got string) error {\n\treturn &SchemaError{Kind: ErrConcatIteratorNotStr, Message: fmt.Sprintf(\"ConcatenateIterator element type must be str, got %s\", got)}\n}\n\nfunc errChoicesNotResolvable(param string) error { //nolint:unused // used by generator.go (not yet written)\n\treturn &SchemaError{\n\t\tKind:    ErrChoicesNotResolvable,\n\t\tMessage: fmt.Sprintf(\"choices for parameter '%s' cannot be statically resolved — use a literal list instead (e.g. choices=[\\\"a\\\", \\\"b\\\"])\", param),\n\t}\n}\n\nfunc errUnresolvableImportedType(name, module string) error {\n\treturn &SchemaError{\n\t\tKind: ErrUnresolvableType,\n\t\tMessage: fmt.Sprintf(\n\t\t\t\"cannot resolve output type '%s' (imported from '%s') — \"+\n\t\t\t\t\"external types cannot be statically analyzed. \"+\n\t\t\t\t\"Define it as a BaseModel subclass in your predict file, or provide a .pyi stub\",\n\t\t\tname, module),\n\t}\n}\n\nfunc errUnresolvableType(name string) error {\n\treturn &SchemaError{\n\t\tKind: ErrUnresolvableType,\n\t\tMessage: fmt.Sprintf(\n\t\t\t\"cannot resolve output type '%s' — \"+\n\t\t\t\t\"it is not a primitive type (str, int, float, bool, Path) \"+\n\t\t\t\t\"and no BaseModel definition was found in the predict file\",\n\t\t\tname),\n\t}\n}\n\nfunc errDefaultNotResolvable(param, expr string) error { //nolint:unused // used by generator.go (not yet written)\n\treturn &SchemaError{\n\t\tKind: ErrDefaultNotResolvable,\n\t\tMessage: fmt.Sprintf(\n\t\t\t\"default value for parameter '%s' cannot be statically resolved: `%s`. \"+\n\t\t\t\t\"Defaults must be literals (string, int, float, bool, None, list) or Input() calls.\", param, expr),\n\t}\n}\n"
  },
  {
    "path": "pkg/schema/generator.go",
    "content": "package schema\n\nimport (\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"os\"\n\t\"strings\"\n)\n\n// Parser is a function that parses source code and extracts predictor info.\n// This is defined as a type to avoid an import cycle between schema and\n// schema/python. The concrete implementation is python.ParsePredictor.\n//\n// sourceDir is the project root directory, used for resolving cross-file\n// imports (e.g. \"from .types import Output\"). Pass \"\" if unknown.\ntype Parser func(source []byte, predictRef string, mode Mode, sourceDir string) (*PredictorInfo, error)\n\n// Generate produces an OpenAPI 3.0.2 JSON schema from a predict/train reference.\n//\n// predictRef has the format \"module.py:ClassName\" (e.g. \"predict.py:Predictor\").\n// sourceDir is the directory containing the source file.\n// mode selects predict vs train.\n// parse is the parser implementation (use python.ParsePredictor).\n//\n// If the COG_OPENAPI_SCHEMA environment variable is set, its value is treated\n// as a path to a pre-built JSON schema file. The file contents are returned\n// directly and no parsing or generation takes place.\nfunc Generate(predictRef string, sourceDir string, mode Mode, parse Parser) ([]byte, error) {\n\t// \"Bring your own schema\" override\n\tif schemaPath := os.Getenv(\"COG_OPENAPI_SCHEMA\"); schemaPath != \"\" {\n\t\tdata, err := os.ReadFile(schemaPath) //nolint:gosec // G703: path from trusted env var\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"COG_OPENAPI_SCHEMA: failed to read %s: %w\", schemaPath, err)\n\t\t}\n\t\treturn data, nil\n\t}\n\n\tfilePath, className, err := parsePredictRef(predictRef)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tfullPath := filePath\n\tif sourceDir != \"\" {\n\t\tfullPath = sourceDir + \"/\" + filePath\n\t}\n\n\tsource, err := os.ReadFile(fullPath)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to read predictor source %s: %w\", fullPath, err)\n\t}\n\n\treturn GenerateFromSource(source, className, mode, parse, sourceDir)\n}\n\n// GenerateFromSource produces an OpenAPI 3.0.2 JSON schema from Python source bytes.\n//\n// predictRef is the class or function name (e.g. \"Predictor\" or \"predict\").\n// parse is the parser implementation (use python.ParsePredictor).\n// sourceDir is the project root for resolving cross-file imports. Pass \"\" if unknown.\n// This is the lower-level API — it does not read files or check COG_OPENAPI_SCHEMA.\nfunc GenerateFromSource(source []byte, predictRef string, mode Mode, parse Parser, sourceDir string) ([]byte, error) {\n\tinfo, err := parse(source, predictRef, mode, sourceDir)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treturn GenerateOpenAPISchema(info)\n}\n\n// GenerateCombined produces an OpenAPI schema for both predict and train (when\n// both are configured) and merges them into a single document. If only one mode\n// is configured, it returns that single schema.\n//\n// If the COG_OPENAPI_SCHEMA environment variable is set, its value is treated\n// as a path to a pre-built JSON schema file and returned directly.\nfunc GenerateCombined(sourceDir string, predictRef string, trainRef string, parse Parser) ([]byte, error) {\n\t// \"Bring your own schema\" override\n\tif schemaPath := os.Getenv(\"COG_OPENAPI_SCHEMA\"); schemaPath != \"\" {\n\t\tdata, err := os.ReadFile(schemaPath) //nolint:gosec // G703: path from trusted env var\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"COG_OPENAPI_SCHEMA: failed to read %s: %w\", schemaPath, err)\n\t\t}\n\t\treturn data, nil\n\t}\n\n\tif predictRef == \"\" && trainRef == \"\" {\n\t\treturn nil, fmt.Errorf(\"no predict or train reference provided\")\n\t}\n\n\t// Single-mode: just generate the one schema\n\tif predictRef == \"\" {\n\t\treturn Generate(trainRef, sourceDir, ModeTrain, parse)\n\t}\n\tif trainRef == \"\" {\n\t\treturn Generate(predictRef, sourceDir, ModePredict, parse)\n\t}\n\n\t// Both modes: generate each and merge\n\tpredictJSON, err := Generate(predictRef, sourceDir, ModePredict, parse)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"predict schema: %w\", err)\n\t}\n\ttrainJSON, err := Generate(trainRef, sourceDir, ModeTrain, parse)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"train schema: %w\", err)\n\t}\n\n\tvar predictSchema, trainSchema map[string]any\n\tif err := json.Unmarshal(predictJSON, &predictSchema); err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to parse predict schema: %w\", err)\n\t}\n\tif err := json.Unmarshal(trainJSON, &trainSchema); err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to parse train schema: %w\", err)\n\t}\n\n\tmerged := MergeSchemas(predictSchema, trainSchema)\n\treturn json.MarshalIndent(merged, \"\", \"  \")\n}\n\n// MergeSchemas merges a predict-mode and train-mode OpenAPI schema into a single\n// combined schema. The predict schema is used as the base; paths and component\n// schemas from the train schema are added to it.\nfunc MergeSchemas(predict, train map[string]any) map[string]any {\n\t// Merge paths\n\tpredictPaths, _ := predict[\"paths\"].(map[string]any)\n\ttrainPaths, _ := train[\"paths\"].(map[string]any)\n\tif predictPaths != nil && trainPaths != nil {\n\t\tfor k, v := range trainPaths {\n\t\t\tif _, exists := predictPaths[k]; !exists {\n\t\t\t\tpredictPaths[k] = v\n\t\t\t}\n\t\t}\n\t}\n\n\t// Merge component schemas\n\tpredictComponents, _ := predict[\"components\"].(map[string]any)\n\ttrainComponents, _ := train[\"components\"].(map[string]any)\n\tif predictComponents != nil && trainComponents != nil {\n\t\tpredictSchemas, _ := predictComponents[\"schemas\"].(map[string]any)\n\t\ttrainSchemas, _ := trainComponents[\"schemas\"].(map[string]any)\n\t\tif predictSchemas != nil && trainSchemas != nil {\n\t\t\tfor k, v := range trainSchemas {\n\t\t\t\tif _, exists := predictSchemas[k]; !exists {\n\t\t\t\t\tpredictSchemas[k] = v\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\treturn predict\n}\n\n// parsePredictRef splits a predict reference like \"predict.py:Predictor\" into\n// the file path and class/function name.\nfunc parsePredictRef(ref string) (filePath string, name string, err error) {\n\tparts := strings.SplitN(ref, \":\", 2)\n\tif len(parts) != 2 || parts[0] == \"\" || parts[1] == \"\" {\n\t\treturn \"\", \"\", errInvalidPredictRef(ref)\n\t}\n\treturn parts[0], parts[1], nil\n}\n"
  },
  {
    "path": "pkg/schema/generator_test.go",
    "content": "package schema\n\nimport (\n\t\"encoding/json\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n)\n\n// mockParser is a test parser that returns a fixed PredictorInfo.\nfunc mockParser(source []byte, predictRef string, mode Mode, sourceDir string) (*PredictorInfo, error) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"prompt\", InputField{\n\t\tName:      \"prompt\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeString, Repetition: Required},\n\t})\n\treturn &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   mode,\n\t}, nil\n}\n\n// failParser always returns an error.\nfunc failParser(_ []byte, _ string, _ Mode, _ string) (*PredictorInfo, error) {\n\treturn nil, NewError(ErrParse, \"mock parse failure\")\n}\n\n// ---------------------------------------------------------------------------\n// parsePredictRef\n// ---------------------------------------------------------------------------\n\nfunc TestParsePredictRef(t *testing.T) {\n\ttests := []struct {\n\t\tinput   string\n\t\tfile    string\n\t\tname    string\n\t\twantErr bool\n\t}{\n\t\t{\"predict.py:Predictor\", \"predict.py\", \"Predictor\", false},\n\t\t{\"src/model.py:MyModel\", \"src/model.py\", \"MyModel\", false},\n\t\t{\"train.py:train\", \"train.py\", \"train\", false},\n\t\t{\"no_colon\", \"\", \"\", true},\n\t\t{\":NoFile\", \"\", \"\", true},\n\t\t{\"no_name:\", \"\", \"\", true},\n\t\t{\"\", \"\", \"\", true},\n\t}\n\tfor _, tt := range tests {\n\t\tt.Run(tt.input, func(t *testing.T) {\n\t\t\tfile, name, err := parsePredictRef(tt.input)\n\t\t\tif tt.wantErr {\n\t\t\t\trequire.Error(t, err)\n\t\t\t\tvar se *SchemaError\n\t\t\t\trequire.ErrorAs(t, err, &se)\n\t\t\t\tassert.Equal(t, ErrInvalidPredictRef, se.Kind)\n\t\t\t} else {\n\t\t\t\trequire.NoError(t, err)\n\t\t\t\tassert.Equal(t, tt.file, file)\n\t\t\t\tassert.Equal(t, tt.name, name)\n\t\t\t}\n\t\t})\n\t}\n}\n\n// ---------------------------------------------------------------------------\n// GenerateFromSource\n// ---------------------------------------------------------------------------\n\nfunc TestGenerateFromSource(t *testing.T) {\n\tdata, err := GenerateFromSource([]byte(\"unused\"), \"Predictor\", ModePredict, mockParser, \"\")\n\trequire.NoError(t, err)\n\n\tvar spec map[string]any\n\trequire.NoError(t, json.Unmarshal(data, &spec))\n\n\tassert.Equal(t, \"3.0.2\", spec[\"openapi\"])\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\tassert.Contains(t, props, \"prompt\")\n}\n\nfunc TestGenerateFromSourceTrainMode(t *testing.T) {\n\tdata, err := GenerateFromSource([]byte(\"unused\"), \"Trainer\", ModeTrain, mockParser, \"\")\n\trequire.NoError(t, err)\n\n\tvar spec map[string]any\n\trequire.NoError(t, json.Unmarshal(data, &spec))\n\n\tassert.NotNil(t, getPath(spec, \"components\", \"schemas\", \"TrainingInput\"))\n\tassert.NotNil(t, getPath(spec, \"paths\", \"/trainings\", \"post\"))\n}\n\nfunc TestGenerateFromSourceParseError(t *testing.T) {\n\t_, err := GenerateFromSource([]byte(\"unused\"), \"Predictor\", ModePredict, failParser, \"\")\n\trequire.Error(t, err)\n\tassert.Contains(t, err.Error(), \"mock parse failure\")\n}\n\n// ---------------------------------------------------------------------------\n// Generate — file-based\n// ---------------------------------------------------------------------------\n\nfunc TestGenerateReadsFile(t *testing.T) {\n\tdir := t.TempDir()\n\terr := os.WriteFile(filepath.Join(dir, \"predict.py\"), []byte(\"class Predictor: pass\"), 0o644)\n\trequire.NoError(t, err)\n\n\tdata, err := Generate(\"predict.py:Predictor\", dir, ModePredict, mockParser)\n\trequire.NoError(t, err)\n\n\tvar spec map[string]any\n\trequire.NoError(t, json.Unmarshal(data, &spec))\n\tassert.Equal(t, \"3.0.2\", spec[\"openapi\"])\n}\n\nfunc TestGenerateMissingFile(t *testing.T) {\n\tdir := t.TempDir()\n\t_, err := Generate(\"missing.py:Predictor\", dir, ModePredict, mockParser)\n\trequire.Error(t, err)\n\tassert.Contains(t, err.Error(), \"failed to read predictor source\")\n}\n\nfunc TestGenerateInvalidRef(t *testing.T) {\n\t_, err := Generate(\"no_colon\", \".\", ModePredict, mockParser)\n\trequire.Error(t, err)\n\tvar se *SchemaError\n\trequire.ErrorAs(t, err, &se)\n\tassert.Equal(t, ErrInvalidPredictRef, se.Kind)\n}\n\n// ---------------------------------------------------------------------------\n// COG_OPENAPI_SCHEMA env var\n// ---------------------------------------------------------------------------\n\nfunc TestGenerateCogOpenAPISchemaEnv(t *testing.T) {\n\t// Write a pre-built schema file\n\tdir := t.TempDir()\n\tschemaContent := `{\"openapi\": \"3.0.2\", \"info\": {\"title\": \"Custom\"}}`\n\tschemaPath := filepath.Join(dir, \"custom_schema.json\")\n\terr := os.WriteFile(schemaPath, []byte(schemaContent), 0o644)\n\trequire.NoError(t, err)\n\n\tt.Setenv(\"COG_OPENAPI_SCHEMA\", schemaPath)\n\n\t// Should return the custom schema without parsing\n\t// (using failParser to prove parsing is skipped)\n\tdata, err := Generate(\"predict.py:Predictor\", \".\", ModePredict, failParser)\n\trequire.NoError(t, err)\n\tassert.Equal(t, schemaContent, string(data))\n}\n\nfunc TestGenerateCogOpenAPISchemaEnvMissingFile(t *testing.T) {\n\tt.Setenv(\"COG_OPENAPI_SCHEMA\", \"/nonexistent/schema.json\")\n\n\t_, err := Generate(\"predict.py:Predictor\", \".\", ModePredict, mockParser)\n\trequire.Error(t, err)\n\tassert.Contains(t, err.Error(), \"COG_OPENAPI_SCHEMA\")\n\tassert.Contains(t, err.Error(), \"failed to read\")\n}\n\nfunc TestGenerateCogOpenAPISchemaEnvNotSet(t *testing.T) {\n\t// Ensure env var is not set\n\tt.Setenv(\"COG_OPENAPI_SCHEMA\", \"\")\n\n\tdir := t.TempDir()\n\terr := os.WriteFile(filepath.Join(dir, \"predict.py\"), []byte(\"class Predictor: pass\"), 0o644)\n\trequire.NoError(t, err)\n\n\t// Should proceed with normal generation (not use env var)\n\tdata, err := Generate(\"predict.py:Predictor\", dir, ModePredict, mockParser)\n\trequire.NoError(t, err)\n\n\tvar spec map[string]any\n\trequire.NoError(t, json.Unmarshal(data, &spec))\n\tassert.Equal(t, \"Cog\", getPath(spec, \"info\", \"title\"))\n}\n"
  },
  {
    "path": "pkg/schema/openapi.go",
    "content": "package schema\n\nimport (\n\t\"encoding/json\"\n\t\"maps\"\n\t\"sort\"\n)\n\n// GenerateOpenAPISchema produces a complete OpenAPI 3.0.2 specification\n// from a PredictorInfo. The returned bytes are compact JSON.\nfunc GenerateOpenAPISchema(info *PredictorInfo) ([]byte, error) {\n\tspec := buildOpenAPISpec(info)\n\n\t// Post-processing: remove title next to $ref, fix nullable anyOf\n\tremoveTitleNextToRef(spec)\n\tfixNullableAnyOf(spec)\n\n\treturn json.Marshal(spec)\n}\n\n// buildOpenAPISpec constructs the full OpenAPI 3.0.2 map.\nfunc buildOpenAPISpec(info *PredictorInfo) map[string]any {\n\tinputSchema, enumSchemas := buildInputSchema(info)\n\toutputSchema := info.Output.JSONSchema()\n\n\tisTrain := info.Mode == ModeTrain\n\n\tvar (\n\t\tendpoint     string\n\t\trequestName  string\n\t\tresponseName string\n\t\tcancelEP     string\n\t\tsummary      string\n\t\tdescription  string\n\t\topID         string\n\t\tcancelOpID   string\n\t\tcancelParam  string\n\t\tinputKey     string\n\t\toutputKey    string\n\t)\n\n\tif isTrain {\n\t\tendpoint = \"/trainings\"\n\t\trequestName = \"TrainingRequest\"\n\t\tresponseName = \"TrainingResponse\"\n\t\tcancelEP = \"/trainings/{training_id}/cancel\"\n\t\tsummary = \"Train\"\n\t\tdescription = \"Run a single training on the model\"\n\t\topID = \"train_trainings_post\"\n\t\tcancelOpID = \"cancel_trainings__training_id__cancel_post\"\n\t\tcancelParam = \"training_id\"\n\t\tinputKey = \"TrainingInput\"\n\t\toutputKey = \"TrainingOutput\"\n\t} else {\n\t\tendpoint = \"/predictions\"\n\t\trequestName = \"PredictionRequest\"\n\t\tresponseName = \"PredictionResponse\"\n\t\tcancelEP = \"/predictions/{prediction_id}/cancel\"\n\t\tsummary = \"Predict\"\n\t\tdescription = \"Run a single prediction on the model\"\n\t\topID = \"predict_predictions_post\"\n\t\tcancelOpID = \"cancel_predictions__prediction_id__cancel_post\"\n\t\tcancelParam = \"prediction_id\"\n\t\tinputKey = \"Input\"\n\t\toutputKey = \"Output\"\n\t}\n\n\t// Build components/schemas\n\tcomponents := newOrderedMapAny()\n\n\t// Input schema\n\tinputSchema[\"title\"] = inputKey\n\tcomponents.Set(inputKey, inputSchema)\n\n\t// Output schema\n\tcomponents.Set(outputKey, outputSchema)\n\n\t// Enum schemas for choices\n\tfor _, es := range enumSchemas {\n\t\tcomponents.Set(es.name, es.schema)\n\t}\n\n\tinputRef := \"#/components/schemas/\" + inputKey\n\toutputRef := \"#/components/schemas/\" + outputKey\n\n\t// Request schema\n\tcomponents.Set(requestName, map[string]any{\n\t\t\"title\": requestName,\n\t\t\"type\":  \"object\",\n\t\t\"properties\": map[string]any{\n\t\t\t\"id\":    map[string]any{\"title\": \"Id\", \"type\": \"string\"},\n\t\t\t\"input\": map[string]any{\"$ref\": inputRef},\n\t\t},\n\t})\n\n\t// Response schema\n\tcomponents.Set(responseName, map[string]any{\n\t\t\"title\": responseName,\n\t\t\"type\":  \"object\",\n\t\t\"properties\": map[string]any{\n\t\t\t\"input\":        map[string]any{\"$ref\": inputRef},\n\t\t\t\"output\":       map[string]any{\"$ref\": outputRef},\n\t\t\t\"id\":           map[string]any{\"title\": \"Id\", \"type\": \"string\"},\n\t\t\t\"version\":      map[string]any{\"title\": \"Version\", \"type\": \"string\"},\n\t\t\t\"created_at\":   map[string]any{\"title\": \"Created At\", \"type\": \"string\", \"format\": \"date-time\"},\n\t\t\t\"started_at\":   map[string]any{\"title\": \"Started At\", \"type\": \"string\", \"format\": \"date-time\"},\n\t\t\t\"completed_at\": map[string]any{\"title\": \"Completed At\", \"type\": \"string\", \"format\": \"date-time\"},\n\t\t\t\"status\":       map[string]any{\"title\": \"Status\", \"type\": \"string\"},\n\t\t\t\"error\":        map[string]any{\"title\": \"Error\", \"type\": \"string\"},\n\t\t\t\"logs\":         map[string]any{\"title\": \"Logs\", \"type\": \"string\"},\n\t\t\t\"metrics\":      map[string]any{\"title\": \"Metrics\", \"type\": \"object\"},\n\t\t},\n\t})\n\n\t// Status enum\n\tcomponents.Set(\"Status\", map[string]any{\n\t\t\"title\":       \"Status\",\n\t\t\"description\": \"An enumeration.\",\n\t\t\"enum\":        []any{\"starting\", \"processing\", \"succeeded\", \"canceled\", \"failed\"},\n\t\t\"type\":        \"string\",\n\t})\n\n\t// Validation error schemas\n\tcomponents.Set(\"HTTPValidationError\", map[string]any{\n\t\t\"title\": \"HTTPValidationError\",\n\t\t\"type\":  \"object\",\n\t\t\"properties\": map[string]any{\n\t\t\t\"detail\": map[string]any{\n\t\t\t\t\"title\": \"Detail\",\n\t\t\t\t\"type\":  \"array\",\n\t\t\t\t\"items\": map[string]any{\"$ref\": \"#/components/schemas/ValidationError\"},\n\t\t\t},\n\t\t},\n\t})\n\n\tcomponents.Set(\"ValidationError\", map[string]any{\n\t\t\"title\":    \"ValidationError\",\n\t\t\"required\": []any{\"loc\", \"msg\", \"type\"},\n\t\t\"type\":     \"object\",\n\t\t\"properties\": map[string]any{\n\t\t\t\"loc\": map[string]any{\n\t\t\t\t\"title\": \"Location\",\n\t\t\t\t\"type\":  \"array\",\n\t\t\t\t\"items\": map[string]any{\n\t\t\t\t\t\"anyOf\": []any{\n\t\t\t\t\t\tmap[string]any{\"type\": \"string\"},\n\t\t\t\t\t\tmap[string]any{\"type\": \"integer\"},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t\"msg\":  map[string]any{\"title\": \"Message\", \"type\": \"string\"},\n\t\t\t\"type\": map[string]any{\"title\": \"Error Type\", \"type\": \"string\"},\n\t\t},\n\t})\n\n\trequestRef := \"#/components/schemas/\" + requestName\n\tresponseRef := \"#/components/schemas/\" + responseName\n\n\t// Build paths\n\tpaths := newOrderedMapAny()\n\n\t// Root\n\tpaths.Set(\"/\", map[string]any{\n\t\t\"get\": map[string]any{\n\t\t\t\"summary\":     \"Root\",\n\t\t\t\"operationId\": \"root__get\",\n\t\t\t\"responses\": map[string]any{\n\t\t\t\t\"200\": map[string]any{\n\t\t\t\t\t\"description\": \"Successful Response\",\n\t\t\t\t\t\"content\":     map[string]any{\"application/json\": map[string]any{\"schema\": map[string]any{}}},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\t// Health check\n\tpaths.Set(\"/health-check\", map[string]any{\n\t\t\"get\": map[string]any{\n\t\t\t\"summary\":     \"Healthcheck\",\n\t\t\t\"operationId\": \"healthcheck_health_check_get\",\n\t\t\t\"responses\": map[string]any{\n\t\t\t\t\"200\": map[string]any{\n\t\t\t\t\t\"description\": \"Successful Response\",\n\t\t\t\t\t\"content\":     map[string]any{\"application/json\": map[string]any{\"schema\": map[string]any{}}},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\t// Main endpoint (predict or train)\n\tpaths.Set(endpoint, map[string]any{\n\t\t\"post\": map[string]any{\n\t\t\t\"summary\":     summary,\n\t\t\t\"description\": description,\n\t\t\t\"operationId\": opID,\n\t\t\t\"requestBody\": map[string]any{\n\t\t\t\t\"content\": map[string]any{\n\t\t\t\t\t\"application/json\": map[string]any{\n\t\t\t\t\t\t\"schema\": map[string]any{\"$ref\": requestRef},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t\t\"responses\": map[string]any{\n\t\t\t\t\"200\": map[string]any{\n\t\t\t\t\t\"description\": \"Successful Response\",\n\t\t\t\t\t\"content\": map[string]any{\n\t\t\t\t\t\t\"application/json\": map[string]any{\n\t\t\t\t\t\t\t\"schema\": map[string]any{\"$ref\": responseRef},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t\t\"422\": map[string]any{\n\t\t\t\t\t\"description\": \"Validation Error\",\n\t\t\t\t\t\"content\": map[string]any{\n\t\t\t\t\t\t\"application/json\": map[string]any{\n\t\t\t\t\t\t\t\"schema\": map[string]any{\"$ref\": \"#/components/schemas/HTTPValidationError\"},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\t// Cancel endpoint\n\tpaths.Set(cancelEP, map[string]any{\n\t\t\"post\": map[string]any{\n\t\t\t\"summary\":     \"Cancel\",\n\t\t\t\"operationId\": cancelOpID,\n\t\t\t\"parameters\": []any{\n\t\t\t\tmap[string]any{\n\t\t\t\t\t\"required\": true,\n\t\t\t\t\t\"schema\":   map[string]any{\"title\": TitleCaseSingle(cancelParam), \"type\": \"string\"},\n\t\t\t\t\t\"name\":     cancelParam,\n\t\t\t\t\t\"in\":       \"path\",\n\t\t\t\t},\n\t\t\t},\n\t\t\t\"responses\": map[string]any{\n\t\t\t\t\"200\": map[string]any{\n\t\t\t\t\t\"description\": \"Successful Response\",\n\t\t\t\t\t\"content\":     map[string]any{\"application/json\": map[string]any{\"schema\": map[string]any{}}},\n\t\t\t\t},\n\t\t\t\t\"422\": map[string]any{\n\t\t\t\t\t\"description\": \"Validation Error\",\n\t\t\t\t\t\"content\": map[string]any{\n\t\t\t\t\t\t\"application/json\": map[string]any{\n\t\t\t\t\t\t\t\"schema\": map[string]any{\"$ref\": \"#/components/schemas/HTTPValidationError\"},\n\t\t\t\t\t\t},\n\t\t\t\t\t},\n\t\t\t\t},\n\t\t\t},\n\t\t},\n\t})\n\n\treturn map[string]any{\n\t\t\"openapi\": \"3.0.2\",\n\t\t\"info\":    map[string]any{\"title\": \"Cog\", \"version\": \"0.1.0\"},\n\t\t\"paths\":   paths,\n\t\t\"components\": map[string]any{\n\t\t\t\"schemas\": components,\n\t\t},\n\t}\n}\n\n// enumSchema pairs a name with its schema for choices fields.\ntype enumSchema struct {\n\tname   string\n\tschema map[string]any\n}\n\n// buildInputSchema builds the Input schema object and any enum schemas for choices.\nfunc buildInputSchema(info *PredictorInfo) (map[string]any, []enumSchema) {\n\tproperties := newOrderedMapAny()\n\tvar required []string\n\tvar enums []enumSchema\n\n\tinfo.Inputs.Entries(func(name string, field InputField) {\n\t\tprop := newOrderedMapAny()\n\n\t\t// x-order for field ordering\n\t\tprop.Set(\"x-order\", field.Order)\n\n\t\tif len(field.Choices) > 0 {\n\t\t\t// Choices -> use allOf with $ref to enum schema\n\t\t\tenumName := TitleCaseSingle(name)\n\t\t\tenumType := field.FieldType.Primitive.JSONType()\n\t\t\ttypeStr, _ := enumType[\"type\"].(string)\n\t\t\tif typeStr == \"\" {\n\t\t\t\ttypeStr = \"string\"\n\t\t\t}\n\n\t\t\tchoiceValues := make([]any, len(field.Choices))\n\t\t\tfor i, c := range field.Choices {\n\t\t\t\tchoiceValues[i] = c.ToJSON()\n\t\t\t}\n\n\t\t\tenums = append(enums, enumSchema{\n\t\t\t\tname: enumName,\n\t\t\t\tschema: map[string]any{\n\t\t\t\t\t\"title\":       enumName,\n\t\t\t\t\t\"description\": \"An enumeration.\",\n\t\t\t\t\t\"enum\":        choiceValues,\n\t\t\t\t\t\"type\":        typeStr,\n\t\t\t\t},\n\t\t\t})\n\n\t\t\tprop.Set(\"allOf\", []any{\n\t\t\t\tmap[string]any{\"$ref\": \"#/components/schemas/\" + enumName},\n\t\t\t})\n\t\t} else {\n\t\t\t// Regular field — inline type\n\t\t\tprop.Set(\"title\", TitleCase(name))\n\t\t\ttypeSchema := field.FieldType.JSONType()\n\t\t\t// Merge type schema keys into prop in sorted order for determinism\n\t\t\tkeys := make([]string, 0, len(typeSchema))\n\t\t\tfor k := range typeSchema {\n\t\t\t\tkeys = append(keys, k)\n\t\t\t}\n\t\t\tsort.Strings(keys)\n\t\t\tfor _, k := range keys {\n\t\t\t\tprop.Set(k, typeSchema[k])\n\t\t\t}\n\t\t}\n\n\t\t// Required?\n\t\tif field.IsRequired() {\n\t\t\trequired = append(required, name)\n\t\t}\n\n\t\t// Default value\n\t\tif field.Default != nil {\n\t\t\tprop.Set(\"default\", field.Default.ToJSON())\n\t\t}\n\n\t\t// Nullable\n\t\tif field.FieldType.Repetition == Optional {\n\t\t\tprop.Set(\"nullable\", true)\n\t\t}\n\n\t\t// Description\n\t\tif field.Description != nil {\n\t\t\tprop.Set(\"description\", *field.Description)\n\t\t}\n\n\t\t// Numeric constraints\n\t\tif field.GE != nil {\n\t\t\tprop.Set(\"minimum\", *field.GE)\n\t\t}\n\t\tif field.LE != nil {\n\t\t\tprop.Set(\"maximum\", *field.LE)\n\t\t}\n\n\t\t// String constraints\n\t\tif field.MinLength != nil {\n\t\t\tprop.Set(\"minLength\", *field.MinLength)\n\t\t}\n\t\tif field.MaxLength != nil {\n\t\t\tprop.Set(\"maxLength\", *field.MaxLength)\n\t\t}\n\t\tif field.Regex != nil {\n\t\t\tprop.Set(\"pattern\", *field.Regex)\n\t\t}\n\n\t\t// Deprecated\n\t\tif field.Deprecated != nil && *field.Deprecated {\n\t\t\tprop.Set(\"deprecated\", true)\n\t\t}\n\n\t\tproperties.Set(name, prop)\n\t})\n\n\tinputSchema := map[string]any{\n\t\t\"title\":      \"Input\",\n\t\t\"type\":       \"object\",\n\t\t\"properties\": properties,\n\t}\n\n\tif len(required) > 0 {\n\t\tinputSchema[\"required\"] = required\n\t}\n\n\treturn inputSchema, enums\n}\n\n// ---------------------------------------------------------------------------\n// Post-processing (mirrors openapi_schema.py fixups)\n// ---------------------------------------------------------------------------\n\n// removeTitleNextToRef removes \"title\" from any map that also has \"$ref\".\n// OpenAPI 3.0 doesn't allow sibling keywords next to $ref.\nfunc removeTitleNextToRef(v any) {\n\tswitch val := v.(type) {\n\tcase map[string]any:\n\t\tif _, hasRef := val[\"$ref\"]; hasRef {\n\t\t\tdelete(val, \"title\")\n\t\t}\n\t\tfor _, child := range val {\n\t\t\tremoveTitleNextToRef(child)\n\t\t}\n\tcase *orderedMapAny:\n\t\tif _, hasRef := val.Get(\"$ref\"); hasRef {\n\t\t\tval.Delete(\"title\")\n\t\t}\n\t\tval.Entries(func(_ string, child any) {\n\t\t\tremoveTitleNextToRef(child)\n\t\t})\n\tcase []any:\n\t\tfor _, child := range val {\n\t\t\tremoveTitleNextToRef(child)\n\t\t}\n\t}\n}\n\n// fixNullableAnyOf converts {\"anyOf\": [{\"type\": T}, {\"type\": \"null\"}]} to\n// {\"type\": T, \"nullable\": true}. OpenAPI 3.0 uses nullable instead of union-with-null.\nfunc fixNullableAnyOf(v any) {\n\tswitch val := v.(type) {\n\tcase map[string]any:\n\t\t// Recurse first\n\t\tfor _, child := range val {\n\t\t\tfixNullableAnyOf(child)\n\t\t}\n\t\t// Check for anyOf with null pattern\n\t\tanyOf, ok := val[\"anyOf\"].([]any)\n\t\tif !ok || len(anyOf) != 2 {\n\t\t\treturn\n\t\t}\n\t\tvar nonNull map[string]any\n\t\thasNull := false\n\t\tfor _, variant := range anyOf {\n\t\t\tm, ok := variant.(map[string]any)\n\t\t\tif !ok {\n\t\t\t\treturn\n\t\t\t}\n\t\t\tif t, _ := m[\"type\"].(string); t == \"null\" {\n\t\t\t\thasNull = true\n\t\t\t} else {\n\t\t\t\tnonNull = m\n\t\t\t}\n\t\t}\n\t\tif hasNull && nonNull != nil {\n\t\t\tdelete(val, \"anyOf\")\n\t\t\tmaps.Copy(val, nonNull)\n\t\t\tval[\"nullable\"] = true\n\t\t}\n\tcase *orderedMapAny:\n\t\t// Recurse first\n\t\tval.Entries(func(_ string, child any) {\n\t\t\tfixNullableAnyOf(child)\n\t\t})\n\t\t// Check for anyOf with null pattern\n\t\tanyOfRaw, ok := val.Get(\"anyOf\")\n\t\tif !ok {\n\t\t\treturn\n\t\t}\n\t\tanyOf, ok := anyOfRaw.([]any)\n\t\tif !ok || len(anyOf) != 2 {\n\t\t\treturn\n\t\t}\n\t\tvar nonNull map[string]any\n\t\thasNull := false\n\t\tfor _, variant := range anyOf {\n\t\t\tm, ok := variant.(map[string]any)\n\t\t\tif !ok {\n\t\t\t\treturn\n\t\t\t}\n\t\t\tif t, _ := m[\"type\"].(string); t == \"null\" {\n\t\t\t\thasNull = true\n\t\t\t} else {\n\t\t\t\tnonNull = m\n\t\t\t}\n\t\t}\n\t\tif hasNull && nonNull != nil {\n\t\t\tval.Delete(\"anyOf\")\n\t\t\tfor k, v := range nonNull {\n\t\t\t\tval.Set(k, v)\n\t\t\t}\n\t\t\tval.Set(\"nullable\", true)\n\t\t}\n\tcase []any:\n\t\tfor _, child := range val {\n\t\t\tfixNullableAnyOf(child)\n\t\t}\n\t}\n}\n\n// ---------------------------------------------------------------------------\n// orderedMapAny — ordered map with JSON marshaling that preserves key order.\n// Used for schema properties where field ordering matters.\n// ---------------------------------------------------------------------------\n\ntype orderedMapAny struct {\n\tkeys   []string\n\tvalues map[string]any\n}\n\nfunc newOrderedMapAny() *orderedMapAny {\n\treturn &orderedMapAny{values: make(map[string]any)}\n}\n\nfunc (m *orderedMapAny) Set(key string, value any) {\n\tif _, exists := m.values[key]; !exists {\n\t\tm.keys = append(m.keys, key)\n\t}\n\tm.values[key] = value\n}\n\nfunc (m *orderedMapAny) Get(key string) (any, bool) {\n\tv, ok := m.values[key]\n\treturn v, ok\n}\n\nfunc (m *orderedMapAny) Delete(key string) {\n\tif _, exists := m.values[key]; !exists {\n\t\treturn\n\t}\n\tdelete(m.values, key)\n\tfor i, k := range m.keys {\n\t\tif k == key {\n\t\t\tm.keys = append(m.keys[:i], m.keys[i+1:]...)\n\t\t\tbreak\n\t\t}\n\t}\n}\n\nfunc (m *orderedMapAny) Entries(fn func(key string, value any)) {\n\tfor _, k := range m.keys {\n\t\tfn(k, m.values[k])\n\t}\n}\n\n// MarshalJSON produces a JSON object with keys in insertion order.\nfunc (m *orderedMapAny) MarshalJSON() ([]byte, error) {\n\tbuf := []byte{'{'}\n\tfor i, k := range m.keys {\n\t\tif i > 0 {\n\t\t\tbuf = append(buf, ',')\n\t\t}\n\t\tkeyBytes, err := json.Marshal(k)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tbuf = append(buf, keyBytes...)\n\t\tbuf = append(buf, ':')\n\t\tvalBytes, err := json.Marshal(m.values[k])\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tbuf = append(buf, valBytes...)\n\t}\n\tbuf = append(buf, '}')\n\treturn buf, nil\n}\n"
  },
  {
    "path": "pkg/schema/openapi_test.go",
    "content": "package schema\n\nimport (\n\t\"encoding/json\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n)\n\n// ---------------------------------------------------------------------------\n// Helpers\n// ---------------------------------------------------------------------------\n\nfunc simplePredictor() *PredictorInfo {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"s\", InputField{\n\t\tName:      \"s\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeString, Repetition: Required},\n\t})\n\n\treturn &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n}\n\nfunc ptr[T any](v T) *T { return &v }\n\n// parseSpec is a test helper that generates the schema and unmarshals\n// it into a generic map for assertion.\nfunc parseSpec(t *testing.T, info *PredictorInfo) map[string]any {\n\tt.Helper()\n\tdata, err := GenerateOpenAPISchema(info)\n\trequire.NoError(t, err)\n\tvar spec map[string]any\n\trequire.NoError(t, json.Unmarshal(data, &spec))\n\treturn spec\n}\n\nfunc getPath(m map[string]any, keys ...string) any {\n\tvar cur any = m\n\tfor _, k := range keys {\n\t\tobj, ok := cur.(map[string]any)\n\t\tif !ok {\n\t\t\treturn nil\n\t\t}\n\t\tcur = obj[k]\n\t}\n\treturn cur\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Top-level structure\n// ---------------------------------------------------------------------------\n\nfunc TestGeneratesValidOpenAPI(t *testing.T) {\n\tspec := parseSpec(t, simplePredictor())\n\n\tassert.Equal(t, \"3.0.2\", spec[\"openapi\"])\n\tassert.Equal(t, \"Cog\", getPath(spec, \"info\", \"title\"))\n\tassert.Equal(t, \"0.1.0\", getPath(spec, \"info\", \"version\"))\n}\n\nfunc TestPredictEndpoints(t *testing.T) {\n\tspec := parseSpec(t, simplePredictor())\n\n\t// Root\n\tassert.NotNil(t, getPath(spec, \"paths\", \"/\", \"get\"))\n\t// Health check\n\tassert.NotNil(t, getPath(spec, \"paths\", \"/health-check\", \"get\"))\n\t// Predictions\n\tpost := getPath(spec, \"paths\", \"/predictions\", \"post\")\n\trequire.NotNil(t, post)\n\tpostMap := post.(map[string]any)\n\tassert.Equal(t, \"Predict\", postMap[\"summary\"])\n\tassert.Equal(t, \"predict_predictions_post\", postMap[\"operationId\"])\n\t// Cancel\n\tassert.NotNil(t, getPath(spec, \"paths\", \"/predictions/{prediction_id}/cancel\", \"post\"))\n}\n\nfunc TestTrainEndpoints(t *testing.T) {\n\tinfo := simplePredictor()\n\tinfo.Mode = ModeTrain\n\tspec := parseSpec(t, info)\n\n\tpost := getPath(spec, \"paths\", \"/trainings\", \"post\")\n\trequire.NotNil(t, post)\n\tpostMap := post.(map[string]any)\n\tassert.Equal(t, \"Train\", postMap[\"summary\"])\n\tassert.Equal(t, \"train_trainings_post\", postMap[\"operationId\"])\n\n\t// Cancel\n\tcancel := getPath(spec, \"paths\", \"/trainings/{training_id}/cancel\", \"post\")\n\trequire.NotNil(t, cancel)\n\n\t// Schema keys use TrainingInput/TrainingOutput\n\tassert.NotNil(t, getPath(spec, \"components\", \"schemas\", \"TrainingInput\"))\n\tassert.NotNil(t, getPath(spec, \"components\", \"schemas\", \"TrainingOutput\"))\n\tassert.NotNil(t, getPath(spec, \"components\", \"schemas\", \"TrainingRequest\"))\n\tassert.NotNil(t, getPath(spec, \"components\", \"schemas\", \"TrainingResponse\"))\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Fixed components\n// ---------------------------------------------------------------------------\n\nfunc TestFixedComponentSchemas(t *testing.T) {\n\tspec := parseSpec(t, simplePredictor())\n\tschemas := getPath(spec, \"components\", \"schemas\").(map[string]any)\n\n\t// PredictionRequest\n\treq := schemas[\"PredictionRequest\"].(map[string]any)\n\tassert.Equal(t, \"PredictionRequest\", req[\"title\"])\n\tprops := req[\"properties\"].(map[string]any)\n\tassert.Equal(t, \"#/components/schemas/Input\", getPath(props, \"input\", \"$ref\"))\n\tassert.Equal(t, \"string\", getPath(props, \"id\", \"type\"))\n\n\t// PredictionResponse\n\tresp := schemas[\"PredictionResponse\"].(map[string]any)\n\tassert.Equal(t, \"PredictionResponse\", resp[\"title\"])\n\trespProps := resp[\"properties\"].(map[string]any)\n\tassert.Equal(t, \"#/components/schemas/Input\", getPath(respProps, \"input\", \"$ref\"))\n\tassert.Equal(t, \"#/components/schemas/Output\", getPath(respProps, \"output\", \"$ref\"))\n\n\t// Status\n\tstatus := schemas[\"Status\"].(map[string]any)\n\tassert.Equal(t, \"string\", status[\"type\"])\n\tenum := status[\"enum\"].([]any)\n\tassert.Contains(t, enum, \"starting\")\n\tassert.Contains(t, enum, \"succeeded\")\n\n\t// Validation errors\n\tassert.NotNil(t, schemas[\"HTTPValidationError\"])\n\tassert.NotNil(t, schemas[\"ValidationError\"])\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Input schema\n// ---------------------------------------------------------------------------\n\nfunc TestInputRequiredField(t *testing.T) {\n\tspec := parseSpec(t, simplePredictor())\n\tinput := getPath(spec, \"components\", \"schemas\", \"Input\").(map[string]any)\n\n\trequired := input[\"required\"].([]any)\n\tassert.Contains(t, required, \"s\")\n}\n\nfunc TestInputOptionalFieldNotRequired(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"name\", InputField{\n\t\tName:      \"name\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeString, Repetition: Optional},\n\t\tDefault:   &DefaultValue{Kind: DefaultNone},\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tinput := getPath(spec, \"components\", \"schemas\", \"Input\").(map[string]any)\n\n\t// Should not have required since there's a default\n\tassert.Nil(t, input[\"required\"])\n\n\t// Should have nullable\n\tprops := input[\"properties\"].(map[string]any)\n\tnameField := props[\"name\"].(map[string]any)\n\tassert.Equal(t, true, nameField[\"nullable\"])\n}\n\nfunc TestInputDefaultValue(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"count\", InputField{\n\t\tName:      \"count\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeInteger, Repetition: Required},\n\t\tDefault:   &DefaultValue{Kind: DefaultInt, Int: 42},\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\tcountField := props[\"count\"].(map[string]any)\n\t// JSON numbers unmarshal as float64\n\tassert.Equal(t, float64(42), countField[\"default\"])\n}\n\nfunc TestInputDescription(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"text\", InputField{\n\t\tName:        \"text\",\n\t\tOrder:       0,\n\t\tFieldType:   FieldType{Primitive: TypeString, Repetition: Required},\n\t\tDescription: ptr(\"The input text\"),\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\ttextField := props[\"text\"].(map[string]any)\n\tassert.Equal(t, \"The input text\", textField[\"description\"])\n}\n\nfunc TestInputNumericConstraints(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"temperature\", InputField{\n\t\tName:      \"temperature\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeFloat, Repetition: Required},\n\t\tGE:        ptr(0.0),\n\t\tLE:        ptr(1.0),\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\ttempField := props[\"temperature\"].(map[string]any)\n\tassert.Equal(t, float64(0), tempField[\"minimum\"])\n\tassert.Equal(t, float64(1), tempField[\"maximum\"])\n}\n\nfunc TestInputStringConstraints(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"name\", InputField{\n\t\tName:      \"name\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeString, Repetition: Required},\n\t\tMinLength: ptr[uint64](1),\n\t\tMaxLength: ptr[uint64](100),\n\t\tRegex:     ptr(\"^[a-z]+$\"),\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\tnameField := props[\"name\"].(map[string]any)\n\tassert.Equal(t, float64(1), nameField[\"minLength\"])\n\tassert.Equal(t, float64(100), nameField[\"maxLength\"])\n\tassert.Equal(t, \"^[a-z]+$\", nameField[\"pattern\"])\n}\n\nfunc TestInputDeprecated(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"old_param\", InputField{\n\t\tName:       \"old_param\",\n\t\tOrder:      0,\n\t\tFieldType:  FieldType{Primitive: TypeString, Repetition: Required},\n\t\tDeprecated: ptr(true),\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\tfield := props[\"old_param\"].(map[string]any)\n\tassert.Equal(t, true, field[\"deprecated\"])\n}\n\nfunc TestInputXOrder(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"first\", InputField{\n\t\tName:      \"first\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeString, Repetition: Required},\n\t})\n\tinputs.Set(\"second\", InputField{\n\t\tName:      \"second\",\n\t\tOrder:     1,\n\t\tFieldType: FieldType{Primitive: TypeInteger, Repetition: Required},\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\tassert.Equal(t, float64(0), props[\"first\"].(map[string]any)[\"x-order\"])\n\tassert.Equal(t, float64(1), props[\"second\"].(map[string]any)[\"x-order\"])\n}\n\nfunc TestInputRepeatedType(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"items\", InputField{\n\t\tName:      \"items\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeString, Repetition: Repeated},\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\titemsField := props[\"items\"].(map[string]any)\n\tassert.Equal(t, \"array\", itemsField[\"type\"])\n\titems := itemsField[\"items\"].(map[string]any)\n\tassert.Equal(t, \"string\", items[\"type\"])\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Choices / Enums\n// ---------------------------------------------------------------------------\n\nfunc TestChoicesGenerateEnum(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"color\", InputField{\n\t\tName:      \"color\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeString, Repetition: Required},\n\t\tChoices: []DefaultValue{\n\t\t\t{Kind: DefaultString, Str: \"red\"},\n\t\t\t{Kind: DefaultString, Str: \"blue\"},\n\t\t},\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\n\t// Enum schema created\n\tschemas := getPath(spec, \"components\", \"schemas\").(map[string]any)\n\tcolorEnum := schemas[\"Color\"].(map[string]any)\n\tassert.Equal(t, \"Color\", colorEnum[\"title\"])\n\tassert.Equal(t, \"string\", colorEnum[\"type\"])\n\tassert.Equal(t, []any{\"red\", \"blue\"}, colorEnum[\"enum\"])\n\n\t// Property uses allOf $ref\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\tcolorProp := props[\"color\"].(map[string]any)\n\tallOf := colorProp[\"allOf\"].([]any)\n\tassert.Len(t, allOf, 1)\n\tref := allOf[0].(map[string]any)\n\tassert.Equal(t, \"#/components/schemas/Color\", ref[\"$ref\"])\n}\n\nfunc TestIntegerChoices(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"size\", InputField{\n\t\tName:      \"size\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeInteger, Repetition: Required},\n\t\tChoices: []DefaultValue{\n\t\t\t{Kind: DefaultInt, Int: 256},\n\t\t\t{Kind: DefaultInt, Int: 512},\n\t\t\t{Kind: DefaultInt, Int: 1024},\n\t\t},\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tschemas := getPath(spec, \"components\", \"schemas\").(map[string]any)\n\tsizeEnum := schemas[\"Size\"].(map[string]any)\n\tassert.Equal(t, \"integer\", sizeEnum[\"type\"])\n\t// JSON numbers are float64\n\tassert.Equal(t, []any{float64(256), float64(512), float64(1024)}, sizeEnum[\"enum\"])\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Output types\n// ---------------------------------------------------------------------------\n\nfunc TestOutputSingle(t *testing.T) {\n\tspec := parseSpec(t, simplePredictor())\n\toutput := getPath(spec, \"components\", \"schemas\", \"Output\").(map[string]any)\n\tassert.Equal(t, \"Output\", output[\"title\"])\n\tassert.Equal(t, \"string\", output[\"type\"])\n}\n\nfunc TestOutputList(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaArrayOf(SchemaPrim(TypeString)),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\toutput := getPath(spec, \"components\", \"schemas\", \"Output\").(map[string]any)\n\tassert.Equal(t, \"Output\", output[\"title\"])\n\tassert.Equal(t, \"array\", output[\"type\"])\n\titems := output[\"items\"].(map[string]any)\n\tassert.Equal(t, \"string\", items[\"type\"])\n}\n\nfunc TestOutputIterator(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaIteratorOf(SchemaPrim(TypeString)),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\toutput := getPath(spec, \"components\", \"schemas\", \"Output\").(map[string]any)\n\tassert.Equal(t, \"array\", output[\"type\"])\n\tassert.Equal(t, \"iterator\", output[\"x-cog-array-type\"])\n}\n\nfunc TestOutputConcatenateIterator(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaConcatIteratorOf(),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\toutput := getPath(spec, \"components\", \"schemas\", \"Output\").(map[string]any)\n\tassert.Equal(t, \"array\", output[\"type\"])\n\tassert.Equal(t, \"iterator\", output[\"x-cog-array-type\"])\n\tassert.Equal(t, \"concatenate\", output[\"x-cog-array-display\"])\n}\n\nfunc TestOutputObject(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tfields := NewOrderedMap[string, SchemaField]()\n\tfields.Set(\"name\", SchemaField{\n\t\tType:     SchemaPrim(TypeString),\n\t\tRequired: true,\n\t})\n\tfields.Set(\"score\", SchemaField{\n\t\tType:     SchemaPrim(TypeFloat),\n\t\tRequired: true,\n\t})\n\tfields.Set(\"notes\", SchemaField{\n\t\tType:     SchemaType{Kind: SchemaPrimitive, Primitive: TypeString, Nullable: true},\n\t\tRequired: false,\n\t\tDefault:  &DefaultValue{Kind: DefaultNone},\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaObjectOf(fields),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\toutput := getPath(spec, \"components\", \"schemas\", \"Output\").(map[string]any)\n\tassert.Equal(t, \"object\", output[\"type\"])\n\tprops := output[\"properties\"].(map[string]any)\n\n\t// name\n\tnameField := props[\"name\"].(map[string]any)\n\tassert.Equal(t, \"string\", nameField[\"type\"])\n\tassert.Equal(t, \"Name\", nameField[\"title\"])\n\n\t// score\n\tscoreField := props[\"score\"].(map[string]any)\n\tassert.Equal(t, \"number\", scoreField[\"type\"])\n\n\t// notes — nullable\n\tnotesField := props[\"notes\"].(map[string]any)\n\tassert.Equal(t, true, notesField[\"nullable\"])\n\n\t// Required should include name and score but not notes\n\trequired := output[\"required\"].([]any)\n\tassert.Contains(t, required, \"name\")\n\tassert.Contains(t, required, \"score\")\n\tassert.NotContains(t, required, \"notes\")\n}\n\nfunc TestOutputPath(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypePath),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\toutput := getPath(spec, \"components\", \"schemas\", \"Output\").(map[string]any)\n\tassert.Equal(t, \"string\", output[\"type\"])\n\tassert.Equal(t, \"uri\", output[\"format\"])\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Post-processing\n// ---------------------------------------------------------------------------\n\nfunc TestRemoveTitleNextToRef(t *testing.T) {\n\tschema := map[string]any{\n\t\t\"title\": \"Foo\",\n\t\t\"$ref\":  \"#/components/schemas/Bar\",\n\t}\n\tremoveTitleNextToRef(schema)\n\tassert.Nil(t, schema[\"title\"])\n\tassert.Equal(t, \"#/components/schemas/Bar\", schema[\"$ref\"])\n}\n\nfunc TestRemoveTitleNextToRefNested(t *testing.T) {\n\tschema := map[string]any{\n\t\t\"properties\": map[string]any{\n\t\t\t\"inner\": map[string]any{\n\t\t\t\t\"title\": \"Inner\",\n\t\t\t\t\"$ref\":  \"#/components/schemas/Foo\",\n\t\t\t},\n\t\t},\n\t}\n\tremoveTitleNextToRef(schema)\n\tinner := schema[\"properties\"].(map[string]any)[\"inner\"].(map[string]any)\n\tassert.Nil(t, inner[\"title\"])\n}\n\nfunc TestFixNullableAnyOf(t *testing.T) {\n\tschema := map[string]any{\n\t\t\"anyOf\": []any{\n\t\t\tmap[string]any{\"type\": \"string\"},\n\t\t\tmap[string]any{\"type\": \"null\"},\n\t\t},\n\t}\n\tfixNullableAnyOf(schema)\n\tassert.Nil(t, schema[\"anyOf\"])\n\tassert.Equal(t, \"string\", schema[\"type\"])\n\tassert.Equal(t, true, schema[\"nullable\"])\n}\n\nfunc TestFixNullableAnyOfNoOp(t *testing.T) {\n\t// anyOf with no null should be left alone\n\tschema := map[string]any{\n\t\t\"anyOf\": []any{\n\t\t\tmap[string]any{\"type\": \"string\"},\n\t\t\tmap[string]any{\"type\": \"integer\"},\n\t\t},\n\t}\n\tfixNullableAnyOf(schema)\n\tassert.NotNil(t, schema[\"anyOf\"])\n\tassert.Nil(t, schema[\"nullable\"])\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Title case helpers\n// ---------------------------------------------------------------------------\n\nfunc TestTitleCaseWords(t *testing.T) {\n\tassert.Equal(t, \"Hello World\", TitleCase(\"hello_world\"))\n\tassert.Equal(t, \"Segmented Image\", TitleCase(\"segmented_image\"))\n\tassert.Equal(t, \"Name\", TitleCase(\"name\"))\n}\n\nfunc TestTitleCaseSingleWord(t *testing.T) {\n\tassert.Equal(t, \"Prediction_id\", TitleCaseSingle(\"prediction_id\"))\n\tassert.Equal(t, \"Color\", TitleCaseSingle(\"color\"))\n\tassert.Equal(t, \"\", TitleCaseSingle(\"\"))\n}\n\n// ---------------------------------------------------------------------------\n// Tests: JSON output is valid and parseable\n// ---------------------------------------------------------------------------\n\nfunc TestOutputIsValidJSON(t *testing.T) {\n\tdata, err := GenerateOpenAPISchema(simplePredictor())\n\trequire.NoError(t, err)\n\n\tvar parsed any\n\trequire.NoError(t, json.Unmarshal(data, &parsed))\n\tassert.NotNil(t, parsed)\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Multiple inputs with various types\n// ---------------------------------------------------------------------------\n\nfunc TestMultipleInputTypes(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinputs.Set(\"text\", InputField{\n\t\tName:      \"text\",\n\t\tOrder:     0,\n\t\tFieldType: FieldType{Primitive: TypeString, Repetition: Required},\n\t})\n\tinputs.Set(\"count\", InputField{\n\t\tName:      \"count\",\n\t\tOrder:     1,\n\t\tFieldType: FieldType{Primitive: TypeInteger, Repetition: Required},\n\t\tDefault:   &DefaultValue{Kind: DefaultInt, Int: 10},\n\t})\n\tinputs.Set(\"image\", InputField{\n\t\tName:      \"image\",\n\t\tOrder:     2,\n\t\tFieldType: FieldType{Primitive: TypePath, Repetition: Required},\n\t})\n\tinputs.Set(\"flag\", InputField{\n\t\tName:      \"flag\",\n\t\tOrder:     3,\n\t\tFieldType: FieldType{Primitive: TypeBool, Repetition: Required},\n\t\tDefault:   &DefaultValue{Kind: DefaultBool, Bool: false},\n\t})\n\tinputs.Set(\"secret_key\", InputField{\n\t\tName:      \"secret_key\",\n\t\tOrder:     4,\n\t\tFieldType: FieldType{Primitive: TypeSecret, Repetition: Optional},\n\t\tDefault:   &DefaultValue{Kind: DefaultNone},\n\t})\n\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tprops := getPath(spec, \"components\", \"schemas\", \"Input\", \"properties\").(map[string]any)\n\n\t// text - string\n\ttextField := props[\"text\"].(map[string]any)\n\tassert.Equal(t, \"string\", textField[\"type\"])\n\n\t// count - integer with default\n\tcountField := props[\"count\"].(map[string]any)\n\tassert.Equal(t, \"integer\", countField[\"type\"])\n\tassert.Equal(t, float64(10), countField[\"default\"])\n\n\t// image - path (URI)\n\timageField := props[\"image\"].(map[string]any)\n\tassert.Equal(t, \"string\", imageField[\"type\"])\n\tassert.Equal(t, \"uri\", imageField[\"format\"])\n\n\t// flag - boolean\n\tflagField := props[\"flag\"].(map[string]any)\n\tassert.Equal(t, \"boolean\", flagField[\"type\"])\n\n\t// secret_key - secret\n\tsecretField := props[\"secret_key\"].(map[string]any)\n\tassert.Equal(t, \"string\", secretField[\"type\"])\n\tassert.Equal(t, \"password\", secretField[\"format\"])\n\tassert.Equal(t, true, secretField[\"writeOnly\"])\n\tassert.Equal(t, true, secretField[\"x-cog-secret\"])\n\tassert.Equal(t, true, secretField[\"nullable\"])\n\n\t// Only text and image should be required (count has default, flag has default, secret has default)\n\trequired := getPath(spec, \"components\", \"schemas\", \"Input\", \"required\").([]any)\n\tassert.Contains(t, required, \"text\")\n\tassert.Contains(t, required, \"image\")\n\tassert.NotContains(t, required, \"count\")\n\tassert.NotContains(t, required, \"flag\")\n\tassert.NotContains(t, required, \"secret_key\")\n}\n\n// ---------------------------------------------------------------------------\n// Tests: Edge cases\n// ---------------------------------------------------------------------------\n\nfunc TestNoInputs(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaPrim(TypeString),\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\tinput := getPath(spec, \"components\", \"schemas\", \"Input\").(map[string]any)\n\tassert.Equal(t, \"object\", input[\"type\"])\n\t// required should not be present when there are no required fields\n\tassert.Nil(t, input[\"required\"])\n}\n\nfunc TestOutputObjectNoFields(t *testing.T) {\n\tinputs := NewOrderedMap[string, InputField]()\n\tinfo := &PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: SchemaType{Kind: SchemaObject},\n\t\tMode:   ModePredict,\n\t}\n\n\tspec := parseSpec(t, info)\n\toutput := getPath(spec, \"components\", \"schemas\", \"Output\").(map[string]any)\n\tassert.Equal(t, \"object\", output[\"type\"])\n}\n\n// ---------------------------------------------------------------------------\n// Tests: orderedMapAny JSON output preserves insertion order\n// ---------------------------------------------------------------------------\n\nfunc TestOrderedMapAnyJSON(t *testing.T) {\n\tm := newOrderedMapAny()\n\tm.Set(\"z\", 1)\n\tm.Set(\"a\", 2)\n\tm.Set(\"m\", 3)\n\n\tdata, err := json.Marshal(m)\n\trequire.NoError(t, err)\n\tassert.Equal(t, `{\"z\":1,\"a\":2,\"m\":3}`, string(data))\n}\n\nfunc TestOrderedMapAnyDelete(t *testing.T) {\n\tm := newOrderedMapAny()\n\tm.Set(\"a\", 1)\n\tm.Set(\"b\", 2)\n\tm.Set(\"c\", 3)\n\tm.Delete(\"b\")\n\n\tdata, err := json.Marshal(m)\n\trequire.NoError(t, err)\n\tassert.Equal(t, `{\"a\":1,\"c\":3}`, string(data))\n}\n"
  },
  {
    "path": "pkg/schema/python/parser.go",
    "content": "// Package python implements a tree-sitter based Python parser for extracting\n// Cog predictor signatures. It walks the concrete syntax tree to extract\n// imports, class definitions, function parameters with type annotations and\n// default values, and Input() call keyword arguments.\n//\n// This parser is Python-specific. Future languages (e.g. Node.js) would get\n// their own parser package under pkg/schema/.\npackage python\n\nimport (\n\t\"context\"\n\t\"errors\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"strconv\"\n\t\"strings\"\n\n\tsitter \"github.com/smacker/go-tree-sitter\"\n\t\"github.com/smacker/go-tree-sitter/python\"\n\n\t\"github.com/replicate/cog/pkg/schema\"\n)\n\n// ParsePredictor parses Python source and extracts predictor information.\n// predictRef is the class or function name (e.g. \"Predictor\" or \"predict\").\n// mode controls whether we look for predict or train method.\n// sourceDir is the project root for resolving cross-file imports. Pass \"\" if unknown.\nfunc ParsePredictor(source []byte, predictRef string, mode schema.Mode, sourceDir string) (*schema.PredictorInfo, error) {\n\tparser := sitter.NewParser()\n\tparser.SetLanguage(python.GetLanguage())\n\n\ttree, err := parser.ParseCtx(context.Background(), nil, source)\n\tif err != nil {\n\t\treturn nil, schema.WrapError(schema.ErrParse, \"tree-sitter parse failed\", err)\n\t}\n\n\troot := tree.RootNode()\n\n\t// 1. Collect imports\n\timports := collectImports(root, source)\n\n\t// 2. Collect module-level variable assignments\n\tmoduleScope := collectModuleScope(root, source)\n\n\t// 3. Collect BaseModel subclasses (local file first, then cross-file)\n\tmodelClasses := collectModelClasses(root, source, imports)\n\tif sourceDir != \"\" {\n\t\tresolveExternalModels(sourceDir, imports, modelClasses)\n\t}\n\n\t// 4. Collect Input() references from class attributes and static methods\n\tinputRegistry := collectInputRegistry(root, source, imports, moduleScope)\n\n\t// 5. Find the target predict/train function\n\tmethodName := \"predict\"\n\tif mode == schema.ModeTrain {\n\t\tmethodName = \"train\"\n\t}\n\n\tfuncNode, err := findTargetFunction(root, source, predictRef, methodName)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// 6. Check if method (has self first param)\n\tparamsNode := funcNode.ChildByFieldName(\"parameters\")\n\tif paramsNode == nil {\n\t\treturn nil, schema.WrapError(schema.ErrParse, \"function has no parameters node\", nil)\n\t}\n\tisMethod := firstParamIsSelf(paramsNode, source)\n\n\t// 7. Extract parameters\n\tinputs, err := extractInputs(paramsNode, source, methodName, isMethod, imports, inputRegistry, moduleScope)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// 8. Extract return type\n\treturnAnn := funcNode.ChildByFieldName(\"return_type\")\n\tif returnAnn == nil {\n\t\treturn nil, schema.WrapError(schema.ErrMissingReturnType, methodName, nil)\n\t}\n\treturnTypeAnn, err := parseTypeAnnotation(returnAnn, source)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\toutput, err := schema.ResolveSchemaType(returnTypeAnn, imports, modelClasses)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &schema.PredictorInfo{\n\t\tInputs: inputs,\n\t\tOutput: output,\n\t\tMode:   mode,\n\t}, nil\n}\n\n// ---------------------------------------------------------------------------\n// Helpers\n// ---------------------------------------------------------------------------\n\n// namedChildren returns all named children of a node.\nfunc namedChildren(n *sitter.Node) []*sitter.Node {\n\tcount := int(n.NamedChildCount())\n\tresult := make([]*sitter.Node, 0, count)\n\tfor i := range count {\n\t\tresult = append(result, n.NamedChild(i))\n\t}\n\treturn result\n}\n\n// allChildren returns all children (named and anonymous) of a node.\nfunc allChildren(n *sitter.Node) []*sitter.Node {\n\tcount := int(n.ChildCount())\n\tresult := make([]*sitter.Node, 0, count)\n\tfor i := range count {\n\t\tresult = append(result, n.Child(i))\n\t}\n\treturn result\n}\n\n// content returns the source text for a node.\nfunc content(n *sitter.Node, source []byte) string {\n\treturn n.Content(source)\n}\n\n// ---------------------------------------------------------------------------\n// Import collection\n// ---------------------------------------------------------------------------\n\nfunc collectImports(root *sitter.Node, source []byte) *schema.ImportContext {\n\tctx := schema.NewImportContext()\n\n\tfor _, child := range namedChildren(root) {\n\t\tif child.Type() == \"import_from_statement\" {\n\t\t\tparseImportFrom(child, source, ctx)\n\t\t}\n\t}\n\n\t// Always include Python builtins\n\tfor _, builtin := range []string{\"str\", \"int\", \"float\", \"bool\", \"list\", \"dict\", \"set\"} {\n\t\tif _, ok := ctx.Names.Get(builtin); !ok {\n\t\t\tctx.Names.Set(builtin, schema.ImportEntry{Module: \"builtins\", Original: builtin})\n\t\t}\n\t}\n\tif _, ok := ctx.Names.Get(\"None\"); !ok {\n\t\tctx.Names.Set(\"None\", schema.ImportEntry{Module: \"builtins\", Original: \"None\"})\n\t}\n\n\treturn ctx\n}\n\nfunc parseImportFrom(node *sitter.Node, source []byte, ctx *schema.ImportContext) {\n\tmoduleNode := node.ChildByFieldName(\"module_name\")\n\tif moduleNode == nil {\n\t\treturn\n\t}\n\tmodule := content(moduleNode, source)\n\n\tfor _, child := range allChildren(node) {\n\t\tswitch child.Type() {\n\t\tcase \"dotted_name\":\n\t\t\t// Single import: `from X import name`\n\t\t\t// Skip if this is the module_name itself\n\t\t\tif child.StartByte() != moduleNode.StartByte() {\n\t\t\t\tname := content(child, source)\n\t\t\t\tctx.Names.Set(name, schema.ImportEntry{Module: module, Original: name})\n\t\t\t}\n\t\tcase \"aliased_import\":\n\t\t\t// Single aliased import: `from X import name as alias`\n\t\t\torigNode := child.ChildByFieldName(\"name\")\n\t\t\taliasNode := child.ChildByFieldName(\"alias\")\n\t\t\torig := \"\"\n\t\t\tif origNode != nil {\n\t\t\t\torig = content(origNode, source)\n\t\t\t}\n\t\t\talias := orig\n\t\t\tif aliasNode != nil {\n\t\t\t\talias = content(aliasNode, source)\n\t\t\t}\n\t\t\tctx.Names.Set(alias, schema.ImportEntry{Module: module, Original: orig})\n\t\tcase \"import_list\":\n\t\t\tfor _, importChild := range allChildren(child) {\n\t\t\t\tswitch importChild.Type() {\n\t\t\t\tcase \"dotted_name\":\n\t\t\t\t\tname := content(importChild, source)\n\t\t\t\t\tctx.Names.Set(name, schema.ImportEntry{Module: module, Original: name})\n\t\t\t\tcase \"aliased_import\":\n\t\t\t\t\torigNode := importChild.ChildByFieldName(\"name\")\n\t\t\t\t\taliasNode := importChild.ChildByFieldName(\"alias\")\n\t\t\t\t\torig := \"\"\n\t\t\t\t\tif origNode != nil {\n\t\t\t\t\t\torig = content(origNode, source)\n\t\t\t\t\t}\n\t\t\t\t\talias := orig\n\t\t\t\t\tif aliasNode != nil {\n\t\t\t\t\t\talias = content(aliasNode, source)\n\t\t\t\t\t}\n\t\t\t\t\tctx.Names.Set(alias, schema.ImportEntry{Module: module, Original: orig})\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n}\n\n// ---------------------------------------------------------------------------\n// Module scope collection\n// ---------------------------------------------------------------------------\n\ntype moduleScope map[string]schema.DefaultValue\n\nfunc collectModuleScope(root *sitter.Node, source []byte) moduleScope {\n\tscope := make(moduleScope)\n\tfor _, child := range namedChildren(root) {\n\t\tvar assign *sitter.Node\n\t\tif child.Type() == \"expression_statement\" {\n\t\t\tif child.NamedChildCount() == 1 {\n\t\t\t\tinner := child.NamedChild(0)\n\t\t\t\tif inner.Type() == \"assignment\" {\n\t\t\t\t\tassign = inner\n\t\t\t\t}\n\t\t\t}\n\t\t} else if child.Type() == \"assignment\" {\n\t\t\tassign = child\n\t\t}\n\t\tif assign == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tleft := assign.ChildByFieldName(\"left\")\n\t\tif left == nil || left.Type() != \"identifier\" {\n\t\t\tcontinue\n\t\t}\n\t\tname := content(left, source)\n\n\t\tright := assign.ChildByFieldName(\"right\")\n\t\tif right == nil {\n\t\t\tcontinue\n\t\t}\n\t\tif val, ok := parseDefaultValue(right, source); ok {\n\t\t\tscope[name] = val\n\t\t}\n\t}\n\treturn scope\n}\n\n// resolveDefaultExpr tries to resolve an expression to a DefaultValue by\n// literal parsing, then falling back to module scope lookup for identifiers.\nfunc resolveDefaultExpr(node *sitter.Node, source []byte, scope moduleScope) (schema.DefaultValue, bool) {\n\tif val, ok := parseDefaultValue(node, source); ok {\n\t\treturn val, true\n\t}\n\tif node.Type() == \"identifier\" {\n\t\tname := content(node, source)\n\t\tif val, ok := scope[name]; ok {\n\t\t\treturn val, true\n\t\t}\n\t}\n\treturn schema.DefaultValue{}, false\n}\n\n// resolveChoicesExpr tries to statically resolve a choices= expression.\nfunc resolveChoicesExpr(node *sitter.Node, source []byte, scope moduleScope) ([]schema.DefaultValue, bool) {\n\tswitch node.Type() {\n\tcase \"list\":\n\t\treturn parseListLiteral(node, source)\n\n\tcase \"identifier\":\n\t\tname := content(node, source)\n\t\tval, ok := scope[name]\n\t\tif !ok {\n\t\t\treturn nil, false\n\t\t}\n\t\tif val.Kind == schema.DefaultList {\n\t\t\treturn val.List, true\n\t\t}\n\t\treturn nil, false\n\n\tcase \"call\":\n\t\treturn resolveChoicesCall(node, source, scope)\n\n\tcase \"binary_operator\":\n\t\t// Only handle + (list concatenation)\n\t\thasPlus := false\n\t\tfor _, c := range allChildren(node) {\n\t\t\tif !c.IsNamed() && content(c, source) == \"+\" {\n\t\t\t\thasPlus = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tif !hasPlus {\n\t\t\treturn nil, false\n\t\t}\n\t\tleft := node.ChildByFieldName(\"left\")\n\t\tright := node.ChildByFieldName(\"right\")\n\t\tif left == nil || right == nil {\n\t\t\treturn nil, false\n\t\t}\n\t\tleftItems, ok := resolveChoicesExpr(left, source, scope)\n\t\tif !ok {\n\t\t\treturn nil, false\n\t\t}\n\t\trightItems, ok := resolveChoicesExpr(right, source, scope)\n\t\tif !ok {\n\t\t\treturn nil, false\n\t\t}\n\t\treturn append(leftItems, rightItems...), true\n\t}\n\treturn nil, false\n}\n\n// resolveChoicesCall resolves list(X.keys()) or list(X.values()).\nfunc resolveChoicesCall(node *sitter.Node, source []byte, scope moduleScope) ([]schema.DefaultValue, bool) {\n\tfuncNode := node.ChildByFieldName(\"function\")\n\tif funcNode == nil || content(funcNode, source) != \"list\" {\n\t\treturn nil, false\n\t}\n\n\targs := node.ChildByFieldName(\"arguments\")\n\tif args == nil {\n\t\treturn nil, false\n\t}\n\n\t// Find the single positional argument\n\tvar arg *sitter.Node\n\tfor _, c := range namedChildren(args) {\n\t\targ = c\n\t\tbreak\n\t}\n\tif arg == nil || arg.Type() != \"call\" {\n\t\treturn nil, false\n\t}\n\n\tinnerFunc := arg.ChildByFieldName(\"function\")\n\tif innerFunc == nil || innerFunc.Type() != \"attribute\" {\n\t\treturn nil, false\n\t}\n\n\tobj := innerFunc.ChildByFieldName(\"object\")\n\tattr := innerFunc.ChildByFieldName(\"attribute\")\n\tif obj == nil || attr == nil || obj.Type() != \"identifier\" {\n\t\treturn nil, false\n\t}\n\n\tvarName := content(obj, source)\n\tmethodName := content(attr, source)\n\n\tdictVal, ok := scope[varName]\n\tif !ok || dictVal.Kind != schema.DefaultDict {\n\t\treturn nil, false\n\t}\n\n\tswitch methodName {\n\tcase \"keys\":\n\t\treturn dictVal.DictKeys, true\n\tcase \"values\":\n\t\treturn dictVal.DictVals, true\n\t}\n\treturn nil, false\n}\n\n// ---------------------------------------------------------------------------\n// BaseModel subclass collection\n// ---------------------------------------------------------------------------\n\nfunc collectModelClasses(root *sitter.Node, source []byte, imports *schema.ImportContext) schema.ModelClassMap {\n\tmodels := schema.NewOrderedMap[string, []schema.ModelField]()\n\n\tfor _, child := range namedChildren(root) {\n\t\tclassNode := unwrapClass(child)\n\t\tif classNode == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tnameNode := classNode.ChildByFieldName(\"name\")\n\t\tif nameNode == nil {\n\t\t\tcontinue\n\t\t}\n\t\tclassName := content(nameNode, source)\n\n\t\tif !inheritsFromBaseModel(classNode, source, imports) {\n\t\t\tcontinue\n\t\t}\n\n\t\tfields := extractClassAnnotations(classNode, source)\n\t\tmodels.Set(className, fields)\n\t}\n\treturn models\n}\n\n// resolveExternalModels looks at imports that brought in names not yet in\n// modelClasses, attempts to find the corresponding .py file on disk, parses\n// it, and merges any BaseModel subclasses into modelClasses.\n//\n// This handles every local import permutation:\n//\n//\tfrom .types import Output          → /types.py\n//\tfrom types import Output           → /types.py\n//\tfrom models.output import Result   → /models/output.py\n//\tfrom .models.output import Result  → /models/output.py\n//\tfrom my_app.types import Foo       → /my_app/types.py\n//\n// Non-local imports (stdlib, pip packages) are skipped because the file\n// won't exist on disk.\nfunc resolveExternalModels(sourceDir string, imports *schema.ImportContext, models schema.ModelClassMap) {\n\t// Track which modules we've already tried so we don't re-parse.\n\ttried := make(map[string]bool)\n\n\timports.Names.Entries(func(localName string, entry schema.ImportEntry) {\n\t\t// Already resolved locally — skip.\n\t\tif _, ok := models.Get(localName); ok {\n\t\t\treturn\n\t\t}\n\n\t\tmodule := entry.Module\n\t\tif !tried[module] {\n\t\t\ttried[module] = true\n\n\t\t\t// Skip known non-local modules.\n\t\t\tif isKnownExternalModule(module) {\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\t// Convert module path to filesystem path and try to find it.\n\t\t\tpyPath := moduleToFilePath(module)\n\t\t\tif pyPath == \"\" {\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\tfullPath := filepath.Join(sourceDir, pyPath)\n\t\t\tsource, err := os.ReadFile(fullPath)\n\t\t\tif err != nil {\n\t\t\t\tif errors.Is(err, os.ErrNotExist) {\n\t\t\t\t\t// File doesn't exist — it's an external package, not local.\n\t\t\t\t\treturn\n\t\t\t\t}\n\t\t\t\tfmt.Fprintf(os.Stderr, \"cog: warning: failed to read %q: %v\\n\", fullPath, err)\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\t// Parse the file and extract BaseModel subclasses.\n\t\t\tparser := sitter.NewParser()\n\t\t\tparser.SetLanguage(python.GetLanguage())\n\t\t\ttree, err := parser.ParseCtx(context.Background(), nil, source)\n\t\t\tif err != nil {\n\t\t\t\tfmt.Fprintf(os.Stderr, \"cog: warning: failed to parse %q: %v\\n\", fullPath, err)\n\t\t\t\treturn\n\t\t\t}\n\n\t\t\tfileImports := collectImports(tree.RootNode(), source)\n\t\t\tfileModels := collectModelClasses(tree.RootNode(), source, fileImports)\n\n\t\t\t// Merge discovered models into the caller's map.\n\t\t\tfileModels.Entries(func(name string, fields []schema.ModelField) {\n\t\t\t\tif _, exists := models.Get(name); !exists {\n\t\t\t\t\tmodels.Set(name, fields)\n\t\t\t\t}\n\t\t\t})\n\t\t}\n\n\t\t// Handle aliases: \"from X import MyOutput as Output\"\n\t\t// localName is \"Output\", entry.Original is \"MyOutput\".\n\t\t// If we resolved \"MyOutput\" from the file, also register it under \"Output\".\n\t\tif localName != entry.Original {\n\t\t\tif fields, ok := models.Get(entry.Original); ok {\n\t\t\t\tif _, exists := models.Get(localName); !exists {\n\t\t\t\t\tmodels.Set(localName, fields)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t})\n}\n\n// moduleToFilePath converts a Python module path to a relative .py file path.\n//\n//\t\".types\"          → \"types.py\"\n//\t\"types\"           → \"types.py\"\n//\t\".models.output\"  → \"models/output.py\"\n//\t\"models.output\"   → \"models/output.py\"\n//\t\"cog\"             → \"cog.py\"  (will fail os.ReadFile → skipped)\nfunc moduleToFilePath(module string) string {\n\t// Strip leading dots (relative import markers).\n\tclean := strings.TrimLeft(module, \".\")\n\tif clean == \"\" {\n\t\treturn \"\"\n\t}\n\t// Replace dots with path separators.\n\tparts := strings.Split(clean, \".\")\n\treturn filepath.Join(parts...) + \".py\"\n}\n\n// isKnownExternalModule returns true for modules that are definitely not\n// local project files — stdlib, well-known packages, etc.\nfunc isKnownExternalModule(module string) bool {\n\t// Extract the top-level package name.\n\ttop := module\n\tif i := strings.Index(module, \".\"); i > 0 {\n\t\ttop = module[:i]\n\t}\n\ttop = strings.TrimLeft(top, \".\")\n\n\tswitch top {\n\tcase \"builtins\", \"typing\", \"typing_extensions\",\n\t\t\"collections\", \"abc\", \"enum\", \"dataclasses\",\n\t\t\"os\", \"sys\", \"io\", \"json\", \"re\", \"math\", \"pathlib\",\n\t\t\"functools\", \"itertools\", \"contextlib\",\n\t\t\"concurrent\", \"asyncio\", \"multiprocessing\", \"threading\",\n\t\t\"logging\", \"warnings\", \"unittest\", \"pytest\",\n\t\t\"numpy\", \"torch\", \"tensorflow\", \"jax\", \"scipy\", \"sklearn\",\n\t\t\"transformers\", \"diffusers\", \"accelerate\", \"safetensors\",\n\t\t\"PIL\", \"cv2\", \"skimage\",\n\t\t\"requests\", \"httpx\", \"aiohttp\", \"fastapi\", \"flask\",\n\t\t\"pydantic\", \"cog\":\n\t\treturn true\n\t}\n\treturn false\n}\n\nfunc unwrapClass(node *sitter.Node) *sitter.Node {\n\tif node.Type() == \"class_definition\" {\n\t\treturn node\n\t}\n\tif node.Type() == \"decorated_definition\" {\n\t\tfor _, c := range namedChildren(node) {\n\t\t\tif c.Type() == \"class_definition\" {\n\t\t\t\treturn c\n\t\t\t}\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc unwrapFunction(node *sitter.Node) *sitter.Node {\n\tif node.Type() == \"function_definition\" {\n\t\treturn node\n\t}\n\tif node.Type() == \"decorated_definition\" {\n\t\tfor _, c := range namedChildren(node) {\n\t\t\tif c.Type() == \"function_definition\" {\n\t\t\t\treturn c\n\t\t\t}\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc inheritsFromBaseModel(classNode *sitter.Node, source []byte, imports *schema.ImportContext) bool {\n\tsupers := classNode.ChildByFieldName(\"superclasses\")\n\tif supers == nil {\n\t\treturn false\n\t}\n\tfor _, child := range allChildren(supers) {\n\t\tswitch child.Type() {\n\t\tcase \"identifier\":\n\t\t\tname := content(child, source)\n\t\t\tif imports.IsBaseModel(name) || name == \"BaseModel\" {\n\t\t\t\treturn true\n\t\t\t}\n\t\tcase \"attribute\":\n\t\t\t// Handle dotted access: pydantic.BaseModel, cog.BaseModel\n\t\t\ttext := content(child, source)\n\t\t\tif strings.HasSuffix(text, \".BaseModel\") {\n\t\t\t\treturn true\n\t\t\t}\n\t\t}\n\t}\n\treturn false\n}\n\nfunc extractClassAnnotations(classNode *sitter.Node, source []byte) []schema.ModelField {\n\tbody := classNode.ChildByFieldName(\"body\")\n\tif body == nil {\n\t\treturn nil\n\t}\n\n\tvar fields []schema.ModelField\n\tfor _, child := range namedChildren(body) {\n\t\tnode := child\n\t\tif child.Type() == \"expression_statement\" && child.NamedChildCount() == 1 {\n\t\t\tnode = child.NamedChild(0)\n\t\t}\n\n\t\tswitch node.Type() {\n\t\tcase \"assignment\":\n\t\t\tif f, ok := parseAnnotatedAssignment(node, source); ok {\n\t\t\t\tfields = append(fields, f)\n\t\t\t}\n\t\tcase \"type\":\n\t\t\tif f, ok := parseBareAnnotation(node, source); ok {\n\t\t\t\tfields = append(fields, f)\n\t\t\t}\n\t\t}\n\t}\n\treturn fields\n}\n\nfunc parseAnnotatedAssignment(node *sitter.Node, source []byte) (schema.ModelField, bool) {\n\tleft := node.ChildByFieldName(\"left\")\n\ttypeNode := node.ChildByFieldName(\"type\")\n\tif left == nil || typeNode == nil || left.Type() != \"identifier\" {\n\t\treturn schema.ModelField{}, false\n\t}\n\n\tname := content(left, source)\n\ttypeAnn, err := parseTypeAnnotation(typeNode, source)\n\tif err != nil {\n\t\treturn schema.ModelField{}, false\n\t}\n\n\tvar def *schema.DefaultValue\n\tif right := node.ChildByFieldName(\"right\"); right != nil {\n\t\tif v, ok := parseDefaultValue(right, source); ok {\n\t\t\tdef = &v\n\t\t}\n\t}\n\n\treturn schema.ModelField{Name: name, Type: typeAnn, Default: def}, true\n}\n\nfunc parseBareAnnotation(node *sitter.Node, source []byte) (schema.ModelField, bool) {\n\ttext := strings.TrimSpace(content(node, source))\n\tparts := strings.SplitN(text, \":\", 2)\n\tif len(parts) != 2 {\n\t\treturn schema.ModelField{}, false\n\t}\n\tname := strings.TrimSpace(parts[0])\n\ttypeStr := strings.TrimSpace(parts[1])\n\n\tif name == \"\" || (name[0] != '_' && (name[0] < 'a' || name[0] > 'z') && (name[0] < 'A' || name[0] > 'Z')) {\n\t\treturn schema.ModelField{}, false\n\t}\n\n\ttypeAnn, ok := parseTypeFromString(typeStr)\n\tif !ok {\n\t\treturn schema.ModelField{}, false\n\t}\n\n\treturn schema.ModelField{Name: name, Type: typeAnn, Default: nil}, true\n}\n\nfunc parseTypeFromString(s string) (schema.TypeAnnotation, bool) {\n\ts = strings.TrimSpace(s)\n\tif s == \"\" {\n\t\treturn schema.TypeAnnotation{}, false\n\t}\n\n\t// Union: X | Y\n\tif strings.Contains(s, \"|\") {\n\t\tparts := strings.Split(s, \"|\")\n\t\tvar members []schema.TypeAnnotation\n\t\tfor _, p := range parts {\n\t\t\tm, ok := parseTypeFromString(strings.TrimSpace(p))\n\t\t\tif !ok {\n\t\t\t\treturn schema.TypeAnnotation{}, false\n\t\t\t}\n\t\t\tmembers = append(members, m)\n\t\t}\n\t\tif len(members) >= 2 {\n\t\t\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotUnion, Args: members}, true\n\t\t}\n\t\treturn schema.TypeAnnotation{}, false\n\t}\n\n\t// Generic: X[Y] or X[Y, Z]\n\tbracketPos := strings.Index(s, \"[\")\n\tif bracketPos >= 0 && strings.HasSuffix(s, \"]\") {\n\t\touter := strings.TrimSpace(s[:bracketPos])\n\t\tinnerStr := s[bracketPos+1 : len(s)-1]\n\n\t\t// Split on top-level commas (handles Union[str, None], etc.)\n\t\tparts := splitTopLevelCommas(innerStr)\n\t\tvar args []schema.TypeAnnotation\n\t\tfor _, p := range parts {\n\t\t\targ, ok := parseTypeFromString(strings.TrimSpace(p))\n\t\t\tif !ok {\n\t\t\t\treturn schema.TypeAnnotation{}, false\n\t\t\t}\n\t\t\targs = append(args, arg)\n\t\t}\n\t\tif len(args) == 0 {\n\t\t\treturn schema.TypeAnnotation{}, false\n\t\t}\n\t\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotGeneric, Name: outer, Args: args}, true\n\t}\n\n\t// Simple identifier\n\tfor _, c := range s {\n\t\tif (c < 'a' || c > 'z') && (c < 'A' || c > 'Z') && (c < '0' || c > '9') && c != '_' {\n\t\t\treturn schema.TypeAnnotation{}, false\n\t\t}\n\t}\n\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotSimple, Name: s}, true\n}\n\n// splitTopLevelCommas splits a string on commas that are not nested inside brackets.\n// e.g. \"str, None\" → [\"str\", \"None\"], \"List[str], None\" → [\"List[str]\", \"None\"]\nfunc splitTopLevelCommas(s string) []string {\n\tvar parts []string\n\tdepth := 0\n\tstart := 0\n\tfor i, c := range s {\n\t\tswitch c {\n\t\tcase '[':\n\t\t\tdepth++\n\t\tcase ']':\n\t\t\tdepth--\n\t\tcase ',':\n\t\t\tif depth == 0 {\n\t\t\t\tparts = append(parts, s[start:i])\n\t\t\t\tstart = i + 1\n\t\t\t}\n\t\t}\n\t}\n\tparts = append(parts, s[start:])\n\treturn parts\n}\n\n// ---------------------------------------------------------------------------\n// InputRegistry — resolves ClassName.attr and ClassName.method(args)\n// ---------------------------------------------------------------------------\n\ntype inputCallInfo struct {\n\tDefault     *schema.DefaultValue\n\tDescription *string\n\tGE          *float64\n\tLE          *float64\n\tMinLength   *uint64\n\tMaxLength   *uint64\n\tRegex       *string\n\tChoices     []schema.DefaultValue\n\tDeprecated  *bool\n}\n\ntype inputMethodInfo struct {\n\tParamNames []string\n\tBaseInfo   inputCallInfo\n}\n\ntype inputRegistry struct {\n\tAttributes map[string]inputCallInfo\n\tMethods    map[string]inputMethodInfo\n}\n\nfunc newInputRegistry() *inputRegistry {\n\treturn &inputRegistry{\n\t\tAttributes: make(map[string]inputCallInfo),\n\t\tMethods:    make(map[string]inputMethodInfo),\n\t}\n}\n\nfunc collectInputRegistry(root *sitter.Node, source []byte, imports *schema.ImportContext, scope moduleScope) *inputRegistry {\n\tregistry := newInputRegistry()\n\n\tfor _, child := range namedChildren(root) {\n\t\tclassNode := unwrapClass(child)\n\t\tif classNode == nil {\n\t\t\tcontinue\n\t\t}\n\t\tnameNode := classNode.ChildByFieldName(\"name\")\n\t\tif nameNode == nil {\n\t\t\tcontinue\n\t\t}\n\t\tclassName := content(nameNode, source)\n\n\t\tbody := classNode.ChildByFieldName(\"body\")\n\t\tif body == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tfor _, stmt := range namedChildren(body) {\n\t\t\tinner := stmt\n\t\t\tif stmt.Type() == \"expression_statement\" && stmt.NamedChildCount() == 1 {\n\t\t\t\tinner = stmt.NamedChild(0)\n\t\t\t}\n\n\t\t\tif inner.Type() == \"assignment\" {\n\t\t\t\tcollectInputAttribute(className, inner, source, imports, scope, registry)\n\t\t\t}\n\n\t\t\tif funcNode := unwrapFunction(inner); funcNode != nil {\n\t\t\t\tcollectInputMethod(className, funcNode, source, imports, scope, registry)\n\t\t\t}\n\t\t}\n\t}\n\n\treturn registry\n}\n\nfunc collectInputAttribute(className string, assignment *sitter.Node, source []byte, imports *schema.ImportContext, scope moduleScope, registry *inputRegistry) {\n\tleft := assignment.ChildByFieldName(\"left\")\n\tif left == nil || left.Type() != \"identifier\" {\n\t\treturn\n\t}\n\tattrName := content(left, source)\n\n\tright := assignment.ChildByFieldName(\"right\")\n\tif right == nil || !isInputCall(right, source, imports) {\n\t\treturn\n\t}\n\n\tkey := className + \".\" + attrName\n\tinfo, err := parseInputCall(right, source, key, scope)\n\tif err != nil {\n\t\treturn\n\t}\n\tregistry.Attributes[key] = info\n}\n\nfunc collectInputMethod(className string, funcNode *sitter.Node, source []byte, imports *schema.ImportContext, scope moduleScope, registry *inputRegistry) {\n\tnameNode := funcNode.ChildByFieldName(\"name\")\n\tif nameNode == nil {\n\t\treturn\n\t}\n\tmethodName := content(nameNode, source)\n\n\tparams := funcNode.ChildByFieldName(\"parameters\")\n\tif params == nil {\n\t\treturn\n\t}\n\n\tvar paramNames []string\n\tfor _, param := range allChildren(params) {\n\t\tswitch param.Type() {\n\t\tcase \"identifier\":\n\t\t\tname := content(param, source)\n\t\t\tif name != \"self\" && name != \"cls\" {\n\t\t\t\tparamNames = append(paramNames, name)\n\t\t\t}\n\t\tcase \"typed_parameter\":\n\t\t\t// typed_parameter has no \"name\" field; first identifier child is the name\n\t\t\tfor j := 0; j < int(param.NamedChildCount()); j++ {\n\t\t\t\tc := param.NamedChild(j)\n\t\t\t\tif c.Type() == \"identifier\" {\n\t\t\t\t\tname := content(c, source)\n\t\t\t\t\tif name != \"self\" && name != \"cls\" {\n\t\t\t\t\t\tparamNames = append(paramNames, name)\n\t\t\t\t\t}\n\t\t\t\t\tbreak\n\t\t\t\t}\n\t\t\t}\n\t\tcase \"typed_default_parameter\", \"default_parameter\":\n\t\t\tif n := param.ChildByFieldName(\"name\"); n != nil {\n\t\t\t\tname := content(n, source)\n\t\t\t\tif name != \"self\" && name != \"cls\" {\n\t\t\t\t\tparamNames = append(paramNames, name)\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\tbody := funcNode.ChildByFieldName(\"body\")\n\tif body == nil {\n\t\treturn\n\t}\n\n\tinputCall := findReturnInputCall(body, source, imports)\n\tif inputCall == nil {\n\t\treturn\n\t}\n\n\tkey := className + \".\" + methodName\n\tinfo, err := parseInputCall(inputCall, source, key, scope)\n\tif err != nil {\n\t\treturn\n\t}\n\tregistry.Methods[key] = inputMethodInfo{ParamNames: paramNames, BaseInfo: info}\n}\n\nfunc findReturnInputCall(body *sitter.Node, source []byte, imports *schema.ImportContext) *sitter.Node {\n\tfor _, child := range namedChildren(body) {\n\t\tif child.Type() == \"return_statement\" {\n\t\t\tif child.NamedChildCount() > 0 {\n\t\t\t\texpr := child.NamedChild(0)\n\t\t\t\tif isInputCall(expr, source, imports) {\n\t\t\t\t\treturn expr\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc resolveInputReference(node *sitter.Node, source []byte, registry *inputRegistry) (inputCallInfo, bool) {\n\tswitch node.Type() {\n\tcase \"attribute\":\n\t\ttext := content(node, source)\n\t\tinfo, ok := registry.Attributes[text]\n\t\treturn info, ok\n\n\tcase \"call\":\n\t\tfuncNode := node.ChildByFieldName(\"function\")\n\t\tif funcNode == nil || funcNode.Type() != \"attribute\" {\n\t\t\treturn inputCallInfo{}, false\n\t\t}\n\t\tkey := content(funcNode, source)\n\t\tmethodInfo, ok := registry.Methods[key]\n\t\tif !ok {\n\t\t\treturn inputCallInfo{}, false\n\t\t}\n\n\t\tresolved := methodInfo.BaseInfo\n\n\t\targs := node.ChildByFieldName(\"arguments\")\n\t\tif args == nil {\n\t\t\treturn resolved, true\n\t\t}\n\n\t\t// Build param_name -> call-site value map\n\t\targValues := make(map[string]*sitter.Node)\n\t\tpositionalIdx := 0\n\t\tfor _, arg := range namedChildren(args) {\n\t\t\tif arg.Type() == \"keyword_argument\" {\n\t\t\t\tnameNode := arg.ChildByFieldName(\"name\")\n\t\t\t\tvalNode := arg.ChildByFieldName(\"value\")\n\t\t\t\tif nameNode != nil && valNode != nil {\n\t\t\t\t\targValues[content(nameNode, source)] = valNode\n\t\t\t\t}\n\t\t\t} else if positionalIdx < len(methodInfo.ParamNames) {\n\t\t\t\targValues[methodInfo.ParamNames[positionalIdx]] = arg\n\t\t\t\tpositionalIdx++\n\t\t\t}\n\t\t}\n\n\t\t// Override with call-site values\n\t\tfor paramName, callNode := range argValues {\n\t\t\tswitch paramName {\n\t\t\tcase \"default\":\n\t\t\t\tif val, ok := parseDefaultValue(callNode, source); ok {\n\t\t\t\t\tresolved.Default = &val\n\t\t\t\t}\n\t\t\tcase \"description\":\n\t\t\t\tif s, ok := parseStringLiteral(callNode, source); ok {\n\t\t\t\t\tresolved.Description = &s\n\t\t\t\t}\n\t\t\tcase \"ge\":\n\t\t\t\tif n, ok := parseNumberLiteral(callNode, source); ok {\n\t\t\t\t\tresolved.GE = &n\n\t\t\t\t}\n\t\t\tcase \"le\":\n\t\t\t\tif n, ok := parseNumberLiteral(callNode, source); ok {\n\t\t\t\t\tresolved.LE = &n\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\treturn resolved, true\n\t}\n\treturn inputCallInfo{}, false\n}\n\n// ---------------------------------------------------------------------------\n// Target function finding\n// ---------------------------------------------------------------------------\n\nfunc findTargetFunction(root *sitter.Node, source []byte, predictRef, methodName string) (*sitter.Node, error) {\n\t// First: look for a class with this name\n\tfor _, child := range namedChildren(root) {\n\t\tclassNode := unwrapClass(child)\n\t\tif classNode == nil {\n\t\t\tcontinue\n\t\t}\n\t\tnameNode := classNode.ChildByFieldName(\"name\")\n\t\tif nameNode != nil && content(nameNode, source) == predictRef {\n\t\t\treturn findMethodInClass(classNode, source, predictRef, methodName)\n\t\t}\n\t}\n\n\t// Second: look for standalone function\n\tfor _, child := range namedChildren(root) {\n\t\tfuncNode := unwrapFunction(child)\n\t\tif funcNode == nil {\n\t\t\tcontinue\n\t\t}\n\t\tnameNode := funcNode.ChildByFieldName(\"name\")\n\t\tif nameNode != nil {\n\t\t\tname := content(nameNode, source)\n\t\t\tif name == predictRef || name == methodName {\n\t\t\t\treturn funcNode, nil\n\t\t\t}\n\t\t}\n\t}\n\n\treturn nil, schema.WrapError(schema.ErrPredictorNotFound, predictRef, nil)\n}\n\nfunc findMethodInClass(classNode *sitter.Node, source []byte, className, methodName string) (*sitter.Node, error) {\n\tbody := classNode.ChildByFieldName(\"body\")\n\tif body == nil {\n\t\treturn nil, schema.WrapError(schema.ErrParse, fmt.Sprintf(\"class %s has no body\", className), nil)\n\t}\n\n\tfor _, child := range namedChildren(body) {\n\t\tfuncNode := unwrapFunction(child)\n\t\tif funcNode == nil {\n\t\t\tcontinue\n\t\t}\n\t\tnameNode := funcNode.ChildByFieldName(\"name\")\n\t\tif nameNode != nil && content(nameNode, source) == methodName {\n\t\t\treturn funcNode, nil\n\t\t}\n\t}\n\n\treturn nil, schema.WrapError(schema.ErrMethodNotFound, fmt.Sprintf(\"%s.%s not found\", className, methodName), nil)\n}\n\n// ---------------------------------------------------------------------------\n// Parameter extraction\n// ---------------------------------------------------------------------------\n\nfunc firstParamIsSelf(params *sitter.Node, source []byte) bool {\n\tfor _, child := range allChildren(params) {\n\t\tif child.Type() == \"identifier\" {\n\t\t\treturn content(child, source) == \"self\"\n\t\t}\n\t}\n\treturn false\n}\n\nfunc extractInputs(\n\tparamsNode *sitter.Node,\n\tsource []byte,\n\tmethodName string,\n\tskipSelf bool,\n\timports *schema.ImportContext,\n\tregistry *inputRegistry,\n\tscope moduleScope,\n) (*schema.OrderedMap[string, schema.InputField], error) {\n\tinputs := schema.NewOrderedMap[string, schema.InputField]()\n\torder := 0\n\tseenSelf := false\n\n\tfor _, child := range allChildren(paramsNode) {\n\t\tswitch child.Type() {\n\t\tcase \"identifier\":\n\t\t\tif !seenSelf && skipSelf {\n\t\t\t\tname := content(child, source)\n\t\t\t\tif name == \"self\" {\n\t\t\t\t\tseenSelf = true\n\t\t\t\t\tcontinue\n\t\t\t\t}\n\t\t\t}\n\n\t\tcase \"typed_parameter\":\n\t\t\tinput, err := parseTypedParameter(child, source, order, methodName, imports)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tinputs.Set(input.Name, input)\n\t\t\torder++\n\n\t\tcase \"typed_default_parameter\":\n\t\t\tinput, err := parseTypedDefaultParameter(child, source, order, methodName, imports, registry, scope)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tinputs.Set(input.Name, input)\n\t\t\torder++\n\n\t\tcase \"default_parameter\":\n\t\t\tnameNode := child.ChildByFieldName(\"name\")\n\t\t\tparamName := \"\"\n\t\t\tif nameNode != nil {\n\t\t\t\tparamName = content(nameNode, source)\n\t\t\t}\n\t\t\treturn nil, schema.WrapError(schema.ErrMissingTypeAnnotation, fmt.Sprintf(\"parameter '%s' on %s has no type annotation\", paramName, methodName), nil)\n\t\t}\n\t}\n\n\treturn inputs, nil\n}\n\nfunc parseTypedParameter(node *sitter.Node, source []byte, order int, methodName string, imports *schema.ImportContext) (schema.InputField, error) {\n\t// typed_parameter has no \"name\" field in the Python grammar.\n\t// Structure is: identifier \":\" type\n\tvar name string\n\tvar typeNode *sitter.Node\n\tfor i := 0; i < int(node.NamedChildCount()); i++ {\n\t\tc := node.NamedChild(i)\n\t\tswitch c.Type() {\n\t\tcase \"identifier\":\n\t\t\tif name == \"\" {\n\t\t\t\tname = content(c, source)\n\t\t\t}\n\t\tcase \"type\":\n\t\t\ttypeNode = c\n\t\t}\n\t}\n\tif name == \"\" {\n\t\treturn schema.InputField{}, schema.WrapError(schema.ErrParse, \"typed_parameter has no identifier\", nil)\n\t}\n\tif typeNode == nil {\n\t\treturn schema.InputField{}, schema.WrapError(schema.ErrMissingTypeAnnotation, fmt.Sprintf(\"parameter '%s' on %s has no type annotation\", name, methodName), nil)\n\t}\n\n\ttypeAnn, err := parseTypeAnnotation(typeNode, source)\n\tif err != nil {\n\t\treturn schema.InputField{}, err\n\t}\n\tfieldType, err := schema.ResolveFieldType(typeAnn, imports)\n\tif err != nil {\n\t\treturn schema.InputField{}, err\n\t}\n\n\treturn schema.InputField{\n\t\tName:      name,\n\t\tOrder:     order,\n\t\tFieldType: fieldType,\n\t}, nil\n}\n\nfunc parseTypedDefaultParameter(\n\tnode *sitter.Node,\n\tsource []byte,\n\torder int,\n\tmethodName string,\n\timports *schema.ImportContext,\n\tregistry *inputRegistry,\n\tscope moduleScope,\n) (schema.InputField, error) {\n\tnameNode := node.ChildByFieldName(\"name\")\n\tif nameNode == nil {\n\t\treturn schema.InputField{}, schema.WrapError(schema.ErrParse, \"typed_default_parameter has no name\", nil)\n\t}\n\tname := content(nameNode, source)\n\n\ttypeNode := node.ChildByFieldName(\"type\")\n\tif typeNode == nil {\n\t\treturn schema.InputField{}, schema.WrapError(schema.ErrMissingTypeAnnotation, fmt.Sprintf(\"parameter '%s' on %s has no type annotation\", name, methodName), nil)\n\t}\n\n\ttypeAnn, err := parseTypeAnnotation(typeNode, source)\n\tif err != nil {\n\t\treturn schema.InputField{}, err\n\t}\n\tfieldType, err := schema.ResolveFieldType(typeAnn, imports)\n\tif err != nil {\n\t\treturn schema.InputField{}, err\n\t}\n\n\tvalNode := node.ChildByFieldName(\"value\")\n\n\tif valNode != nil {\n\t\t// 1. Direct Input() call\n\t\tif isInputCall(valNode, source, imports) {\n\t\t\tinfo, err := parseInputCall(valNode, source, name, scope)\n\t\t\tif err != nil {\n\t\t\t\treturn schema.InputField{}, err\n\t\t\t}\n\t\t\treturn schema.InputField{\n\t\t\t\tName:        name,\n\t\t\t\tOrder:       order,\n\t\t\t\tFieldType:   fieldType,\n\t\t\t\tDefault:     info.Default,\n\t\t\t\tDescription: info.Description,\n\t\t\t\tGE:          info.GE,\n\t\t\t\tLE:          info.LE,\n\t\t\t\tMinLength:   info.MinLength,\n\t\t\t\tMaxLength:   info.MaxLength,\n\t\t\t\tRegex:       info.Regex,\n\t\t\t\tChoices:     info.Choices,\n\t\t\t\tDeprecated:  info.Deprecated,\n\t\t\t}, nil\n\t\t}\n\n\t\t// 2. Reference to Input() via class attribute or static method\n\t\tif info, ok := resolveInputReference(valNode, source, registry); ok {\n\t\t\treturn schema.InputField{\n\t\t\t\tName:        name,\n\t\t\t\tOrder:       order,\n\t\t\t\tFieldType:   fieldType,\n\t\t\t\tDefault:     info.Default,\n\t\t\t\tDescription: info.Description,\n\t\t\t\tGE:          info.GE,\n\t\t\t\tLE:          info.LE,\n\t\t\t\tMinLength:   info.MinLength,\n\t\t\t\tMaxLength:   info.MaxLength,\n\t\t\t\tRegex:       info.Regex,\n\t\t\t\tChoices:     info.Choices,\n\t\t\t\tDeprecated:  info.Deprecated,\n\t\t\t}, nil\n\t\t}\n\n\t\t// 3. Plain default — must be statically resolvable\n\t\tif def, ok := resolveDefaultExpr(valNode, source, scope); ok {\n\t\t\treturn schema.InputField{\n\t\t\t\tName:      name,\n\t\t\t\tOrder:     order,\n\t\t\t\tFieldType: fieldType,\n\t\t\t\tDefault:   &def,\n\t\t\t}, nil\n\t\t}\n\n\t\t// Can't resolve — hard error\n\t\tvalText := content(valNode, source)\n\t\treturn schema.InputField{}, schema.WrapError(schema.ErrDefaultNotResolvable,\n\t\t\tfmt.Sprintf(\"parameter '%s': default `%s` cannot be statically resolved\", name, valText), nil)\n\t}\n\n\t// No default — required parameter\n\treturn schema.InputField{\n\t\tName:      name,\n\t\tOrder:     order,\n\t\tFieldType: fieldType,\n\t}, nil\n}\n\n// ---------------------------------------------------------------------------\n// Type annotation parsing\n// ---------------------------------------------------------------------------\n\nfunc parseTypeAnnotation(node *sitter.Node, source []byte) (schema.TypeAnnotation, error) {\n\t// Unwrap `type` wrapper node\n\tn := node\n\tif n.Type() == \"type\" && n.NamedChildCount() > 0 {\n\t\tn = n.NamedChild(0)\n\t}\n\n\tswitch n.Type() {\n\tcase \"identifier\":\n\t\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotSimple, Name: content(n, source)}, nil\n\n\tcase \"subscript\":\n\t\tvalue := n.ChildByFieldName(\"value\")\n\t\tif value == nil {\n\t\t\treturn schema.TypeAnnotation{}, schema.WrapError(schema.ErrParse, \"subscript has no value\", nil)\n\t\t}\n\t\touter := content(value, source)\n\n\t\tvar args []schema.TypeAnnotation\n\t\tfor _, child := range namedChildren(n) {\n\t\t\t// Skip the outer identifier (the value field)\n\t\t\tif child.StartByte() == value.StartByte() {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\targ, err := parseTypeAnnotation(child, source)\n\t\t\tif err != nil {\n\t\t\t\treturn schema.TypeAnnotation{}, err\n\t\t\t}\n\t\t\targs = append(args, arg)\n\t\t}\n\n\t\tif len(args) == 0 {\n\t\t\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotSimple, Name: outer}, nil\n\t\t}\n\t\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotGeneric, Name: outer, Args: args}, nil\n\n\tcase \"binary_operator\":\n\t\tleft := n.ChildByFieldName(\"left\")\n\t\tright := n.ChildByFieldName(\"right\")\n\t\tif left == nil || right == nil {\n\t\t\treturn schema.TypeAnnotation{}, schema.WrapError(schema.ErrParse, \"binary_operator missing operand\", nil)\n\t\t}\n\n\t\t// Check that operator is |\n\t\tisUnion := false\n\t\tfor _, c := range allChildren(n) {\n\t\t\tif !c.IsNamed() && content(c, source) == \"|\" {\n\t\t\t\tisUnion = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tif !isUnion {\n\t\t\treturn schema.TypeAnnotation{}, errUnsupported(\"non-union binary operator in type annotation\")\n\t\t}\n\n\t\tleftAnn, err := parseTypeAnnotation(left, source)\n\t\tif err != nil {\n\t\t\treturn schema.TypeAnnotation{}, err\n\t\t}\n\t\trightAnn, err := parseTypeAnnotation(right, source)\n\t\tif err != nil {\n\t\t\treturn schema.TypeAnnotation{}, err\n\t\t}\n\n\t\t// Flatten nested unions\n\t\tvar members []schema.TypeAnnotation\n\t\tif leftAnn.Kind == schema.TypeAnnotUnion {\n\t\t\tmembers = append(members, leftAnn.Args...)\n\t\t} else {\n\t\t\tmembers = append(members, leftAnn)\n\t\t}\n\t\tif rightAnn.Kind == schema.TypeAnnotUnion {\n\t\t\tmembers = append(members, rightAnn.Args...)\n\t\t} else {\n\t\t\tmembers = append(members, rightAnn)\n\t\t}\n\n\t\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotUnion, Args: members}, nil\n\n\tcase \"none\":\n\t\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotSimple, Name: \"None\"}, nil\n\n\tcase \"attribute\":\n\t\treturn schema.TypeAnnotation{Kind: schema.TypeAnnotSimple, Name: content(n, source)}, nil\n\n\tcase \"string\", \"concatenated_string\":\n\t\ttext := content(n, source)\n\t\tinner := strings.TrimLeft(text, \"\\\"'\")\n\t\tinner = strings.TrimRight(inner, \"\\\"'\")\n\t\tif ann, ok := parseTypeFromString(inner); ok {\n\t\t\treturn ann, nil\n\t\t}\n\t\treturn schema.TypeAnnotation{}, errUnsupported(fmt.Sprintf(\"string annotation: %s\", text))\n\n\tdefault:\n\t\ttext := content(n, source)\n\t\tif ann, ok := parseTypeFromString(text); ok {\n\t\t\treturn ann, nil\n\t\t}\n\t\treturn schema.TypeAnnotation{}, errUnsupported(fmt.Sprintf(\"%s: %s\", n.Type(), text))\n\t}\n}\n\nfunc errUnsupported(msg string) error {\n\treturn &schema.SchemaError{Kind: schema.ErrUnsupportedType, Message: fmt.Sprintf(\"unsupported type: %s\", msg)}\n}\n\n// ---------------------------------------------------------------------------\n// Input() call parsing\n// ---------------------------------------------------------------------------\n\nfunc isInputCall(node *sitter.Node, source []byte, imports *schema.ImportContext) bool {\n\tif node.Type() != \"call\" {\n\t\treturn false\n\t}\n\tfuncNode := node.ChildByFieldName(\"function\")\n\tif funcNode == nil {\n\t\treturn false\n\t}\n\tname := content(funcNode, source)\n\tif name == \"Input\" {\n\t\treturn true\n\t}\n\tif e, ok := imports.Names.Get(name); ok {\n\t\treturn e.Module == \"cog\" && e.Original == \"Input\"\n\t}\n\treturn false\n}\n\nfunc parseInputCall(node *sitter.Node, source []byte, paramName string, scope moduleScope) (inputCallInfo, error) {\n\tvar info inputCallInfo\n\n\targs := node.ChildByFieldName(\"arguments\")\n\tif args == nil {\n\t\treturn info, nil\n\t}\n\n\tfor _, child := range namedChildren(args) {\n\t\tif child.Type() != \"keyword_argument\" {\n\t\t\tcontinue\n\t\t}\n\t\tkeyNode := child.ChildByFieldName(\"name\")\n\t\tvalNode := child.ChildByFieldName(\"value\")\n\t\tif keyNode == nil || valNode == nil {\n\t\t\tcontinue\n\t\t}\n\n\t\tkey := content(keyNode, source)\n\t\tswitch key {\n\t\tcase \"default\":\n\t\t\tval, ok := resolveDefaultExpr(valNode, source, scope)\n\t\t\tif !ok {\n\t\t\t\tnone := schema.DefaultValue{Kind: schema.DefaultNone}\n\t\t\t\tval = none\n\t\t\t}\n\t\t\tinfo.Default = &val\n\t\tcase \"default_factory\":\n\t\t\treturn inputCallInfo{}, schema.WrapError(schema.ErrDefaultFactoryNotSupported, fmt.Sprintf(\"parameter '%s': default_factory is not supported in static schema generation\", paramName), nil)\n\t\tcase \"description\":\n\t\t\tif s, ok := parseStringLiteral(valNode, source); ok {\n\t\t\t\tinfo.Description = &s\n\t\t\t}\n\t\tcase \"ge\":\n\t\t\tif n, ok := parseNumberLiteral(valNode, source); ok {\n\t\t\t\tinfo.GE = &n\n\t\t\t}\n\t\tcase \"le\":\n\t\t\tif n, ok := parseNumberLiteral(valNode, source); ok {\n\t\t\t\tinfo.LE = &n\n\t\t\t}\n\t\tcase \"min_length\":\n\t\t\tif n, ok := parseNumberLiteral(valNode, source); ok {\n\t\t\t\tu := uint64(n)\n\t\t\t\tinfo.MinLength = &u\n\t\t\t}\n\t\tcase \"max_length\":\n\t\t\tif n, ok := parseNumberLiteral(valNode, source); ok {\n\t\t\t\tu := uint64(n)\n\t\t\t\tinfo.MaxLength = &u\n\t\t\t}\n\t\tcase \"regex\":\n\t\t\tif s, ok := parseStringLiteral(valNode, source); ok {\n\t\t\t\tinfo.Regex = &s\n\t\t\t}\n\t\tcase \"choices\":\n\t\t\tif items, ok := parseListLiteral(valNode, source); ok {\n\t\t\t\tinfo.Choices = items\n\t\t\t} else if items, ok := resolveChoicesExpr(valNode, source, scope); ok {\n\t\t\t\tinfo.Choices = items\n\t\t\t} else {\n\t\t\t\treturn inputCallInfo{}, schema.WrapError(schema.ErrChoicesNotResolvable, fmt.Sprintf(\"parameter '%s': choices expression cannot be statically resolved\", paramName), nil)\n\t\t\t}\n\t\tcase \"deprecated\":\n\t\t\tif b, ok := parseBoolLiteral(valNode, source); ok {\n\t\t\t\tinfo.Deprecated = &b\n\t\t\t}\n\t\t}\n\t}\n\n\treturn info, nil\n}\n\n// ---------------------------------------------------------------------------\n// Literal parsing\n// ---------------------------------------------------------------------------\n\nfunc parseDefaultValue(node *sitter.Node, source []byte) (schema.DefaultValue, bool) {\n\tswitch node.Type() {\n\tcase \"none\":\n\t\treturn schema.DefaultValue{Kind: schema.DefaultNone}, true\n\tcase \"true\":\n\t\treturn schema.DefaultValue{Kind: schema.DefaultBool, Bool: true}, true\n\tcase \"false\":\n\t\treturn schema.DefaultValue{Kind: schema.DefaultBool, Bool: false}, true\n\tcase \"integer\":\n\t\ttext := content(node, source)\n\t\tn, err := strconv.ParseInt(text, 0, 64)\n\t\tif err != nil {\n\t\t\treturn schema.DefaultValue{}, false\n\t\t}\n\t\treturn schema.DefaultValue{Kind: schema.DefaultInt, Int: n}, true\n\tcase \"float\":\n\t\ttext := content(node, source)\n\t\tf, err := strconv.ParseFloat(text, 64)\n\t\tif err != nil {\n\t\t\treturn schema.DefaultValue{}, false\n\t\t}\n\t\treturn schema.DefaultValue{Kind: schema.DefaultFloat, Float: f}, true\n\tcase \"string\", \"concatenated_string\":\n\t\ts, ok := parseStringLiteral(node, source)\n\t\tif !ok {\n\t\t\treturn schema.DefaultValue{}, false\n\t\t}\n\t\treturn schema.DefaultValue{Kind: schema.DefaultString, Str: s}, true\n\tcase \"list\":\n\t\titems, ok := parseListLiteral(node, source)\n\t\tif !ok {\n\t\t\treturn schema.DefaultValue{}, false\n\t\t}\n\t\treturn schema.DefaultValue{Kind: schema.DefaultList, List: items}, true\n\tcase \"dictionary\":\n\t\tkeys, vals, ok := parseDictLiteral(node, source)\n\t\tif !ok {\n\t\t\treturn schema.DefaultValue{}, false\n\t\t}\n\t\treturn schema.DefaultValue{Kind: schema.DefaultDict, DictKeys: keys, DictVals: vals}, true\n\tcase \"set\":\n\t\titems, ok := parseSetLiteral(node, source)\n\t\tif !ok {\n\t\t\treturn schema.DefaultValue{}, false\n\t\t}\n\t\treturn schema.DefaultValue{Kind: schema.DefaultSet, List: items}, true\n\tcase \"unary_operator\":\n\t\ttext := strings.TrimSpace(content(node, source))\n\t\tif n, err := strconv.ParseInt(text, 0, 64); err == nil {\n\t\t\treturn schema.DefaultValue{Kind: schema.DefaultInt, Int: n}, true\n\t\t}\n\t\tif f, err := strconv.ParseFloat(text, 64); err == nil {\n\t\t\treturn schema.DefaultValue{Kind: schema.DefaultFloat, Float: f}, true\n\t\t}\n\t\treturn schema.DefaultValue{}, false\n\tcase \"tuple\":\n\t\tvar items []schema.DefaultValue\n\t\tfor _, child := range namedChildren(node) {\n\t\t\tif val, ok := parseDefaultValue(child, source); ok {\n\t\t\t\titems = append(items, val)\n\t\t\t}\n\t\t}\n\t\treturn schema.DefaultValue{Kind: schema.DefaultList, List: items}, true\n\t}\n\treturn schema.DefaultValue{}, false\n}\n\nfunc parseStringLiteral(node *sitter.Node, source []byte) (string, bool) {\n\ttext := content(node, source)\n\tif strings.HasPrefix(text, `\"\"\"`) || strings.HasPrefix(text, `'''`) {\n\t\tif len(text) >= 6 {\n\t\t\treturn text[3 : len(text)-3], true\n\t\t}\n\t\treturn \"\", false\n\t}\n\tif strings.HasPrefix(text, `\"`) || strings.HasPrefix(text, `'`) {\n\t\tif len(text) >= 2 {\n\t\t\treturn text[1 : len(text)-1], true\n\t\t}\n\t\treturn \"\", false\n\t}\n\tif strings.HasPrefix(text, `r\"`) || strings.HasPrefix(text, `r'`) {\n\t\tif len(text) >= 3 {\n\t\t\treturn text[2 : len(text)-1], true\n\t\t}\n\t\treturn \"\", false\n\t}\n\treturn \"\", false\n}\n\nfunc parseNumberLiteral(node *sitter.Node, source []byte) (float64, bool) {\n\ttext := strings.TrimSpace(content(node, source))\n\tf, err := strconv.ParseFloat(text, 64)\n\tif err != nil {\n\t\treturn 0, false\n\t}\n\treturn f, true\n}\n\nfunc parseBoolLiteral(node *sitter.Node, source []byte) (bool, bool) {\n\tswitch node.Type() {\n\tcase \"true\":\n\t\treturn true, true\n\tcase \"false\":\n\t\treturn false, true\n\t}\n\ttext := content(node, source)\n\tswitch text {\n\tcase \"True\":\n\t\treturn true, true\n\tcase \"False\":\n\t\treturn false, true\n\t}\n\treturn false, false\n}\n\nfunc parseListLiteral(node *sitter.Node, source []byte) ([]schema.DefaultValue, bool) {\n\tif node.Type() != \"list\" {\n\t\treturn nil, false\n\t}\n\tvar items []schema.DefaultValue\n\tfor _, child := range namedChildren(node) {\n\t\tval, ok := parseDefaultValue(child, source)\n\t\tif !ok {\n\t\t\treturn nil, false\n\t\t}\n\t\titems = append(items, val)\n\t}\n\treturn items, true\n}\n\nfunc parseDictLiteral(node *sitter.Node, source []byte) ([]schema.DefaultValue, []schema.DefaultValue, bool) {\n\tif node.Type() != \"dictionary\" {\n\t\treturn nil, nil, false\n\t}\n\tvar keys, vals []schema.DefaultValue\n\tfor _, child := range namedChildren(node) {\n\t\tif child.Type() == \"pair\" {\n\t\t\tkeyNode := child.ChildByFieldName(\"key\")\n\t\t\tvalNode := child.ChildByFieldName(\"value\")\n\t\t\tif keyNode == nil || valNode == nil {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t\tk, ok1 := parseDefaultValue(keyNode, source)\n\t\t\tv, ok2 := parseDefaultValue(valNode, source)\n\t\t\tif ok1 && ok2 {\n\t\t\t\tkeys = append(keys, k)\n\t\t\t\tvals = append(vals, v)\n\t\t\t}\n\t\t}\n\t}\n\treturn keys, vals, true\n}\n\nfunc parseSetLiteral(node *sitter.Node, source []byte) ([]schema.DefaultValue, bool) {\n\tif node.Type() != \"set\" {\n\t\treturn nil, false\n\t}\n\tvar items []schema.DefaultValue\n\tfor _, child := range namedChildren(node) {\n\t\tval, ok := parseDefaultValue(child, source)\n\t\tif !ok {\n\t\t\treturn nil, false\n\t\t}\n\t\titems = append(items, val)\n\t}\n\treturn items, true\n}\n"
  },
  {
    "path": "pkg/schema/python/parser_fuzz_test.go",
    "content": "package python\n\nimport (\n\t\"testing\"\n\n\tschema \"github.com/replicate/cog/pkg/schema\"\n)\n\n// FuzzParsePredictor feeds arbitrary bytes as Python source to the parser.\n// The parser should never panic regardless of input — it may return errors.\nfunc FuzzParsePredictor(f *testing.F) {\n\t// Seed corpus: valid and invalid Python snippets.\n\tf.Add([]byte(`\nfrom cog import BasePredictor\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> str:\n        return x\n`), \"Predictor\", uint8(0))\n\n\tf.Add([]byte(`\nfrom cog import BasePredictor\nfrom pydantic import BaseModel\nclass Output(BaseModel):\n    text: str\n    score: float = 0.0\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Output:\n        pass\n`), \"Predictor\", uint8(0))\n\n\tf.Add([]byte(`\nfrom cog import BasePredictor\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> dict[str, list[int]]:\n        return {}\n`), \"Predictor\", uint8(0))\n\n\tf.Add([]byte(`\nfrom cog import BasePredictor, ConcatenateIterator\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> ConcatenateIterator[str]:\n        yield x\n`), \"Predictor\", uint8(0))\n\n\t// Training mode.\n\tf.Add([]byte(`\nfrom cog import BasePredictor\nclass Predictor(BasePredictor):\n    def train(self, x: str) -> str:\n        return x\n`), \"Predictor\", uint8(1))\n\n\t// No predictor class at all.\n\tf.Add([]byte(`print(\"hello\")`), \"Predictor\", uint8(0))\n\n\t// Empty source.\n\tf.Add([]byte{}, \"Predictor\", uint8(0))\n\n\t// Garbage bytes.\n\tf.Add([]byte{0xff, 0xfe, 0x00, 0x01, 0x80, 0x90}, \"Predictor\", uint8(0))\n\n\tf.Fuzz(func(t *testing.T, source []byte, predictRef string, modeRaw uint8) {\n\t\tmode := schema.ModePredict\n\t\tif modeRaw%2 == 1 {\n\t\t\tmode = schema.ModeTrain\n\t\t}\n\n\t\t// Must not panic regardless of input.\n\t\t_, _ = ParsePredictor(source, predictRef, mode, \"\")\n\t})\n}\n\n// FuzzParseTypeAnnotation exercises the type annotation parser with\n// arbitrary annotation strings embedded in a predict signature.\nfunc FuzzParseTypeAnnotation(f *testing.F) {\n\ttypes := []string{\n\t\t\"str\", \"int\", \"float\", \"bool\", \"Path\",\n\t\t\"dict\", \"dict[str, int]\", \"dict[str, list[str]]\",\n\t\t\"list[str]\", \"list[dict[str, float]]\",\n\t\t\"Optional[str]\", \"Optional[dict[str, int]]\",\n\t\t\"Iterator[str]\", \"ConcatenateIterator[str]\",\n\t\t\"dict[str, dict[str, dict[str, int]]]\",\n\t\t\"Any\", \"None\", \"list\",\n\t}\n\tfor _, typ := range types {\n\t\tf.Add(typ)\n\t}\n\n\tf.Fuzz(func(t *testing.T, typeName string) {\n\t\t// Build a minimal predict.py with the fuzzed return type.\n\t\tsource := []byte(\"from cog import BasePredictor\\nfrom typing import *\\nclass Predictor(BasePredictor):\\n    def predict(self, x: str) -> \" + typeName + \":\\n        pass\\n\")\n\t\t// Must not panic.\n\t\t_, _ = ParsePredictor(source, \"Predictor\", schema.ModePredict, \"\")\n\t})\n}\n"
  },
  {
    "path": "pkg/schema/python/parser_test.go",
    "content": "package python\n\nimport (\n\t\"errors\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/schema\"\n)\n\n// helper that parses in predict mode and fails on error.\nfunc parse(t *testing.T, source, predictRef string) *schema.PredictorInfo {\n\tt.Helper()\n\tinfo, err := ParsePredictor([]byte(source), predictRef, schema.ModePredict, \"\")\n\trequire.NoError(t, err)\n\treturn info\n}\n\n// helper to parse and expect an error.\nfunc parseErr(t *testing.T, source, predictRef string, mode schema.Mode) *schema.SchemaError {\n\tt.Helper()\n\t_, err := ParsePredictor([]byte(source), predictRef, mode, \"\")\n\trequire.Error(t, err)\n\tvar se *schema.SchemaError\n\trequire.True(t, errors.As(err, &se), \"expected *schema.SchemaError, got %T: %v\", err, err)\n\treturn se\n}\n\n// ---------------------------------------------------------------------------\n// Basic predictor tests\n// ---------------------------------------------------------------------------\n\nfunc TestSimpleStringPredictor(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        return \"hello \" + s\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, 1, info.Inputs.Len())\n\n\ts, ok := info.Inputs.Get(\"s\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, s.FieldType.Primitive)\n\trequire.Equal(t, schema.Required, s.FieldType.Repetition)\n\trequire.Nil(t, s.Default)\n\trequire.True(t, s.IsRequired())\n\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Kind)\n\trequire.Equal(t, schema.TypeString, info.Output.Primitive)\n}\n\nfunc TestMultipleInputsWithDefaults(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input, Path\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        image: Path = Input(description=\"Grayscale input image\"),\n        scale: float = Input(description=\"Factor to scale image by\", ge=0, le=10, default=1.5),\n    ) -> Path:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, 2, info.Inputs.Len())\n\n\timage, ok := info.Inputs.Get(\"image\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypePath, image.FieldType.Primitive)\n\trequire.Nil(t, image.Default)\n\trequire.NotNil(t, image.Description)\n\trequire.Equal(t, \"Grayscale input image\", *image.Description)\n\trequire.True(t, image.IsRequired())\n\n\tscale, ok := info.Inputs.Get(\"scale\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeFloat, scale.FieldType.Primitive)\n\trequire.NotNil(t, scale.Default)\n\trequire.Equal(t, schema.DefaultFloat, scale.Default.Kind)\n\trequire.Equal(t, 1.5, scale.Default.Float)\n\trequire.NotNil(t, scale.GE)\n\trequire.Equal(t, 0.0, *scale.GE)\n\trequire.NotNil(t, scale.LE)\n\trequire.Equal(t, 10.0, *scale.LE)\n\trequire.False(t, scale.IsRequired())\n}\n\n// ---------------------------------------------------------------------------\n// Optional / union inputs\n// ---------------------------------------------------------------------------\n\nfunc TestOptionalInputPipeNone(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input, Path\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        test_image: Path | None = Input(description=\"Test image\", default=None),\n    ) -> Path:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\timg, ok := info.Inputs.Get(\"test_image\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.Optional, img.FieldType.Repetition)\n\trequire.Equal(t, schema.TypePath, img.FieldType.Primitive)\n\trequire.NotNil(t, img.Default)\n\trequire.Equal(t, schema.DefaultNone, img.Default.Kind)\n}\n\nfunc TestOptionalInputTyping(t *testing.T) {\n\tsource := `\nfrom typing import Optional\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        name: Optional[str] = Input(default=None),\n    ) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tname, ok := info.Inputs.Get(\"name\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.Optional, name.FieldType.Repetition)\n\trequire.Equal(t, schema.TypeString, name.FieldType.Primitive)\n}\n\n// ---------------------------------------------------------------------------\n// List inputs\n// ---------------------------------------------------------------------------\n\nfunc TestListInput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, paths: list[str] = Input(description=\"Paths\")) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tpaths, ok := info.Inputs.Get(\"paths\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.Repeated, paths.FieldType.Repetition)\n\trequire.Equal(t, schema.TypeString, paths.FieldType.Primitive)\n}\n\n// ---------------------------------------------------------------------------\n// Choices\n// ---------------------------------------------------------------------------\n\nfunc TestChoicesLiteralList(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, color: str = Input(choices=[\"red\", \"green\", \"blue\"])) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tcolor, ok := info.Inputs.Get(\"color\")\n\trequire.True(t, ok)\n\trequire.Len(t, color.Choices, 3)\n\trequire.Equal(t, \"red\", color.Choices[0].Str)\n\trequire.Equal(t, \"green\", color.Choices[1].Str)\n\trequire.Equal(t, \"blue\", color.Choices[2].Str)\n}\n\nfunc TestChoicesModuleLevelListVar(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nMY_CHOICES = [\"x\", \"y\", \"z\"]\n\nclass Predictor(BasePredictor):\n    def predict(self, v: str = Input(choices=MY_CHOICES)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tv, ok := info.Inputs.Get(\"v\")\n\trequire.True(t, ok)\n\trequire.Len(t, v.Choices, 3)\n\trequire.Equal(t, \"x\", v.Choices[0].Str)\n\trequire.Equal(t, \"y\", v.Choices[1].Str)\n\trequire.Equal(t, \"z\", v.Choices[2].Str)\n}\n\nfunc TestChoicesListDictKeys(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nASPECT_RATIOS = {\n    \"1:1\": (1024, 1024),\n    \"16:9\": (1344, 768),\n    \"2:3\": (832, 1216),\n}\n\nclass Predictor(BasePredictor):\n    def predict(self, ar: str = Input(choices=list(ASPECT_RATIOS.keys()), default=\"1:1\")) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tar, ok := info.Inputs.Get(\"ar\")\n\trequire.True(t, ok)\n\trequire.Len(t, ar.Choices, 3)\n\trequire.Equal(t, \"1:1\", ar.Choices[0].Str)\n\trequire.Equal(t, \"16:9\", ar.Choices[1].Str)\n\trequire.Equal(t, \"2:3\", ar.Choices[2].Str)\n}\n\nfunc TestChoicesListDictValues(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nLABELS = {\"fast\": \"Fast Mode\", \"slow\": \"Slow Mode\"}\n\nclass Predictor(BasePredictor):\n    def predict(self, m: str = Input(choices=list(LABELS.values()))) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tm, ok := info.Inputs.Get(\"m\")\n\trequire.True(t, ok)\n\trequire.Len(t, m.Choices, 2)\n\trequire.Equal(t, \"Fast Mode\", m.Choices[0].Str)\n\trequire.Equal(t, \"Slow Mode\", m.Choices[1].Str)\n}\n\nfunc TestChoicesDictKeysPlusLiteral(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nSIZES = {\"small\": 256, \"large\": 1024}\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str = Input(choices=list(SIZES.keys()) + [\"custom\"])) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\ts, ok := info.Inputs.Get(\"s\")\n\trequire.True(t, ok)\n\trequire.Len(t, s.Choices, 3)\n\trequire.Equal(t, \"small\", s.Choices[0].Str)\n\trequire.Equal(t, \"large\", s.Choices[1].Str)\n\trequire.Equal(t, \"custom\", s.Choices[2].Str)\n}\n\nfunc TestChoicesIntegerDictKeys(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nSTEP_LABELS = {1: \"one\", 2: \"two\", 4: \"four\"}\n\nclass Predictor(BasePredictor):\n    def predict(self, steps: int = Input(choices=list(STEP_LABELS.keys()), default=1)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tsteps, ok := info.Inputs.Get(\"steps\")\n\trequire.True(t, ok)\n\trequire.Len(t, steps.Choices, 3)\n\trequire.Equal(t, schema.DefaultInt, steps.Choices[0].Kind)\n\trequire.Equal(t, int64(1), steps.Choices[0].Int)\n\trequire.Equal(t, int64(2), steps.Choices[1].Int)\n\trequire.Equal(t, int64(4), steps.Choices[2].Int)\n}\n\nfunc TestChoicesConcatTwoVars(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nBASE = [\"a\", \"b\"]\nEXTRA = [\"c\"]\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str = Input(choices=BASE + EXTRA)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tx, ok := info.Inputs.Get(\"x\")\n\trequire.True(t, ok)\n\trequire.Len(t, x.Choices, 3)\n\trequire.Equal(t, \"a\", x.Choices[0].Str)\n\trequire.Equal(t, \"b\", x.Choices[1].Str)\n\trequire.Equal(t, \"c\", x.Choices[2].Str)\n}\n\n// ---------------------------------------------------------------------------\n// Choices error cases\n// ---------------------------------------------------------------------------\n\nfunc TestChoicesVarNotAListErrors(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nNOT_A_LIST = \"oops\"\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str = Input(choices=NOT_A_LIST)) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrChoicesNotResolvable, se.Kind)\n}\n\nfunc TestChoicesUndefinedVarErrors(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str = Input(choices=DOES_NOT_EXIST)) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrChoicesNotResolvable, se.Kind)\n}\n\nfunc TestChoicesArbitraryCallErrors(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str = Input(choices=get_choices())) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrChoicesNotResolvable, se.Kind)\n}\n\nfunc TestChoicesListComprehensionErrors(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str = Input(choices=[f\"{i}x\" for i in range(5)])) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrChoicesNotResolvable, se.Kind)\n}\n\nfunc TestChoicesErrorIncludesParamName(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, my_param: str = Input(choices=some_func())) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Contains(t, se.Message, \"my_param\")\n}\n\nfunc TestChoicesNestedVarNotInScope(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\ndef helper():\n    NESTED = [\"a\", \"b\"]\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str = Input(choices=NESTED)) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrChoicesNotResolvable, se.Kind)\n}\n\n// ---------------------------------------------------------------------------\n// Standalone function\n// ---------------------------------------------------------------------------\n\nfunc TestStandaloneFunction(t *testing.T) {\n\tsource := `\nfrom cog import Input\n\ndef predict(text: str = Input(default=\"world\")) -> str:\n    return f\"hello {text}\"\n`\n\tinfo := parse(t, source, \"predict\")\n\trequire.Equal(t, 1, info.Inputs.Len())\n\n\ttext, ok := info.Inputs.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, text.Default)\n\trequire.Equal(t, schema.DefaultString, text.Default.Kind)\n\trequire.Equal(t, \"world\", text.Default.Str)\n}\n\n// ---------------------------------------------------------------------------\n// Output types\n// ---------------------------------------------------------------------------\n\nfunc TestIteratorOutput(t *testing.T) {\n\tsource := `\nfrom typing import Iterator\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, count: int) -> Iterator[str]:\n        for i in range(count):\n            yield f\"chunk {i}\"\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaIterator, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Elem)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Elem.Kind)\n\trequire.Equal(t, schema.TypeString, info.Output.Elem.Primitive)\n}\n\nfunc TestConcatenateIteratorOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, ConcatenateIterator\n\nclass Predictor(BasePredictor):\n    def predict(self, prompt: str) -> ConcatenateIterator[str]:\n        yield \"hello \"\n        yield \"world\"\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaConcatIterator, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Elem)\n\trequire.Equal(t, schema.TypeString, info.Output.Elem.Primitive)\n}\n\nfunc TestConcatenateIteratorNotStrErrors(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, ConcatenateIterator\n\nclass Predictor(BasePredictor):\n    def predict(self, n: int) -> ConcatenateIterator[int]:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrConcatIteratorNotStr, se.Kind)\n}\n\nfunc TestListOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Path\n\nclass Predictor(BasePredictor):\n    def predict(self, n: int) -> list[Path]:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaArray, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Items)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Items.Kind)\n\trequire.Equal(t, schema.TypePath, info.Output.Items.Primitive)\n}\n\nfunc TestBaseModelOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, BaseModel\n\nclass ModelOutput(BaseModel):\n    text: str\n    score: float\n\nclass Predictor(BasePredictor):\n    def predict(self, prompt: str) -> ModelOutput:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Fields)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\ttext, ok := info.Output.Fields.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaPrimitive, text.Type.Kind)\n\trequire.Equal(t, schema.TypeString, text.Type.Primitive)\n\n\tscore, ok := info.Output.Fields.Get(\"score\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaPrimitive, score.Type.Kind)\n\trequire.Equal(t, schema.TypeFloat, score.Type.Primitive)\n}\n\nfunc TestOptionalOutputErrors(t *testing.T) {\n\tsource := `\nfrom typing import Optional\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> Optional[str]:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrOptionalOutput, se.Kind)\n}\n\nfunc TestOptionalOutputPipeNoneErrors(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str | None:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrOptionalOutput, se.Kind)\n}\n\n// ---------------------------------------------------------------------------\n// Train mode\n// ---------------------------------------------------------------------------\n\nfunc TestTrainMode(t *testing.T) {\n\tsource := `\nfrom cog import Input, Path\n\ndef train(n: int) -> Path:\n    pass\n`\n\tinfo, err := ParsePredictor([]byte(source), \"train\", schema.ModeTrain, \"\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, schema.ModeTrain, info.Mode)\n\trequire.Equal(t, 1, info.Inputs.Len())\n}\n\n// ---------------------------------------------------------------------------\n// Non-BasePredictor class (just has predict method)\n// ---------------------------------------------------------------------------\n\nfunc TestNonBasePredictor(t *testing.T) {\n\tsource := `\nfrom cog import Input\n\nclass Predictor:\n    def predict(self, text: str = Input(default=\"hello\")) -> str:\n        return f\"hello {text}\"\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, 1, info.Inputs.Len())\n\ttext, ok := info.Inputs.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, text.Default)\n\trequire.Equal(t, \"hello\", text.Default.Str)\n}\n\n// ---------------------------------------------------------------------------\n// default_factory hard error\n// ---------------------------------------------------------------------------\n\nfunc TestDefaultFactoryError(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, items: list[str] = Input(default_factory=list)) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrDefaultFactoryNotSupported, se.Kind)\n}\n\n// ---------------------------------------------------------------------------\n// Module-scope default resolution\n// ---------------------------------------------------------------------------\n\nfunc TestDefaultModuleLevelStringInInput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nDEFAULT_RATIO = \"1:1\"\n\nclass Predictor(BasePredictor):\n    def predict(self, ar: str = Input(default=DEFAULT_RATIO)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tar, ok := info.Inputs.Get(\"ar\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, ar.Default)\n\trequire.Equal(t, schema.DefaultString, ar.Default.Kind)\n\trequire.Equal(t, \"1:1\", ar.Default.Str)\n}\n\nfunc TestDefaultModuleLevelIntInInput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nDEFAULT_STEPS = 50\n\nclass Predictor(BasePredictor):\n    def predict(self, steps: int = Input(default=DEFAULT_STEPS)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tsteps, ok := info.Inputs.Get(\"steps\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, steps.Default)\n\trequire.Equal(t, schema.DefaultInt, steps.Default.Kind)\n\trequire.Equal(t, int64(50), steps.Default.Int)\n}\n\nfunc TestDefaultModuleLevelListInInput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nDEFAULT_TAGS = [\"a\", \"b\"]\n\nclass Predictor(BasePredictor):\n    def predict(self, tags: list[str] = Input(default=DEFAULT_TAGS)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\ttags, ok := info.Inputs.Get(\"tags\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, tags.Default)\n\trequire.Equal(t, schema.DefaultList, tags.Default.Kind)\n\trequire.Len(t, tags.Default.List, 2)\n\trequire.Equal(t, \"a\", tags.Default.List[0].Str)\n\trequire.Equal(t, \"b\", tags.Default.List[1].Str)\n}\n\nfunc TestDefaultModuleLevelVarPlain(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nMY_DEFAULT = \"hello\"\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str = MY_DEFAULT) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\ttext, ok := info.Inputs.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, text.Default)\n\trequire.Equal(t, schema.DefaultString, text.Default.Kind)\n\trequire.Equal(t, \"hello\", text.Default.Str)\n}\n\nfunc TestDefaultUndefinedVarPlainErrors(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str = UNDEFINED_VAR) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Contains(t, se.Message, \"cannot be statically resolved\")\n}\n\n// ---------------------------------------------------------------------------\n// InputRegistry — class attribute reference\n// ---------------------------------------------------------------------------\n\nfunc TestInputRegistryAttribute(t *testing.T) {\n\tsource := `\nfrom dataclasses import dataclass\nfrom cog import BasePredictor, Input\n\nRATIOS = {\"1:1\": (1024, 1024), \"16:9\": (1344, 768)}\n\n@dataclass(frozen=True)\nclass Inputs:\n    ar = Input(description=\"Aspect ratio\", choices=list(RATIOS.keys()), default=\"1:1\")\n\nclass Predictor(BasePredictor):\n    def predict(self, ar: str = Inputs.ar) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tar, ok := info.Inputs.Get(\"ar\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, ar.Description)\n\trequire.Equal(t, \"Aspect ratio\", *ar.Description)\n\trequire.Len(t, ar.Choices, 2)\n\trequire.Equal(t, \"1:1\", ar.Choices[0].Str)\n\trequire.Equal(t, \"16:9\", ar.Choices[1].Str)\n\trequire.NotNil(t, ar.Default)\n\trequire.Equal(t, \"1:1\", ar.Default.Str)\n}\n\n// ---------------------------------------------------------------------------\n// InputRegistry — static method reference\n// ---------------------------------------------------------------------------\n\nfunc TestInputRegistryMethod(t *testing.T) {\n\tsource := `\nfrom dataclasses import dataclass\nfrom cog import BasePredictor, Input\n\n@dataclass(frozen=True)\nclass Inputs:\n    @staticmethod\n    def guidance(default: float) -> Input:\n        return Input(description=\"Guidance scale\", ge=0.0, le=20.0, default=default)\n\nclass Predictor(BasePredictor):\n    def predict(self, guidance_scale: float = Inputs.guidance(7.5)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tgs, ok := info.Inputs.Get(\"guidance_scale\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, gs.Description)\n\trequire.Equal(t, \"Guidance scale\", *gs.Description)\n\trequire.NotNil(t, gs.GE)\n\trequire.Equal(t, 0.0, *gs.GE)\n\trequire.NotNil(t, gs.LE)\n\trequire.Equal(t, 20.0, *gs.LE)\n\trequire.NotNil(t, gs.Default)\n\trequire.Equal(t, schema.DefaultFloat, gs.Default.Kind)\n\trequire.Equal(t, 7.5, gs.Default.Float)\n}\n\n// ---------------------------------------------------------------------------\n// Error cases: missing annotations, predictor not found, etc.\n// ---------------------------------------------------------------------------\n\nfunc TestPredictorNotFound(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Other(BasePredictor):\n    def predict(self, s: str) -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrPredictorNotFound, se.Kind)\n}\n\nfunc TestMethodNotFound(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def setup(self):\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrMethodNotFound, se.Kind)\n}\n\nfunc TestMissingReturnType(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str):\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrMissingReturnType, se.Kind)\n}\n\nfunc TestMissingTypeAnnotation(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s=\"hello\") -> str:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrMissingTypeAnnotation, se.Kind)\n}\n\n// ---------------------------------------------------------------------------\n// All input types\n// ---------------------------------------------------------------------------\n\nfunc TestAllPrimitiveInputTypes(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\tpyType   string\n\t\texpected schema.PrimitiveType\n\t}{\n\t\t{\"str\", \"str\", schema.TypeString},\n\t\t{\"int\", \"int\", schema.TypeInteger},\n\t\t{\"float\", \"float\", schema.TypeFloat},\n\t\t{\"bool\", \"bool\", schema.TypeBool},\n\t\t{\"Path\", \"Path\", schema.TypePath},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tsource := `\nfrom cog import BasePredictor, Path\n\nclass Predictor(BasePredictor):\n    def predict(self, x: ` + tt.pyType + `) -> str:\n        pass\n`\n\t\t\tinfo := parse(t, source, \"Predictor\")\n\t\t\tx, ok := info.Inputs.Get(\"x\")\n\t\t\trequire.True(t, ok)\n\t\t\trequire.Equal(t, tt.expected, x.FieldType.Primitive)\n\t\t\trequire.Equal(t, schema.Required, x.FieldType.Repetition)\n\t\t})\n\t}\n}\n\n// ---------------------------------------------------------------------------\n// Input() with constraints\n// ---------------------------------------------------------------------------\n\nfunc TestInputConstraints(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(\n        self,\n        text: str = Input(\n            description=\"Input text\",\n            min_length=1,\n            max_length=100,\n            regex=\"^[a-z]+$\",\n        ),\n    ) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\ttext, ok := info.Inputs.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, text.Description)\n\trequire.Equal(t, \"Input text\", *text.Description)\n\trequire.NotNil(t, text.MinLength)\n\trequire.Equal(t, uint64(1), *text.MinLength)\n\trequire.NotNil(t, text.MaxLength)\n\trequire.Equal(t, uint64(100), *text.MaxLength)\n\trequire.NotNil(t, text.Regex)\n\trequire.Equal(t, \"^[a-z]+$\", *text.Regex)\n}\n\n// ---------------------------------------------------------------------------\n// Negative numbers and booleans as defaults\n// ---------------------------------------------------------------------------\n\nfunc TestNegativeNumberDefault(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, temp: float = Input(default=-1.5)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\ttemp, ok := info.Inputs.Get(\"temp\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, temp.Default)\n\trequire.Equal(t, schema.DefaultFloat, temp.Default.Kind)\n\trequire.Equal(t, -1.5, temp.Default.Float)\n}\n\nfunc TestBoolDefault(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, flag: bool = Input(default=True)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tflag, ok := info.Inputs.Get(\"flag\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, flag.Default)\n\trequire.Equal(t, schema.DefaultBool, flag.Default.Kind)\n\trequire.True(t, flag.Default.Bool)\n}\n\n// ---------------------------------------------------------------------------\n// Parameter ordering\n// ---------------------------------------------------------------------------\n\nfunc TestParameterOrdering(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, alpha: str, beta: int, gamma: float = Input(default=1.0)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, 3, info.Inputs.Len())\n\n\t// Check insertion order\n\tkeys := info.Inputs.Keys()\n\trequire.Equal(t, \"alpha\", keys[0])\n\trequire.Equal(t, \"beta\", keys[1])\n\trequire.Equal(t, \"gamma\", keys[2])\n\n\talpha, _ := info.Inputs.Get(\"alpha\")\n\trequire.Equal(t, 0, alpha.Order)\n\tbeta, _ := info.Inputs.Get(\"beta\")\n\trequire.Equal(t, 1, beta.Order)\n\tgamma, _ := info.Inputs.Get(\"gamma\")\n\trequire.Equal(t, 2, gamma.Order)\n}\n\n// ---------------------------------------------------------------------------\n// Deprecated flag\n// ---------------------------------------------------------------------------\n\nfunc TestDeprecatedInput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, old_param: str = Input(deprecated=True, default=\"x\")) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\told, ok := info.Inputs.Get(\"old_param\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, old.Deprecated)\n\trequire.True(t, *old.Deprecated)\n}\n\n// ---------------------------------------------------------------------------\n// File type (deprecated alias for Path)\n// ---------------------------------------------------------------------------\n\nfunc TestFileType(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, File\n\nclass Predictor(BasePredictor):\n    def predict(self, f: File) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tf, ok := info.Inputs.Get(\"f\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeFile, f.FieldType.Primitive)\n}\n\n// ---------------------------------------------------------------------------\n// Secret type\n// ---------------------------------------------------------------------------\n\nfunc TestSecretType(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Secret\n\nclass Predictor(BasePredictor):\n    def predict(self, token: Secret) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\ttoken, ok := info.Inputs.Get(\"token\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeSecret, token.FieldType.Primitive)\n}\n\n// ---------------------------------------------------------------------------\n// Multiple classes — finds the right one\n// ---------------------------------------------------------------------------\n\nfunc TestMultipleClassesFindsTarget(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, BaseModel\n\nclass Output(BaseModel):\n    text: str\n\nclass Helper:\n    pass\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, 1, info.Inputs.Len())\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Kind)\n}\n\n// ---------------------------------------------------------------------------\n// BaseModel with defaults\n// ---------------------------------------------------------------------------\n\nfunc TestBaseModelOutputWithDefaults(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, BaseModel\n\nclass Result(BaseModel):\n    text: str\n    confidence: float = 0.0\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\n\tconf, ok := info.Output.Fields.Get(\"confidence\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, conf.Default)\n\trequire.Equal(t, schema.DefaultFloat, conf.Default.Kind)\n\trequire.Equal(t, 0.0, conf.Default.Float)\n}\n\n// ---------------------------------------------------------------------------\n// Pydantic BaseModel output\n// ---------------------------------------------------------------------------\n\nfunc TestPydanticBaseModelOutput(t *testing.T) {\n\tsource := `\nfrom pydantic import BaseModel as PydanticBaseModel\nfrom cog import BasePredictor\n\nclass Result(PydanticBaseModel):\n    name: str\n    score: float\n    tags: list[str]\n\nclass Predictor(BasePredictor):\n    def predict(self, name: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Fields)\n\trequire.Equal(t, 3, info.Output.Fields.Len())\n\n\tname, ok := info.Output.Fields.Get(\"name\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaPrimitive, name.Type.Kind)\n\trequire.Equal(t, schema.TypeString, name.Type.Primitive)\n\n\tscore, ok := info.Output.Fields.Get(\"score\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaPrimitive, score.Type.Kind)\n\trequire.Equal(t, schema.TypeFloat, score.Type.Primitive)\n}\n\nfunc TestPydanticBaseModelDottedOutput(t *testing.T) {\n\tsource := `\nimport pydantic\nfrom cog import BasePredictor\n\nclass Result(pydantic.BaseModel):\n    text: str\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\n\ttext, ok := info.Output.Fields.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaPrimitive, text.Type.Kind)\n\trequire.Equal(t, schema.TypeString, text.Type.Primitive)\n}\n\nfunc TestPydanticBaseModelDirectImport(t *testing.T) {\n\tsource := `\nfrom pydantic import BaseModel\nfrom cog import BasePredictor\n\nclass Output(BaseModel):\n    value: int\n\nclass Predictor(BasePredictor):\n    def predict(self, x: int) -> Output:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\n\tval, ok := info.Output.Fields.Get(\"value\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaPrimitive, val.Type.Kind)\n\trequire.Equal(t, schema.TypeInteger, val.Type.Primitive)\n}\n\n// ---------------------------------------------------------------------------\n// Unparameterized dict/list output (opaque JSON)\n// ---------------------------------------------------------------------------\n\nfunc TestDictOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input, Path\n\nclass Predictor(BasePredictor):\n    def predict(self, image: Path = Input(description=\"Image\")) -> dict:\n        return {\"class\": \"hotdog\", \"score\": 0.95}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaAny, info.Output.Kind)\n}\n\nfunc TestParameterizedDictOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str = Input(description=\"Text\")) -> dict[str, dict[str, str]]:\n        return {\"inputs\": {\"text\": text}}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\t// dict[str, dict[str, str]] → SchemaDict with nested SchemaDict value type\n\trequire.Equal(t, schema.SchemaDict, info.Output.Kind)\n\trequire.NotNil(t, info.Output.ValueType)\n\trequire.Equal(t, schema.SchemaDict, info.Output.ValueType.Kind)\n\trequire.NotNil(t, info.Output.ValueType.ValueType)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.ValueType.ValueType.Kind)\n\trequire.Equal(t, schema.TypeString, info.Output.ValueType.ValueType.Primitive)\n}\n\nfunc TestBareListOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, s: str) -> list:\n        return [1, 2, 3]\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaArray, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Items)\n\trequire.Equal(t, schema.SchemaAny, info.Output.Items.Kind)\n}\n\n// ---------------------------------------------------------------------------\n// No-input predictor (only self)\n// ---------------------------------------------------------------------------\n\nfunc TestNoInputPredictor(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self) -> str:\n        return \"hello\"\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, 0, info.Inputs.Len())\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Kind)\n}\n\n// ---------------------------------------------------------------------------\n// Falsy defaults (False, 0, 0.0, \"\")\n// ---------------------------------------------------------------------------\n\nfunc TestDefaultFalse(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, flag: bool = False) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tf, ok := info.Inputs.Get(\"flag\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, f.Default)\n\trequire.Equal(t, schema.DefaultBool, f.Default.Kind)\n\trequire.Equal(t, false, f.Default.Bool)\n\trequire.False(t, f.IsRequired())\n}\n\nfunc TestDefaultZeroInt(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, count: int = 0) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tc, ok := info.Inputs.Get(\"count\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, c.Default)\n\trequire.Equal(t, schema.DefaultInt, c.Default.Kind)\n\trequire.Equal(t, int64(0), c.Default.Int)\n\trequire.False(t, c.IsRequired())\n}\n\nfunc TestDefaultZeroFloat(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, weight: float = 0.0) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tw, ok := info.Inputs.Get(\"weight\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, w.Default)\n\trequire.Equal(t, schema.DefaultFloat, w.Default.Kind)\n\trequire.Equal(t, 0.0, w.Default.Float)\n\trequire.False(t, w.IsRequired())\n}\n\nfunc TestDefaultEmptyString(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: str = \"\") -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\ttext, ok := info.Inputs.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, text.Default)\n\trequire.Equal(t, schema.DefaultString, text.Default.Kind)\n\trequire.Equal(t, \"\", text.Default.Str)\n\trequire.False(t, text.IsRequired())\n}\n\nfunc TestDefaultNegativeInt(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, offset: int = -1) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\to, ok := info.Inputs.Get(\"offset\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, o.Default)\n\trequire.Equal(t, schema.DefaultInt, o.Default.Kind)\n\trequire.Equal(t, int64(-1), o.Default.Int)\n}\n\n// ---------------------------------------------------------------------------\n// Async iterators\n// ---------------------------------------------------------------------------\n\nfunc TestAsyncIteratorOutput(t *testing.T) {\n\tsource := `\nfrom typing import AsyncIterator\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    async def predict(self, s: str) -> AsyncIterator[str]:\n        yield s\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaIterator, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Elem)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Elem.Kind)\n\trequire.Equal(t, schema.TypeString, info.Output.Elem.Primitive)\n}\n\nfunc TestAsyncConcatenateIteratorOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, ConcatenateIterator\n\nclass Predictor(BasePredictor):\n    async def predict(self, s: str) -> ConcatenateIterator[str]:\n        yield s\n`\n\t// Note: AsyncConcatenateIterator is also valid via typing import,\n\t// but ConcatenateIterator in async context works the same way\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaConcatIterator, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Elem)\n\trequire.Equal(t, schema.TypeString, info.Output.Elem.Primitive)\n}\n\n// ---------------------------------------------------------------------------\n// typing.List and typing.Union syntax\n// ---------------------------------------------------------------------------\n\nfunc TestTypingListCapitalL(t *testing.T) {\n\tsource := `\nfrom typing import List\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, items: List[str]) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\titems, ok := info.Inputs.Get(\"items\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, items.FieldType.Primitive)\n\trequire.Equal(t, schema.Repeated, items.FieldType.Repetition)\n}\n\nfunc TestTypingUnionStrNone(t *testing.T) {\n\tsource := `\nfrom typing import Union\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, text: Union[str, None] = None) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\ttext, ok := info.Inputs.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, text.FieldType.Primitive)\n\trequire.Equal(t, schema.Optional, text.FieldType.Repetition)\n\trequire.False(t, text.IsRequired())\n}\n\n// ---------------------------------------------------------------------------\n// All-optional inputs (no required array)\n// ---------------------------------------------------------------------------\n\nfunc TestAllOptionalInputs(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Input\n\nclass Predictor(BasePredictor):\n    def predict(self, a: str = \"x\", b: int = Input(default=5)) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, 2, info.Inputs.Len())\n\n\ta, ok := info.Inputs.Get(\"a\")\n\trequire.True(t, ok)\n\trequire.False(t, a.IsRequired())\n\n\tb, ok := info.Inputs.Get(\"b\")\n\trequire.True(t, ok)\n\trequire.False(t, b.IsRequired())\n}\n\n// ---------------------------------------------------------------------------\n// list[Path] as input\n// ---------------------------------------------------------------------------\n\nfunc TestListPathInput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor, Path\n\nclass Predictor(BasePredictor):\n    def predict(self, files: list[Path]) -> str:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\tfiles, ok := info.Inputs.Get(\"files\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypePath, files.FieldType.Primitive)\n\trequire.Equal(t, schema.Repeated, files.FieldType.Repetition)\n}\n\n// ---------------------------------------------------------------------------\n// Recursive / nested output types\n// ---------------------------------------------------------------------------\n\nfunc TestDictStrStrOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> dict[str, str]:\n        return {\"key\": \"value\"}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaDict, info.Output.Kind)\n\trequire.NotNil(t, info.Output.ValueType)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.ValueType.Kind)\n\trequire.Equal(t, schema.TypeString, info.Output.ValueType.Primitive)\n}\n\nfunc TestDictStrIntOutput(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> dict[str, int]:\n        return {\"count\": 42}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaDict, info.Output.Kind)\n\trequire.NotNil(t, info.Output.ValueType)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.ValueType.Kind)\n\trequire.Equal(t, schema.TypeInteger, info.Output.ValueType.Primitive)\n}\n\nfunc TestNestedDictOutput(t *testing.T) {\n\t// dict[str, dict[str, str]]\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> dict[str, dict[str, str]]:\n        return {\"outer\": {\"inner\": \"value\"}}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaDict, info.Output.Kind)\n\trequire.NotNil(t, info.Output.ValueType)\n\trequire.Equal(t, schema.SchemaDict, info.Output.ValueType.Kind)\n\trequire.NotNil(t, info.Output.ValueType.ValueType)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.ValueType.ValueType.Kind)\n\trequire.Equal(t, schema.TypeString, info.Output.ValueType.ValueType.Primitive)\n}\n\nfunc TestDictOfListOutput(t *testing.T) {\n\t// dict[str, list[int]]\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> dict[str, list[int]]:\n        return {\"numbers\": [1, 2, 3]}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaDict, info.Output.Kind)\n\trequire.NotNil(t, info.Output.ValueType)\n\trequire.Equal(t, schema.SchemaArray, info.Output.ValueType.Kind)\n\trequire.NotNil(t, info.Output.ValueType.Items)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.ValueType.Items.Kind)\n\trequire.Equal(t, schema.TypeInteger, info.Output.ValueType.Items.Primitive)\n}\n\nfunc TestListOfDictOutput(t *testing.T) {\n\t// list[dict[str, str]]\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> list[dict[str, str]]:\n        return [{\"key\": \"value\"}]\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaArray, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Items)\n\trequire.Equal(t, schema.SchemaDict, info.Output.Items.Kind)\n\trequire.NotNil(t, info.Output.Items.ValueType)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Items.ValueType.Kind)\n\trequire.Equal(t, schema.TypeString, info.Output.Items.ValueType.Primitive)\n}\n\nfunc TestListOfListOutput(t *testing.T) {\n\t// list[list[float]]\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> list[list[float]]:\n        return [[1.0, 2.0], [3.0]]\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaArray, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Items)\n\trequire.Equal(t, schema.SchemaArray, info.Output.Items.Kind)\n\trequire.NotNil(t, info.Output.Items.Items)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Items.Items.Kind)\n\trequire.Equal(t, schema.TypeFloat, info.Output.Items.Items.Primitive)\n}\n\nfunc TestTripleNestedDictOutput(t *testing.T) {\n\t// dict[str, dict[str, dict[str, int]]]\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> dict[str, dict[str, dict[str, int]]]:\n        return {\"a\": {\"b\": {\"c\": 1}}}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaDict, info.Output.Kind)\n\n\tlevel2 := info.Output.ValueType\n\trequire.NotNil(t, level2)\n\trequire.Equal(t, schema.SchemaDict, level2.Kind)\n\n\tlevel3 := level2.ValueType\n\trequire.NotNil(t, level3)\n\trequire.Equal(t, schema.SchemaDict, level3.Kind)\n\n\tleaf := level3.ValueType\n\trequire.NotNil(t, leaf)\n\trequire.Equal(t, schema.SchemaPrimitive, leaf.Kind)\n\trequire.Equal(t, schema.TypeInteger, leaf.Primitive)\n}\n\nfunc TestListOfDictOfListOutput(t *testing.T) {\n\t// list[dict[str, list[str]]]\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> list[dict[str, list[str]]]:\n        return [{\"tags\": [\"a\", \"b\"]}]\n`\n\tinfo := parse(t, source, \"Predictor\")\n\t// list[...]\n\trequire.Equal(t, schema.SchemaArray, info.Output.Kind)\n\t// dict[str, ...]\n\tdictType := info.Output.Items\n\trequire.NotNil(t, dictType)\n\trequire.Equal(t, schema.SchemaDict, dictType.Kind)\n\t// list[str]\n\tinnerList := dictType.ValueType\n\trequire.NotNil(t, innerList)\n\trequire.Equal(t, schema.SchemaArray, innerList.Kind)\n\t// str\n\trequire.NotNil(t, innerList.Items)\n\trequire.Equal(t, schema.SchemaPrimitive, innerList.Items.Kind)\n\trequire.Equal(t, schema.TypeString, innerList.Items.Primitive)\n}\n\nfunc TestIteratorOfDictOutput(t *testing.T) {\n\t// Iterator[dict[str, str]] — iterator yielding dicts\n\tsource := `\nfrom typing import Iterator\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Iterator[dict[str, str]]:\n        yield {\"key\": \"value\"}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaIterator, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Elem)\n\trequire.Equal(t, schema.SchemaDict, info.Output.Elem.Kind)\n\trequire.NotNil(t, info.Output.Elem.ValueType)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Elem.ValueType.Kind)\n\trequire.Equal(t, schema.TypeString, info.Output.Elem.ValueType.Primitive)\n}\n\nfunc TestIteratorOfListOutput(t *testing.T) {\n\t// Iterator[list[int]] — iterator yielding lists\n\tsource := `\nfrom typing import Iterator\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Iterator[list[int]]:\n        yield [1, 2, 3]\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaIterator, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Elem)\n\trequire.Equal(t, schema.SchemaArray, info.Output.Elem.Kind)\n\trequire.NotNil(t, info.Output.Elem.Items)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Elem.Items.Kind)\n\trequire.Equal(t, schema.TypeInteger, info.Output.Elem.Items.Primitive)\n}\n\nfunc TestDictOfPathOutput(t *testing.T) {\n\t// dict[str, Path] — dict with file URIs as values\n\tsource := `\nfrom cog import BasePredictor, Path\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> dict[str, Path]:\n        return {\"file\": Path(\"output.png\")}\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaDict, info.Output.Kind)\n\trequire.NotNil(t, info.Output.ValueType)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.ValueType.Kind)\n\trequire.Equal(t, schema.TypePath, info.Output.ValueType.Primitive)\n}\n\nfunc TestListOfPathOutput(t *testing.T) {\n\t// list[Path]\n\tsource := `\nfrom cog import BasePredictor, Path\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> list[Path]:\n        return [Path(\"a.png\"), Path(\"b.png\")]\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaArray, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Items)\n\trequire.Equal(t, schema.SchemaPrimitive, info.Output.Items.Kind)\n\trequire.Equal(t, schema.TypePath, info.Output.Items.Primitive)\n}\n\n// ---------------------------------------------------------------------------\n// Unresolvable output type errors\n// ---------------------------------------------------------------------------\n\nfunc TestUnresolvableImportedTypeError(t *testing.T) {\n\tsource := `\nfrom some_random_package import WeirdType\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> WeirdType:\n        return None\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrUnresolvableType, se.Kind)\n\trequire.Contains(t, se.Message, \"WeirdType\")\n\trequire.Contains(t, se.Message, \"some_random_package\")\n\trequire.Contains(t, se.Message, \".pyi stub\")\n}\n\nfunc TestUnresolvableUndefinedTypeError(t *testing.T) {\n\tsource := `\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> MysteryType:\n        return None\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrUnresolvableType, se.Kind)\n\trequire.Contains(t, se.Message, \"MysteryType\")\n\trequire.Contains(t, se.Message, \"not a primitive type\")\n\trequire.Contains(t, se.Message, \"BaseModel\")\n}\n\nfunc TestUnresolvableDottedImportTypeError(t *testing.T) {\n\tsource := `\nfrom transformers import AutoTokenizer\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> AutoTokenizer:\n        return None\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrUnresolvableType, se.Kind)\n\trequire.Contains(t, se.Message, \"AutoTokenizer\")\n\trequire.Contains(t, se.Message, \"transformers\")\n}\n\nfunc TestUnresolvableTypeTorchTensor(t *testing.T) {\n\tsource := `\nfrom torch import Tensor\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Tensor:\n        return None\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrUnresolvableType, se.Kind)\n\trequire.Contains(t, se.Message, \"Tensor\")\n\trequire.Contains(t, se.Message, \"torch\")\n}\n\nfunc TestUnresolvableTypeNumpyArray(t *testing.T) {\n\tsource := `\nfrom numpy import ndarray\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> ndarray:\n        return None\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrUnresolvableType, se.Kind)\n\trequire.Contains(t, se.Message, \"ndarray\")\n\trequire.Contains(t, se.Message, \"numpy\")\n}\n\nfunc TestDictWithUnresolvableValueTypeErrors(t *testing.T) {\n\t// Regression: dict[str, Tensor] used to silently collapse to SchemaAny.\n\t// Now it propagates the error from the value type resolution.\n\tsource := `\nfrom torch import Tensor\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> dict[str, Tensor]:\n        return {}\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrUnresolvableType, se.Kind)\n\trequire.Contains(t, se.Message, \"Tensor\")\n}\n\nfunc TestModelFieldDictWithUnresolvableValueTypeErrors(t *testing.T) {\n\t// Same bug but inside a BaseModel field.\n\tsource := `\nfrom torch import Tensor\nfrom pydantic import BaseModel\nfrom cog import BasePredictor\n\nclass Result(BaseModel):\n    tensors: dict[str, Tensor]\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`\n\tse := parseErr(t, source, \"Predictor\", schema.ModePredict)\n\trequire.Equal(t, schema.ErrUnresolvableType, se.Kind)\n\trequire.Contains(t, se.Message, \"Tensor\")\n}\n\n// ---------------------------------------------------------------------------\n// Pydantic output still works after migration\n// ---------------------------------------------------------------------------\n\nfunc TestPydanticV1CompatOutput(t *testing.T) {\n\tsource := `\nfrom pydantic.v1 import BaseModel\nfrom cog import BasePredictor\n\nclass Result(BaseModel):\n    text: str\n    score: float\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.NotNil(t, info.Output.Fields)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\ttext, ok := info.Output.Fields.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaPrimitive, text.Type.Kind)\n\trequire.Equal(t, schema.TypeString, text.Type.Primitive)\n\n\tscore, ok := info.Output.Fields.Get(\"score\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaPrimitive, score.Type.Kind)\n\trequire.Equal(t, schema.TypeFloat, score.Type.Primitive)\n}\n\nfunc TestPydanticOutputWithOptionalField(t *testing.T) {\n\tsource := `\nfrom pydantic import BaseModel\nfrom typing import Optional\nfrom cog import BasePredictor\n\nclass Result(BaseModel):\n    text: str\n    error: Optional[str] = None\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\ttext, ok := info.Output.Fields.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.True(t, text.Required)\n\n\terrField, ok := info.Output.Fields.Get(\"error\")\n\trequire.True(t, ok)\n\trequire.True(t, errField.Type.Nullable)\n}\n\nfunc TestPydanticOutputDefaultedFieldNotNullable(t *testing.T) {\n\t// Regression: a field with a default but NOT Optional must NOT be nullable.\n\t// Previously !Required was incorrectly mapped to nullable in JSON Schema.\n\tsource := `\nfrom pydantic import BaseModel\nfrom cog import BasePredictor\n\nclass Result(BaseModel):\n    text: str\n    debug: bool = False\n    count: int = 0\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\n\t// text: required, not nullable\n\ttext, ok := info.Output.Fields.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.True(t, text.Required)\n\trequire.False(t, text.Type.Nullable)\n\n\t// debug: has default so not required, but NOT nullable (not Optional)\n\tdebug, ok := info.Output.Fields.Get(\"debug\")\n\trequire.True(t, ok)\n\trequire.False(t, debug.Required, \"defaulted field should not be required\")\n\trequire.False(t, debug.Type.Nullable, \"non-Optional defaulted field must not be nullable\")\n\n\t// count: same — defaulted, not nullable\n\tcount, ok := info.Output.Fields.Get(\"count\")\n\trequire.True(t, ok)\n\trequire.False(t, count.Required)\n\trequire.False(t, count.Type.Nullable)\n\n\t// Verify JSON Schema output doesn't include \"nullable\" for these fields\n\tjs := info.Output.JSONSchema()\n\tprops, ok := js[\"properties\"].(map[string]any)\n\trequire.True(t, ok)\n\tdebugProp, ok := props[\"debug\"].(map[string]any)\n\trequire.True(t, ok)\n\t_, hasNullable := debugProp[\"nullable\"]\n\trequire.False(t, hasNullable, \"JSON Schema for defaulted non-Optional field must not have nullable\")\n}\n\nfunc TestPydanticOutputOptionalFieldNullable(t *testing.T) {\n\tsource := `\nfrom pydantic import BaseModel\nfrom typing import Optional\nfrom cog import BasePredictor\n\nclass Result(BaseModel):\n    text: str\n    error: Optional[str] = None\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\n\terrField, ok := info.Output.Fields.Get(\"error\")\n\trequire.True(t, ok)\n\trequire.True(t, errField.Type.Nullable, \"Optional field should be nullable\")\n\trequire.False(t, errField.Required, \"Optional field with default should not be required\")\n\n\t// Verify JSON Schema output includes \"nullable\" for Optional field\n\tjs := info.Output.JSONSchema()\n\tprops, ok := js[\"properties\"].(map[string]any)\n\trequire.True(t, ok)\n\terrProp, ok := props[\"error\"].(map[string]any)\n\trequire.True(t, ok)\n\trequire.Equal(t, true, errProp[\"nullable\"])\n}\n\nfunc TestPydanticOutputWithListField(t *testing.T) {\n\tsource := `\nfrom pydantic import BaseModel\nfrom cog import BasePredictor\n\nclass Result(BaseModel):\n    tags: list[str]\n    scores: list[float]\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\ttags, ok := info.Output.Fields.Get(\"tags\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaArray, tags.Type.Kind)\n\trequire.NotNil(t, tags.Type.Items)\n\trequire.Equal(t, schema.TypeString, tags.Type.Items.Primitive)\n}\n\nfunc TestPydanticOutputWithDictField(t *testing.T) {\n\tsource := `\nfrom pydantic import BaseModel\nfrom cog import BasePredictor\n\nclass Result(BaseModel):\n    metadata: dict[str, int]\n    nested: dict[str, list[str]]\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\t// metadata: dict[str, int]\n\tmetadata, ok := info.Output.Fields.Get(\"metadata\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaDict, metadata.Type.Kind)\n\trequire.NotNil(t, metadata.Type.ValueType)\n\trequire.Equal(t, schema.SchemaPrimitive, metadata.Type.ValueType.Kind)\n\trequire.Equal(t, schema.TypeInteger, metadata.Type.ValueType.Primitive)\n\n\t// nested: dict[str, list[str]]\n\tnested, ok := info.Output.Fields.Get(\"nested\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaDict, nested.Type.Kind)\n\trequire.NotNil(t, nested.Type.ValueType)\n\trequire.Equal(t, schema.SchemaArray, nested.Type.ValueType.Kind)\n\trequire.NotNil(t, nested.Type.ValueType.Items)\n\trequire.Equal(t, schema.TypeString, nested.Type.ValueType.Items.Primitive)\n}\n\nfunc TestPydanticOutputWithOptionalDictField(t *testing.T) {\n\tsource := `\nfrom typing import Optional\nfrom pydantic import BaseModel\nfrom cog import BasePredictor\n\nclass Result(BaseModel):\n    data: Optional[dict[str, float]]\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`\n\tinfo := parse(t, source, \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\n\tdata, ok := info.Output.Fields.Get(\"data\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaDict, data.Type.Kind)\n\trequire.True(t, data.Type.Nullable)\n\trequire.False(t, data.Required)\n\trequire.NotNil(t, data.Type.ValueType)\n\trequire.Equal(t, schema.TypeFloat, data.Type.ValueType.Primitive)\n}\n\n// ---------------------------------------------------------------------------\n// Cross-file model resolution\n// ---------------------------------------------------------------------------\n\n// writeFile is a test helper that creates a file in dir with the given content.\nfunc writeFile(t *testing.T, dir, name, content string) {\n\tt.Helper()\n\tfull := filepath.Join(dir, name)\n\trequire.NoError(t, os.MkdirAll(filepath.Dir(full), 0o755))\n\trequire.NoError(t, os.WriteFile(full, []byte(content), 0o644))\n}\n\n// parseFile is a test helper that parses a file from disk with sourceDir context.\nfunc parseFile(t *testing.T, dir, filename, predictRef string) *schema.PredictorInfo {\n\tt.Helper()\n\tsource, err := os.ReadFile(filepath.Join(dir, filename))\n\trequire.NoError(t, err)\n\tinfo, err := ParsePredictor(source, predictRef, schema.ModePredict, dir)\n\trequire.NoError(t, err)\n\treturn info\n}\n\nfunc TestCrossFileBaseModelSameDir(t *testing.T) {\n\t// from types import Output — Output defined in types.py in same dir\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"types.py\", `\nfrom pydantic import BaseModel\n\nclass Output(BaseModel):\n    text: str\n    score: float\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom types import Output\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Output:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\ttext, ok := info.Output.Fields.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, text.Type.Primitive)\n\n\tscore, ok := info.Output.Fields.Get(\"score\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeFloat, score.Type.Primitive)\n}\n\nfunc TestCrossFileRelativeImport(t *testing.T) {\n\t// from .types import Output — relative dot import\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"types.py\", `\nfrom cog import BaseModel\n\nclass Output(BaseModel):\n    label: str\n    confidence: float\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom .types import Output\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Output:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\tlabel, ok := info.Output.Fields.Get(\"label\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, label.Type.Primitive)\n}\n\nfunc TestCrossFileSubpackageImport(t *testing.T) {\n\t// from models.output import Result — nested package\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"models/output.py\", `\nfrom pydantic import BaseModel\n\nclass Result(BaseModel):\n    answer: str\n    tokens: int\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom models.output import Result\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\tanswer, ok := info.Output.Fields.Get(\"answer\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, answer.Type.Primitive)\n\n\ttokens, ok := info.Output.Fields.Get(\"tokens\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeInteger, tokens.Type.Primitive)\n}\n\nfunc TestCrossFileRelativeSubpackage(t *testing.T) {\n\t// from .models.output import Result — relative + nested\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"models/output.py\", `\nfrom pydantic import BaseModel\n\nclass Result(BaseModel):\n    name: str\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom .models.output import Result\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Result:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 1, info.Output.Fields.Len())\n\n\tname, ok := info.Output.Fields.Get(\"name\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, name.Type.Primitive)\n}\n\nfunc TestCrossFileMultipleModelsFromSameFile(t *testing.T) {\n\t// Two BaseModel classes in the same external file\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"schema_types.py\", `\nfrom pydantic import BaseModel\n\nclass Metadata(BaseModel):\n    version: str\n    author: str\n\nclass Prediction(BaseModel):\n    result: str\n    score: float\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom schema_types import Prediction\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Prediction:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\tresult, ok := info.Output.Fields.Get(\"result\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, result.Type.Primitive)\n\n\tscore, ok := info.Output.Fields.Get(\"score\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeFloat, score.Type.Primitive)\n}\n\nfunc TestCrossFileWithOptionalField(t *testing.T) {\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"output.py\", `\nfrom typing import Optional\nfrom pydantic import BaseModel\n\nclass Output(BaseModel):\n    text: str\n    error: Optional[str] = None\n    debug: bool = False\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom output import Output\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Output:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 3, info.Output.Fields.Len())\n\n\ttext, ok := info.Output.Fields.Get(\"text\")\n\trequire.True(t, ok)\n\trequire.True(t, text.Required)\n\n\terrField, ok := info.Output.Fields.Get(\"error\")\n\trequire.True(t, ok)\n\trequire.True(t, errField.Type.Nullable)\n\n\tdebug, ok := info.Output.Fields.Get(\"debug\")\n\trequire.True(t, ok)\n\trequire.NotNil(t, debug.Default)\n\trequire.Equal(t, schema.DefaultBool, debug.Default.Kind)\n\trequire.Equal(t, false, debug.Default.Bool)\n}\n\nfunc TestCrossFileAliasedImport(t *testing.T) {\n\t// from output_types import MyOutput as Output\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"output_types.py\", `\nfrom pydantic import BaseModel\n\nclass MyOutput(BaseModel):\n    value: int\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom output_types import MyOutput as Output\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Output:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 1, info.Output.Fields.Len())\n\n\tval, ok := info.Output.Fields.Get(\"value\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeInteger, val.Type.Primitive)\n}\n\nfunc TestCrossFileExternalPackageStillErrors(t *testing.T) {\n\t// Importing from a package that doesn't exist locally should still error\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"predict.py\", `\nfrom transformers import AutoModelForSequenceClassification\nfrom cog import BasePredictor\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> AutoModelForSequenceClassification:\n        pass\n`)\n\tsource, err := os.ReadFile(filepath.Join(dir, \"predict.py\"))\n\trequire.NoError(t, err)\n\t_, err = ParsePredictor(source, \"Predictor\", schema.ModePredict, dir)\n\trequire.Error(t, err)\n\tvar se *schema.SchemaError\n\trequire.True(t, errors.As(err, &se))\n\trequire.Equal(t, schema.ErrUnresolvableType, se.Kind)\n\trequire.Contains(t, se.Message, \"transformers\")\n}\n\nfunc TestCrossFileLocalPrecedesExternal(t *testing.T) {\n\t// A local file shadows an external package name.\n\t// E.g. user has a local \"utils.py\" and does \"from utils import Output\"\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"utils.py\", `\nfrom cog import BaseModel\n\nclass Output(BaseModel):\n    msg: str\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom utils import Output\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Output:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 1, info.Output.Fields.Len())\n\n\tmsg, ok := info.Output.Fields.Get(\"msg\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.TypeString, msg.Type.Primitive)\n}\n\nfunc TestCrossFileListFieldInExternalModel(t *testing.T) {\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"types.py\", `\nfrom pydantic import BaseModel\n\nclass Output(BaseModel):\n    tags: list[str]\n    scores: list[float]\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom types import Output\n\nclass Predictor(BasePredictor):\n    def predict(self, x: str) -> Output:\n        pass\n`)\n\tinfo := parseFile(t, dir, \"predict.py\", \"Predictor\")\n\trequire.Equal(t, schema.SchemaObject, info.Output.Kind)\n\trequire.Equal(t, 2, info.Output.Fields.Len())\n\n\ttags, ok := info.Output.Fields.Get(\"tags\")\n\trequire.True(t, ok)\n\trequire.Equal(t, schema.SchemaArray, tags.Type.Kind)\n\trequire.Equal(t, schema.TypeString, tags.Type.Items.Primitive)\n}\n\nfunc TestCrossFileEndToEndSchemaGeneration(t *testing.T) {\n\t// Full end-to-end: Generate() reads predict.py from disk,\n\t// resolves Output from types.py, and produces valid OpenAPI JSON.\n\tdir := t.TempDir()\n\n\twriteFile(t, dir, \"types.py\", `\nfrom pydantic import BaseModel\n\nclass Output(BaseModel):\n    text: str\n    score: float\n`)\n\twriteFile(t, dir, \"predict.py\", `\nfrom cog import BasePredictor\nfrom types import Output\n\nclass Predictor(BasePredictor):\n    def predict(self, prompt: str) -> Output:\n        pass\n`)\n\n\tdata, err := schema.Generate(\"predict.py:Predictor\", dir, schema.ModePredict, ParsePredictor)\n\trequire.NoError(t, err)\n\trequire.Contains(t, string(data), `\"openapi\"`)\n\trequire.Contains(t, string(data), `\"Output\"`)\n\trequire.Contains(t, string(data), `\"text\"`)\n\trequire.Contains(t, string(data), `\"score\"`)\n\trequire.Contains(t, string(data), `\"object\"`)\n}\n"
  },
  {
    "path": "pkg/schema/schema_type.go",
    "content": "package schema\n\nimport \"fmt\"\n\n// SchemaType is a recursive algebraic data type representing any type that\n// can appear in a Cog predictor's output (or, in the future, input) position.\n//\n// It replaces the flat OutputType/PrimitiveType system with a composable\n// type tree that can represent dict[str, list[int]], nested BaseModel\n// subclasses, TypedDicts, and types resolved from .pyi stubs — all without\n// running Python.\ntype SchemaType struct {\n\tKind SchemaTypeKind\n\n\t// Primitive: for Kind=SchemaPrimitive — one of the base scalar types.\n\tPrimitive PrimitiveType\n\n\t// Array: for Kind=SchemaArray — the element type.\n\tItems *SchemaType\n\n\t// Dict: for Kind=SchemaDict — key and value types.\n\t// KeyType is always string in JSON Schema, but we track it for completeness.\n\tKeyType   *SchemaType\n\tValueType *SchemaType\n\n\t// Object: for Kind=SchemaObject — named fields with types and defaults.\n\tFields *OrderedMap[string, SchemaField]\n\n\t// Iterator/ConcatIterator: for Kind=SchemaIterator|SchemaConcatIterator.\n\t// The yielded element type.\n\tElem *SchemaType\n\n\t// Nullable: wraps any type to allow null.\n\tNullable bool\n}\n\n// SchemaTypeKind tags the active variant in SchemaType.\ntype SchemaTypeKind int\n\nconst (\n\t// SchemaPrimitive is a scalar type: bool, int, float, str, Path, File, Secret.\n\tSchemaPrimitive SchemaTypeKind = iota\n\t// SchemaAny is an opaque JSON value (unparameterized dict, Any, etc).\n\tSchemaAny\n\t// SchemaArray is a homogeneous list/array.\n\tSchemaArray\n\t// SchemaDict is a string-keyed dictionary with a typed value.\n\tSchemaDict\n\t// SchemaObject is a product type with named fields (BaseModel, TypedDict, dataclass).\n\tSchemaObject\n\t// SchemaIterator is a cog Iterator[T] — array with x-cog-array-type=iterator.\n\tSchemaIterator\n\t// SchemaConcatIterator is a cog ConcatenateIterator[str] — streaming text.\n\tSchemaConcatIterator\n)\n\n// SchemaField is a named field within a SchemaObject.\ntype SchemaField struct {\n\tType     SchemaType\n\tDefault  *DefaultValue\n\tRequired bool\n}\n\n// JSONSchema converts a SchemaType to its JSON Schema representation.\n// This is used for the \"Output\" component in the OpenAPI spec.\nfunc (s SchemaType) JSONSchema() map[string]any {\n\treturn s.jsonSchema(true)\n}\n\nfunc (s SchemaType) jsonSchema(topLevel bool) map[string]any {\n\tresult := s.coreSchema()\n\n\tif topLevel {\n\t\tresult[\"title\"] = \"Output\"\n\t}\n\n\tif s.Nullable {\n\t\tresult[\"nullable\"] = true\n\t}\n\n\treturn result\n}\n\nfunc (s SchemaType) coreSchema() map[string]any {\n\tswitch s.Kind {\n\tcase SchemaPrimitive:\n\t\treturn s.Primitive.JSONType()\n\n\tcase SchemaAny:\n\t\treturn map[string]any{\"type\": \"object\"}\n\n\tcase SchemaArray:\n\t\titems := map[string]any{\"type\": \"object\"}\n\t\tif s.Items != nil {\n\t\t\titems = s.Items.jsonSchema(false)\n\t\t}\n\t\tresult := map[string]any{\n\t\t\t\"type\":  \"array\",\n\t\t\t\"items\": items,\n\t\t}\n\t\treturn result\n\n\tcase SchemaDict:\n\t\tresult := map[string]any{\"type\": \"object\"}\n\t\tif s.ValueType != nil {\n\t\t\tresult[\"additionalProperties\"] = s.ValueType.jsonSchema(false)\n\t\t}\n\t\treturn result\n\n\tcase SchemaObject:\n\t\tif s.Fields == nil {\n\t\t\treturn map[string]any{\"type\": \"object\"}\n\t\t}\n\t\tproperties := make(map[string]any)\n\t\tvar required []string\n\t\ts.Fields.Entries(func(name string, field SchemaField) {\n\t\t\tprop := field.Type.jsonSchema(false)\n\t\t\tprop[\"title\"] = TitleCase(name)\n\t\t\tif field.Type.Nullable {\n\t\t\t\tprop[\"nullable\"] = true\n\t\t\t}\n\t\t\tif field.Required && field.Default == nil {\n\t\t\t\trequired = append(required, name)\n\t\t\t}\n\t\t\tproperties[name] = prop\n\t\t})\n\t\tresult := map[string]any{\n\t\t\t\"type\":       \"object\",\n\t\t\t\"properties\": properties,\n\t\t}\n\t\tif len(required) > 0 {\n\t\t\tresult[\"required\"] = required\n\t\t}\n\t\treturn result\n\n\tcase SchemaIterator:\n\t\titems := map[string]any{\"type\": \"object\"}\n\t\tif s.Elem != nil {\n\t\t\titems = s.Elem.jsonSchema(false)\n\t\t}\n\t\treturn map[string]any{\n\t\t\t\"type\":             \"array\",\n\t\t\t\"items\":            items,\n\t\t\t\"x-cog-array-type\": \"iterator\",\n\t\t}\n\n\tcase SchemaConcatIterator:\n\t\titems := map[string]any{\"type\": \"object\"}\n\t\tif s.Elem != nil {\n\t\t\titems = s.Elem.jsonSchema(false)\n\t\t}\n\t\treturn map[string]any{\n\t\t\t\"type\":                \"array\",\n\t\t\t\"items\":               items,\n\t\t\t\"x-cog-array-type\":    \"iterator\",\n\t\t\t\"x-cog-array-display\": \"concatenate\",\n\t\t}\n\t}\n\n\treturn map[string]any{\"type\": \"object\"}\n}\n\n// ---------------------------------------------------------------------------\n// Constructors — convenience functions for building SchemaType values.\n// ---------------------------------------------------------------------------\n\n// SchemaPrim creates a primitive SchemaType.\nfunc SchemaPrim(p PrimitiveType) SchemaType {\n\treturn SchemaType{Kind: SchemaPrimitive, Primitive: p}\n}\n\n// SchemaAnyType creates an opaque JSON object type.\nfunc SchemaAnyType() SchemaType {\n\treturn SchemaType{Kind: SchemaAny}\n}\n\n// SchemaArrayOf creates an array type with the given element type.\nfunc SchemaArrayOf(elem SchemaType) SchemaType {\n\treturn SchemaType{Kind: SchemaArray, Items: &elem}\n}\n\n// SchemaDictOf creates a dict type with string keys and the given value type.\nfunc SchemaDictOf(value SchemaType) SchemaType {\n\tk := SchemaPrim(TypeString)\n\treturn SchemaType{Kind: SchemaDict, KeyType: &k, ValueType: &value}\n}\n\n// SchemaIteratorOf creates an iterator type with the given element type.\nfunc SchemaIteratorOf(elem SchemaType) SchemaType {\n\treturn SchemaType{Kind: SchemaIterator, Elem: &elem}\n}\n\n// SchemaConcatIteratorOf creates a concatenate iterator type (always str).\nfunc SchemaConcatIteratorOf() SchemaType {\n\telem := SchemaPrim(TypeString)\n\treturn SchemaType{Kind: SchemaConcatIterator, Elem: &elem}\n}\n\n// SchemaObjectOf creates an object type from an ordered map of fields.\nfunc SchemaObjectOf(fields *OrderedMap[string, SchemaField]) SchemaType {\n\treturn SchemaType{Kind: SchemaObject, Fields: fields}\n}\n\n// ---------------------------------------------------------------------------\n// ResolveSchemaType — recursive output type resolver (replaces ResolveOutputType).\n// ---------------------------------------------------------------------------\n\n// ResolveSchemaType resolves a Python type annotation into a SchemaType.\n// Unlike the legacy ResolveOutputType, this handles arbitrary nesting:\n//\n//\tdict[str, list[dict[str, int]]]  →  SchemaDict{ValueType: SchemaArray{Items: SchemaDict{...}}}\n//\tlist[dict[str, str]]             →  SchemaArray{Items: SchemaDict{ValueType: SchemaPrim(TypeString)}}\n//\n// It also resolves BaseModel subclasses and cog iterators.\nfunc ResolveSchemaType(ann TypeAnnotation, ctx *ImportContext, models ModelClassMap) (SchemaType, error) {\n\tswitch ann.Kind {\n\tcase TypeAnnotSimple:\n\t\treturn resolveSimpleSchemaType(ann, ctx, models)\n\tcase TypeAnnotGeneric:\n\t\treturn resolveGenericSchemaType(ann, ctx, models)\n\tcase TypeAnnotUnion:\n\t\treturn resolveUnionSchemaType(ann)\n\t}\n\treturn SchemaType{}, errUnsupportedType(\"unknown type annotation\")\n}\n\nfunc resolveSimpleSchemaType(ann TypeAnnotation, ctx *ImportContext, models ModelClassMap) (SchemaType, error) {\n\t// Check for BaseModel subclass\n\tif fields, ok := models.Get(ann.Name); ok {\n\t\treturn resolveModelToSchemaType(fields, ctx, models)\n\t}\n\n\t// Unparameterized dict → opaque JSON object\n\tif ann.Name == \"Any\" || ann.Name == \"dict\" || ann.Name == \"Dict\" {\n\t\treturn SchemaAnyType(), nil\n\t}\n\n\t// Unparameterized list → array of opaque objects\n\tif ann.Name == \"list\" || ann.Name == \"List\" {\n\t\treturn SchemaArrayOf(SchemaAnyType()), nil\n\t}\n\n\tprim, ok := PrimitiveFromName(ann.Name)\n\tif !ok {\n\t\t// Check if this name was imported from an external package\n\t\tif entry, imported := ctx.Names.Get(ann.Name); imported {\n\t\t\treturn SchemaType{}, errUnresolvableImportedType(ann.Name, entry.Module)\n\t\t}\n\t\treturn SchemaType{}, errUnresolvableType(ann.Name)\n\t}\n\treturn SchemaPrim(prim), nil\n}\n\nfunc resolveGenericSchemaType(ann TypeAnnotation, ctx *ImportContext, models ModelClassMap) (SchemaType, error) {\n\touter := ann.Name\n\n\t// dict[K, V] — recursively resolve value type\n\tif outer == \"dict\" || outer == \"Dict\" {\n\t\tif len(ann.Args) == 2 {\n\t\t\tvalType, err := ResolveSchemaType(ann.Args[1], ctx, models)\n\t\t\tif err != nil {\n\t\t\t\treturn SchemaType{}, fmt.Errorf(\"resolving dict value type: %w\", err)\n\t\t\t}\n\t\t\treturn SchemaDictOf(valType), nil\n\t\t}\n\t\t// Bare dict (no type args) → opaque\n\t\tif len(ann.Args) == 0 {\n\t\t\treturn SchemaAnyType(), nil\n\t\t}\n\t\treturn SchemaType{}, errUnsupportedType(\"dict expects 0 or 2 type arguments\")\n\t}\n\n\t// Optional[X] → rejected as output type (nullable outputs not supported)\n\tif outer == \"Optional\" {\n\t\tif len(ann.Args) != 1 {\n\t\t\treturn SchemaType{}, errUnsupportedType(\"Optional expects exactly 1 type argument\")\n\t\t}\n\t\t// Optional is not allowed as an output type\n\t\treturn SchemaType{}, errOptionalOutput()\n\t}\n\n\t// Union[X, Y] → delegate\n\tif outer == \"Union\" {\n\t\treturn resolveUnionSchemaType(TypeAnnotation{Kind: TypeAnnotUnion, Args: ann.Args})\n\t}\n\n\t// list[X] / List[X]\n\tif outer == \"List\" || outer == \"list\" {\n\t\tif len(ann.Args) != 1 {\n\t\t\treturn SchemaType{}, errUnsupportedType(\"list expects exactly 1 type argument\")\n\t\t}\n\t\telemType, err := ResolveSchemaType(ann.Args[0], ctx, models)\n\t\tif err != nil {\n\t\t\treturn SchemaType{}, err\n\t\t}\n\t\treturn SchemaArrayOf(elemType), nil\n\t}\n\n\t// Cog iterators — single type arg, recursively resolved (supports nested types)\n\tif outer == \"Iterator\" || outer == \"AsyncIterator\" {\n\t\tif len(ann.Args) != 1 {\n\t\t\treturn SchemaType{}, errUnsupportedType(\"Iterator expects exactly 1 type argument\")\n\t\t}\n\t\telemType, err := ResolveSchemaType(ann.Args[0], ctx, models)\n\t\tif err != nil {\n\t\t\treturn SchemaType{}, err\n\t\t}\n\t\treturn SchemaIteratorOf(elemType), nil\n\t}\n\n\tif outer == \"ConcatenateIterator\" || outer == \"AsyncConcatenateIterator\" {\n\t\tif len(ann.Args) != 1 {\n\t\t\treturn SchemaType{}, errUnsupportedType(\"ConcatenateIterator expects exactly 1 type argument\")\n\t\t}\n\t\tinner := ann.Args[0]\n\t\tif inner.Kind != TypeAnnotSimple {\n\t\t\treturn SchemaType{}, errUnsupportedType(\"ConcatenateIterator element type must be a simple type\")\n\t\t}\n\t\tprim, ok := PrimitiveFromName(inner.Name)\n\t\tif !ok || prim != TypeString {\n\t\t\treturn SchemaType{}, errConcatIteratorNotStr(inner.Name)\n\t\t}\n\t\treturn SchemaConcatIteratorOf(), nil\n\t}\n\n\treturn SchemaType{}, errUnsupportedType(fmt.Sprintf(\"%s[...] is not a supported output type\", outer))\n}\n\nfunc resolveUnionSchemaType(ann TypeAnnotation) (SchemaType, error) {\n\tif _, ok := UnwrapOptional(ann); ok {\n\t\treturn SchemaType{}, errOptionalOutput()\n\t}\n\treturn SchemaType{}, errUnsupportedType(\"union types are not supported as output\")\n}\n\n// resolveModelToSchemaType converts a BaseModel's fields into a SchemaObject.\n// Fields are resolved via resolveFieldSchemaType which supports the full recursive\n// SchemaType system (dict[str, list[int]], nested BaseModels, etc.) plus Optional\n// wrapping (which is valid for fields but not for top-level output types).\nfunc resolveModelToSchemaType(modelFields []ModelField, ctx *ImportContext, models ModelClassMap) (SchemaType, error) {\n\tfields := NewOrderedMap[string, SchemaField]()\n\tfor _, f := range modelFields {\n\t\tst, required, err := resolveFieldSchemaType(f.Type, ctx, models)\n\t\tif err != nil {\n\t\t\treturn SchemaType{}, fmt.Errorf(\"field %q: %w\", f.Name, err)\n\t\t}\n\t\tif f.Default != nil {\n\t\t\trequired = false\n\t\t}\n\t\tfields.Set(f.Name, SchemaField{\n\t\t\tType:     st,\n\t\t\tDefault:  f.Default,\n\t\t\tRequired: required,\n\t\t})\n\t}\n\treturn SchemaObjectOf(fields), nil\n}\n\n// resolveFieldSchemaType resolves a type annotation for a model field.\n// Unlike ResolveSchemaType (which rejects Optional as a top-level output),\n// this allows Optional[X] and Union[X, None] for fields, setting Nullable.\nfunc resolveFieldSchemaType(ann TypeAnnotation, ctx *ImportContext, models ModelClassMap) (SchemaType, bool, error) {\n\tif inner, ok := UnwrapOptional(ann); ok {\n\t\tst, err := ResolveSchemaType(inner, ctx, models)\n\t\tif err != nil {\n\t\t\treturn SchemaType{}, false, err\n\t\t}\n\t\tst.Nullable = true\n\t\treturn st, false, nil\n\t}\n\n\tst, err := ResolveSchemaType(ann, ctx, models)\n\tif err != nil {\n\t\treturn SchemaType{}, false, err\n\t}\n\treturn st, true, nil\n}\n"
  },
  {
    "path": "pkg/schema/schema_type_fuzz_test.go",
    "content": "package schema\n\nimport (\n\t\"testing\"\n)\n\n// FuzzResolveSchemaType builds arbitrary TypeAnnotation trees from fuzz input\n// and verifies that ResolveSchemaType never panics.\nfunc FuzzResolveSchemaType(f *testing.F) {\n\t// Seed corpus — known-good and known-tricky inputs.\n\tseeds := []TypeAnnotation{\n\t\t{Kind: TypeAnnotSimple, Name: \"str\"},\n\t\t{Kind: TypeAnnotSimple, Name: \"int\"},\n\t\t{Kind: TypeAnnotSimple, Name: \"dict\"},\n\t\t{Kind: TypeAnnotSimple, Name: \"list\"},\n\t\t{Kind: TypeAnnotSimple, Name: \"Any\"},\n\t\t{Kind: TypeAnnotSimple, Name: \"UnknownType\"},\n\t\t{Kind: TypeAnnotSimple, Name: \"\"},\n\t\t{Kind: TypeAnnotGeneric, Name: \"dict\", Args: []TypeAnnotation{\n\t\t\t{Kind: TypeAnnotSimple, Name: \"str\"},\n\t\t\t{Kind: TypeAnnotSimple, Name: \"int\"},\n\t\t}},\n\t\t{Kind: TypeAnnotGeneric, Name: \"list\", Args: []TypeAnnotation{\n\t\t\t{Kind: TypeAnnotSimple, Name: \"str\"},\n\t\t}},\n\t\t{Kind: TypeAnnotGeneric, Name: \"Optional\", Args: []TypeAnnotation{\n\t\t\t{Kind: TypeAnnotSimple, Name: \"str\"},\n\t\t}},\n\t\t{Kind: TypeAnnotGeneric, Name: \"Iterator\", Args: []TypeAnnotation{\n\t\t\t{Kind: TypeAnnotGeneric, Name: \"dict\", Args: []TypeAnnotation{\n\t\t\t\t{Kind: TypeAnnotSimple, Name: \"str\"},\n\t\t\t\t{Kind: TypeAnnotGeneric, Name: \"list\", Args: []TypeAnnotation{\n\t\t\t\t\t{Kind: TypeAnnotSimple, Name: \"int\"},\n\t\t\t\t}},\n\t\t\t}},\n\t\t}},\n\t\t{Kind: TypeAnnotUnion, Args: []TypeAnnotation{\n\t\t\t{Kind: TypeAnnotSimple, Name: \"str\"},\n\t\t\t{Kind: TypeAnnotSimple, Name: \"None\"},\n\t\t}},\n\t}\n\n\t// Add byte-encoded seeds.\n\tfor _, s := range seeds {\n\t\tb := encodeAnnotation(s)\n\t\tf.Add(b)\n\t}\n\n\tctx := NewImportContext()\n\tmodels := NewOrderedMap[string, []ModelField]()\n\n\tf.Fuzz(func(t *testing.T, data []byte) {\n\t\tann, _ := decodeAnnotation(data, 0, 0)\n\t\t// Must not panic regardless of input.\n\t\tst, err := ResolveSchemaType(ann, ctx, models)\n\t\tif err == nil {\n\t\t\t// If resolution succeeded, JSONSchema must not panic.\n\t\t\t_ = st.JSONSchema()\n\t\t}\n\t})\n}\n\n// FuzzJSONSchema constructs random SchemaType trees and ensures\n// JSONSchema() never panics.\nfunc FuzzJSONSchema(f *testing.F) {\n\tf.Add([]byte{0})\n\tf.Add([]byte{1})\n\tf.Add([]byte{2, 0, 3, 's', 't', 'r'})\n\tf.Add([]byte{3, 2, 0, 3, 's', 't', 'r'})\n\tf.Add([]byte{4, 1, 2, 0, 3, 'i', 'n', 't'})\n\n\tf.Fuzz(func(t *testing.T, data []byte) {\n\t\tst, _ := decodeSchemaType(data, 0, 0)\n\t\t// Must not panic.\n\t\t_ = st.JSONSchema()\n\t\t_ = st.jsonSchema(false)\n\t})\n}\n\n// ---------------------------------------------------------------------------\n// Annotation encoder/decoder — deterministic mapping from bytes to trees.\n// ---------------------------------------------------------------------------\n\nconst maxFuzzDepth = 8\n\n// encodeAnnotation serializes a TypeAnnotation to bytes.\nfunc encodeAnnotation(ann TypeAnnotation) []byte {\n\tbuf := append([]byte{byte(ann.Kind), byte(len(ann.Name))}, []byte(ann.Name)...)\n\tbuf = append(buf, byte(len(ann.Args)))\n\tfor _, a := range ann.Args {\n\t\tbuf = append(buf, encodeAnnotation(a)...)\n\t}\n\treturn buf\n}\n\n// decodeAnnotation deserializes bytes into a TypeAnnotation tree.\n// Returns the annotation and number of bytes consumed.\nfunc decodeAnnotation(data []byte, offset int, depth int) (TypeAnnotation, int) {\n\tif depth > maxFuzzDepth || offset >= len(data) {\n\t\treturn TypeAnnotation{Kind: TypeAnnotSimple, Name: \"str\"}, offset\n\t}\n\n\tkind := TypeAnnotationKind(data[offset] % 3)\n\toffset++\n\n\t// Read name length and name.\n\tnameLen := 0\n\tif offset < len(data) {\n\t\tnameLen = int(data[offset]) % 32 // cap name length\n\t\toffset++\n\t}\n\tif offset+nameLen > len(data) {\n\t\tnameLen = len(data) - offset\n\t}\n\tname := string(data[offset : offset+nameLen])\n\toffset += nameLen\n\n\t// Read args count.\n\tnumArgs := 0\n\tif offset < len(data) {\n\t\tnumArgs = int(data[offset]) % 4 // cap at 3 args\n\t\toffset++\n\t}\n\n\tvar args []TypeAnnotation\n\tfor i := 0; i < numArgs && offset < len(data); i++ {\n\t\targ, newOffset := decodeAnnotation(data, offset, depth+1)\n\t\targs = append(args, arg)\n\t\toffset = newOffset\n\t}\n\n\treturn TypeAnnotation{Kind: kind, Name: name, Args: args}, offset\n}\n\n// decodeSchemaType builds a SchemaType tree from bytes.\nfunc decodeSchemaType(data []byte, offset int, depth int) (SchemaType, int) {\n\tif depth > maxFuzzDepth || offset >= len(data) {\n\t\treturn SchemaPrim(TypeString), offset\n\t}\n\n\tkind := SchemaTypeKind(data[offset] % 7)\n\toffset++\n\n\tswitch kind {\n\tcase SchemaPrimitive:\n\t\tprim := PrimitiveType(0)\n\t\tif offset < len(data) {\n\t\t\tprim = PrimitiveType(data[offset] % 9)\n\t\t\toffset++\n\t\t}\n\t\tst := SchemaPrim(prim)\n\t\tif offset < len(data) && data[offset]%2 == 1 {\n\t\t\tst.Nullable = true\n\t\t}\n\t\tif offset < len(data) {\n\t\t\toffset++\n\t\t}\n\t\treturn st, offset\n\n\tcase SchemaAny:\n\t\treturn SchemaAnyType(), offset\n\n\tcase SchemaArray:\n\t\titems, newOffset := decodeSchemaType(data, offset, depth+1)\n\t\treturn SchemaArrayOf(items), newOffset\n\n\tcase SchemaDict:\n\t\tval, newOffset := decodeSchemaType(data, offset, depth+1)\n\t\treturn SchemaDictOf(val), newOffset\n\n\tcase SchemaObject:\n\t\tnumFields := 0\n\t\tif offset < len(data) {\n\t\t\tnumFields = int(data[offset]) % 5\n\t\t\toffset++\n\t\t}\n\t\tfields := NewOrderedMap[string, SchemaField]()\n\t\tfor i := 0; i < numFields && offset < len(data); i++ {\n\t\t\tnameLen := int(data[offset]) % 8\n\t\t\toffset++\n\t\t\tif offset+nameLen > len(data) {\n\t\t\t\tnameLen = len(data) - offset\n\t\t\t}\n\t\t\tname := string(data[offset : offset+nameLen])\n\t\t\toffset += nameLen\n\t\t\tft, newOffset := decodeSchemaType(data, offset, depth+1)\n\t\t\trequired := false\n\t\t\tif newOffset < len(data) {\n\t\t\t\trequired = data[newOffset]%2 == 0\n\t\t\t\tnewOffset++\n\t\t\t}\n\t\t\tfields.Set(name, SchemaField{Type: ft, Required: required})\n\t\t\toffset = newOffset\n\t\t}\n\t\treturn SchemaObjectOf(fields), offset\n\n\tcase SchemaIterator:\n\t\telem, newOffset := decodeSchemaType(data, offset, depth+1)\n\t\treturn SchemaIteratorOf(elem), newOffset\n\n\tcase SchemaConcatIterator:\n\t\treturn SchemaConcatIteratorOf(), offset\n\n\tdefault:\n\t\treturn SchemaPrim(TypeString), offset\n\t}\n}\n"
  },
  {
    "path": "pkg/schema/types.go",
    "content": "package schema\n\nimport (\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"strings\"\n)\n\n// Mode selects whether to extract predict or train signatures.\ntype Mode int\n\nconst (\n\tModePredict Mode = iota\n\tModeTrain\n)\n\n// PrimitiveType maps Python types to JSON Schema types.\ntype PrimitiveType int\n\nconst (\n\tTypeBool PrimitiveType = iota\n\tTypeFloat\n\tTypeInteger\n\tTypeString\n\tTypePath   // cog.Path — {\"type\":\"string\",\"format\":\"uri\"}\n\tTypeFile   // cog.File (deprecated) — same wire format as Path\n\tTypeSecret // cog.Secret — write-only, masked\n\tTypeAny    // typing.Any or unresolved\n)\n\n// JSONType returns the JSON Schema fragment for this primitive.\nfunc (p PrimitiveType) JSONType() map[string]any {\n\tswitch p {\n\tcase TypeBool:\n\t\treturn map[string]any{\"type\": \"boolean\"}\n\tcase TypeFloat:\n\t\treturn map[string]any{\"type\": \"number\"}\n\tcase TypeInteger:\n\t\treturn map[string]any{\"type\": \"integer\"}\n\tcase TypeString:\n\t\treturn map[string]any{\"type\": \"string\"}\n\tcase TypePath, TypeFile:\n\t\treturn map[string]any{\"type\": \"string\", \"format\": \"uri\"}\n\tcase TypeSecret:\n\t\treturn map[string]any{\"type\": \"string\", \"format\": \"password\", \"writeOnly\": true, \"x-cog-secret\": true}\n\tcase TypeAny:\n\t\treturn map[string]any{\"type\": \"object\"}\n\tdefault:\n\t\treturn map[string]any{\"type\": \"object\"}\n\t}\n}\n\nfunc (p PrimitiveType) String() string {\n\tnames := [...]string{\"bool\", \"float\", \"int\", \"str\", \"Path\", \"File\", \"Secret\", \"Any\"}\n\tif int(p) < len(names) {\n\t\treturn names[p]\n\t}\n\treturn \"unknown\"\n}\n\n// PrimitiveFromName resolves a simple type name to a PrimitiveType.\nfunc PrimitiveFromName(name string) (PrimitiveType, bool) {\n\tswitch name {\n\tcase \"bool\":\n\t\treturn TypeBool, true\n\tcase \"float\":\n\t\treturn TypeFloat, true\n\tcase \"int\":\n\t\treturn TypeInteger, true\n\tcase \"str\":\n\t\treturn TypeString, true\n\tcase \"Path\":\n\t\treturn TypePath, true\n\tcase \"File\":\n\t\treturn TypeFile, true\n\tcase \"Secret\":\n\t\treturn TypeSecret, true\n\tcase \"Any\":\n\t\treturn TypeAny, true\n\tdefault:\n\t\treturn 0, false\n\t}\n}\n\n// Repetition describes cardinality of a field.\ntype Repetition int\n\nconst (\n\tRequired Repetition = iota\n\tOptional\n\tRepeated // list[X]\n)\n\n// FieldType combines a primitive type with its cardinality.\ntype FieldType struct {\n\tPrimitive  PrimitiveType\n\tRepetition Repetition\n}\n\n// JSONType returns the JSON Schema fragment for this field type.\nfunc (ft FieldType) JSONType() map[string]any {\n\tif ft.Repetition == Repeated {\n\t\treturn map[string]any{\n\t\t\t\"type\":  \"array\",\n\t\t\t\"items\": ft.Primitive.JSONType(),\n\t\t}\n\t}\n\treturn ft.Primitive.JSONType()\n}\n\n// DefaultValue represents a statically-parsed Python literal.\ntype DefaultValue struct {\n\tKind     DefaultKind\n\tBool     bool\n\tInt      int64\n\tFloat    float64\n\tStr      string\n\tList     []DefaultValue\n\tDictKeys []DefaultValue // parallel with DictVals\n\tDictVals []DefaultValue\n}\n\n// DefaultKind tags the active field in DefaultValue.\ntype DefaultKind int\n\nconst (\n\tDefaultNone DefaultKind = iota\n\tDefaultBool\n\tDefaultInt\n\tDefaultFloat\n\tDefaultString\n\tDefaultList\n\tDefaultDict\n\tDefaultSet\n)\n\n// ToJSON converts a DefaultValue to its JSON representation.\nfunc (d DefaultValue) ToJSON() any {\n\tswitch d.Kind {\n\tcase DefaultNone:\n\t\treturn nil\n\tcase DefaultBool:\n\t\treturn d.Bool\n\tcase DefaultInt:\n\t\treturn d.Int\n\tcase DefaultFloat:\n\t\treturn d.Float\n\tcase DefaultString:\n\t\treturn d.Str\n\tcase DefaultList, DefaultSet:\n\t\titems := make([]any, len(d.List))\n\t\tfor i, v := range d.List {\n\t\t\titems[i] = v.ToJSON()\n\t\t}\n\t\treturn items\n\tcase DefaultDict:\n\t\tm := make(map[string]any, len(d.DictKeys))\n\t\tfor i := range d.DictKeys {\n\t\t\tkey := fmt.Sprintf(\"%v\", d.DictKeys[i].ToJSON())\n\t\t\tif d.DictKeys[i].Kind == DefaultString {\n\t\t\t\tkey = d.DictKeys[i].Str\n\t\t\t}\n\t\t\tm[key] = d.DictVals[i].ToJSON()\n\t\t}\n\t\treturn m\n\tdefault:\n\t\treturn nil\n\t}\n}\n\n// MarshalJSON implements json.Marshaler for DefaultValue.\nfunc (d DefaultValue) MarshalJSON() ([]byte, error) {\n\treturn json.Marshal(d.ToJSON())\n}\n\n// InputField represents one parameter of predict/train.\ntype InputField struct {\n\tName        string\n\tOrder       int\n\tFieldType   FieldType\n\tDefault     *DefaultValue\n\tDescription *string\n\tGE          *float64\n\tLE          *float64\n\tMinLength   *uint64\n\tMaxLength   *uint64\n\tRegex       *string\n\tChoices     []DefaultValue\n\tDeprecated  *bool\n}\n\n// IsRequired returns true if this field is required in the schema.\nfunc (f *InputField) IsRequired() bool {\n\treturn f.Default == nil && (f.FieldType.Repetition == Required || f.FieldType.Repetition == Repeated)\n}\n\n// PredictorInfo is the top-level extraction result.\ntype PredictorInfo struct {\n\tInputs *OrderedMap[string, InputField]\n\tOutput SchemaType\n\tMode   Mode\n}\n\n// TypeAnnotation is a parsed Python type annotation (intermediate, before resolution).\ntype TypeAnnotation struct {\n\tKind TypeAnnotationKind\n\tName string           // for Simple\n\tArgs []TypeAnnotation // for Generic (outer=Name, args=Args) or Union (members=Args)\n}\n\n// TypeAnnotationKind tags the variant.\ntype TypeAnnotationKind int\n\nconst (\n\tTypeAnnotSimple TypeAnnotationKind = iota\n\tTypeAnnotGeneric\n\tTypeAnnotUnion\n)\n\n// ImportContext tracks what names are imported from which modules.\ntype ImportContext struct {\n\t// Names maps local name → (module, original_name)\n\tNames *OrderedMap[string, ImportEntry]\n}\n\n// ImportEntry records where a name was imported from.\ntype ImportEntry struct {\n\tModule   string\n\tOriginal string\n}\n\n// NewImportContext creates an empty ImportContext.\nfunc NewImportContext() *ImportContext {\n\treturn &ImportContext{Names: NewOrderedMap[string, ImportEntry]()}\n}\n\n// IsCogType returns true if name was imported from the \"cog\" module.\nfunc (ctx *ImportContext) IsCogType(name string) bool {\n\tif e, ok := ctx.Names.Get(name); ok {\n\t\treturn e.Module == \"cog\"\n\t}\n\treturn false\n}\n\n// IsTypingType returns true if name was imported from \"typing\" or \"typing_extensions\".\nfunc (ctx *ImportContext) IsTypingType(name string) bool {\n\tif e, ok := ctx.Names.Get(name); ok {\n\t\treturn e.Module == \"typing\" || e.Module == \"typing_extensions\"\n\t}\n\treturn false\n}\n\n// IsBaseModel returns true if name resolves to cog.BaseModel or pydantic.BaseModel.\nfunc (ctx *ImportContext) IsBaseModel(name string) bool {\n\tif e, ok := ctx.Names.Get(name); ok {\n\t\treturn (e.Module == \"cog\" || e.Module == \"pydantic\" || e.Module == \"pydantic.v1\") && e.Original == \"BaseModel\"\n\t}\n\treturn false\n}\n\n// IsBasePredictor returns true if name resolves to cog.BasePredictor.\nfunc (ctx *ImportContext) IsBasePredictor(name string) bool {\n\tif e, ok := ctx.Names.Get(name); ok {\n\t\treturn e.Module == \"cog\" && e.Original == \"BasePredictor\"\n\t}\n\treturn false\n}\n\n// ResolveFieldType resolves a TypeAnnotation into a FieldType.\nfunc ResolveFieldType(ann TypeAnnotation, ctx *ImportContext) (FieldType, error) {\n\tswitch ann.Kind {\n\tcase TypeAnnotSimple:\n\t\tprim, ok := PrimitiveFromName(ann.Name)\n\t\tif !ok {\n\t\t\treturn FieldType{}, errUnsupportedType(ann.Name)\n\t\t}\n\t\treturn FieldType{Primitive: prim, Repetition: Required}, nil\n\n\tcase TypeAnnotGeneric:\n\t\touter := ann.Name\n\t\tif outer == \"Optional\" {\n\t\t\tif len(ann.Args) != 1 {\n\t\t\t\treturn FieldType{}, errUnsupportedType(fmt.Sprintf(\"Optional expects exactly 1 type argument, got %d\", len(ann.Args)))\n\t\t\t}\n\t\t\tinner, err := ResolveFieldType(ann.Args[0], ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn FieldType{}, err\n\t\t\t}\n\t\t\treturn FieldType{Primitive: inner.Primitive, Repetition: Optional}, nil\n\t\t}\n\t\tif outer == \"Union\" {\n\t\t\t// typing.Union[X, Y] → treat as union type\n\t\t\treturn ResolveFieldType(TypeAnnotation{Kind: TypeAnnotUnion, Args: ann.Args}, ctx)\n\t\t}\n\t\tif outer == \"List\" || outer == \"list\" {\n\t\t\tif len(ann.Args) != 1 {\n\t\t\t\treturn FieldType{}, errUnsupportedType(fmt.Sprintf(\"List expects exactly 1 type argument, got %d\", len(ann.Args)))\n\t\t\t}\n\t\t\tinner, err := ResolveFieldType(ann.Args[0], ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn FieldType{}, err\n\t\t\t}\n\t\t\tif inner.Repetition != Required {\n\t\t\t\treturn FieldType{}, errUnsupportedType(\"nested generics like List[Optional[X]] are not supported\")\n\t\t\t}\n\t\t\treturn FieldType{Primitive: inner.Primitive, Repetition: Repeated}, nil\n\t\t}\n\t\treturn FieldType{}, errUnsupportedType(fmt.Sprintf(\"%s[...] is not a supported input type\", outer))\n\n\tcase TypeAnnotUnion:\n\t\tif inner, ok := UnwrapOptional(ann); ok {\n\t\t\tft, err := ResolveFieldType(inner, ctx)\n\t\t\tif err != nil {\n\t\t\t\treturn FieldType{}, err\n\t\t\t}\n\t\t\treturn FieldType{Primitive: ft.Primitive, Repetition: Optional}, nil\n\t\t}\n\t\treturn FieldType{}, errUnsupportedType(\"union types other than X | None are not supported\")\n\t}\n\treturn FieldType{}, errUnsupportedType(\"unknown type annotation\")\n}\n\n// UnwrapOptional checks if a type annotation represents Optional[X] or Union[X, None].\n// If so, it returns the inner type and true. Otherwise it returns the original and false.\nfunc UnwrapOptional(ann TypeAnnotation) (TypeAnnotation, bool) {\n\t// Optional[X]\n\tif ann.Kind == TypeAnnotGeneric && ann.Name == \"Optional\" && len(ann.Args) == 1 {\n\t\treturn ann.Args[0], true\n\t}\n\t// Union[X, None] or X | None\n\targs := ann.Args\n\tif (ann.Kind == TypeAnnotGeneric && ann.Name == \"Union\" || ann.Kind == TypeAnnotUnion) && len(args) == 2 {\n\t\tfor i := range args {\n\t\t\tif args[i].Kind == TypeAnnotSimple && args[i].Name == \"None\" {\n\t\t\t\treturn args[1-i], true\n\t\t\t}\n\t\t}\n\t}\n\treturn ann, false\n}\n\n// ModelClassMap maps class names to their fields.\ntype ModelClassMap = *OrderedMap[string, []ModelField]\n\n// ModelField is a field extracted from a BaseModel subclass.\ntype ModelField struct {\n\tName    string\n\tType    TypeAnnotation\n\tDefault *DefaultValue\n}\n\n// TitleCase converts snake_case to Title Case.\nfunc TitleCase(s string) string {\n\tparts := strings.Split(s, \"_\")\n\tfor i, p := range parts {\n\t\tif len(p) > 0 {\n\t\t\tparts[i] = strings.ToUpper(p[:1]) + p[1:]\n\t\t}\n\t}\n\treturn strings.Join(parts, \" \")\n}\n\n// TitleCaseSingle title-cases a single word (first letter uppercase).\nfunc TitleCaseSingle(s string) string {\n\tif len(s) == 0 {\n\t\treturn s\n\t}\n\treturn strings.ToUpper(s[:1]) + s[1:]\n}\n\n// OrderedMap is a simple insertion-ordered map.\ntype OrderedMap[K comparable, V any] struct {\n\tkeys   []K\n\tvalues map[K]V\n}\n\n// NewOrderedMap creates a new empty OrderedMap.\nfunc NewOrderedMap[K comparable, V any]() *OrderedMap[K, V] {\n\treturn &OrderedMap[K, V]{values: make(map[K]V)}\n}\n\n// Set inserts or updates a key-value pair.\nfunc (m *OrderedMap[K, V]) Set(key K, value V) {\n\tif _, exists := m.values[key]; !exists {\n\t\tm.keys = append(m.keys, key)\n\t}\n\tm.values[key] = value\n}\n\n// Get returns the value for a key and whether it exists.\nfunc (m *OrderedMap[K, V]) Get(key K) (V, bool) {\n\tv, ok := m.values[key]\n\treturn v, ok\n}\n\n// Keys returns keys in insertion order.\nfunc (m *OrderedMap[K, V]) Keys() []K {\n\treturn m.keys\n}\n\n// Len returns the number of entries.\nfunc (m *OrderedMap[K, V]) Len() int {\n\treturn len(m.keys)\n}\n\n// Entries iterates over key-value pairs in insertion order.\nfunc (m *OrderedMap[K, V]) Entries(fn func(key K, value V)) {\n\tfor _, k := range m.keys {\n\t\tfn(k, m.values[k])\n\t}\n}\n"
  },
  {
    "path": "pkg/update/state.go",
    "content": "package update\n\nimport (\n\t\"encoding/json\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"time\"\n\n\t\"github.com/mitchellh/go-homedir\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n\t\"github.com/replicate/cog/pkg/util/files\"\n)\n\ntype state struct {\n\tMessage     string    `json:\"message\"`\n\tLastChecked time.Time `json:\"lastChecked\"`\n\tVersion     string    `json:\"version\"`\n}\n\n// loadState loads the update check state from disk, returning defaults if it does not exist\nfunc loadState() (*state, error) {\n\tstate := state{}\n\n\tp, err := statePath()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\texists, err := files.Exists(p)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tif !exists {\n\t\treturn &state, nil\n\t}\n\ttext, err := os.ReadFile(p)\n\tif err != nil {\n\t\tconsole.Debugf(\"Failed to read %s: %s\", p, err)\n\t\treturn &state, nil\n\t}\n\n\terr = json.Unmarshal(text, &state)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &state, nil\n}\n\n// writeState saves analytics state to disk\nfunc writeState(s *state) error {\n\tstatePath, err := statePath()\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tbytes, err := json.MarshalIndent(s, \"\", \" \")\n\tif err != nil {\n\t\treturn err\n\t}\n\tdir := filepath.Dir(statePath)\n\tif err := os.MkdirAll(dir, 0o700); err != nil {\n\t\treturn err\n\t}\n\n\terr = os.WriteFile(statePath, bytes, 0o600)\n\tif err != nil {\n\t\treturn err\n\t}\n\treturn nil\n}\n\nfunc userDir() (string, error) {\n\treturn homedir.Expand(\"~/.config/cog\")\n}\n\nfunc statePath() (string, error) {\n\tdir, err := userDir()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\treturn filepath.Join(dir, \"update-state.json\"), nil\n}\n"
  },
  {
    "path": "pkg/update/update.go",
    "content": "package update\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"net/http\"\n\t\"os\"\n\t\"runtime\"\n\t\"time\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc isUpdateEnabled() bool {\n\treturn os.Getenv(\"COG_NO_UPDATE_CHECK\") == \"\"\n}\n\n// DisplayAndCheckForRelease will display an update message if an update is available and will check for a new update in the background\n// The result of that check will then be displayed the next time the user runs Cog\n// Returns errors which the caller is assumed to ignore so as not to break the client\nfunc DisplayAndCheckForRelease(ctx context.Context) error {\n\tif !isUpdateEnabled() {\n\t\treturn fmt.Errorf(\"update check disabled\")\n\t}\n\n\ts, err := loadState()\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tif s.Version != global.Version {\n\t\tconsole.Debugf(\"Resetting update message because Cog has been upgraded\")\n\t\treturn writeState(&state{Message: \"\", LastChecked: time.Now(), Version: global.Version})\n\t}\n\n\tif time.Since(s.LastChecked) > time.Hour {\n\t\tstartCheckingForRelease(ctx)\n\t}\n\tif s.Message != \"\" {\n\t\tconsole.Info(s.Message)\n\t\tconsole.Info(\"\")\n\t}\n\treturn nil\n}\n\nfunc startCheckingForRelease(ctx context.Context) {\n\tgo func() {\n\t\tconsole.Debugf(\"Checking for updates...\")\n\t\tctx, cancel := context.WithTimeout(ctx, time.Second)\n\t\tdefer cancel()\n\t\tswitch r, err := checkForRelease(ctx); {\n\t\tcase err == nil:\n\t\t\tif r == nil {\n\t\t\t\tbreak\n\t\t\t}\n\t\t\tif err := writeState(&state{Message: r.Message, LastChecked: time.Now(), Version: global.Version}); err != nil {\n\t\t\t\tconsole.Debugf(\"Failed to write state: %s\", err)\n\t\t\t}\n\n\t\t\tconsole.Debugf(\"result of update check: %v\", r.Message)\n\t\tcase errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded):\n\t\t\tbreak\n\t\tdefault:\n\t\t\tconsole.Debugf(\"failed querying for new release: %v\", err)\n\t\t}\n\t}()\n}\n\ntype updateCheckResponse struct {\n\tMessage string `json:\"message\"`\n}\n\nfunc checkForRelease(ctx context.Context) (*updateCheckResponse, error) {\n\treq, err := http.NewRequestWithContext(ctx, \"GET\", \"https://update.cog.run/v1/check\", nil)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\treq.Header.Add(\"Accept\", \"application/json\")\n\tq := req.URL.Query()\n\tq.Add(\"version\", global.Version)\n\tq.Add(\"commit\", global.Commit)\n\tq.Add(\"os\", runtime.GOOS)\n\tq.Add(\"arch\", runtime.GOARCH)\n\treq.URL.RawQuery = q.Encode()\n\n\tresp, err := http.DefaultClient.Do(req) //nolint:gosec // G704: URL is built from hardcoded base + version params\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdefer resp.Body.Close()\n\n\tvar response updateCheckResponse\n\tif err := json.NewDecoder(resp.Body).Decode(&response); err != nil {\n\t\treturn &response, err\n\t}\n\n\treturn &response, nil\n}\n"
  },
  {
    "path": "pkg/util/console/console.go",
    "content": "// Package console provides a standard interface for user- and machine-interface with the console\npackage console\n\nimport (\n\t\"fmt\"\n\t\"math\"\n\t\"os\"\n\t\"strings\"\n\t\"sync\"\n\t\"unicode/utf8\"\n\n\t\"github.com/logrusorgru/aurora\"\n\t\"github.com/mattn/go-isatty\"\n\t\"golang.org/x/term\"\n)\n\n// ShouldUseColor returns true if color output should be enabled, based on\n// environment detection. It checks (in order):\n//   - NO_COLOR env var is set and non-empty → no color\n//   - COG_NO_COLOR env var is set and non-empty → no color\n//   - TERM=dumb → no color\n//   - stderr is not a TTY → no color\n//\n// This follows the NO_COLOR standard (https://no-color.org/) and common CLI\n// conventions. The --no-color flag is handled separately at the CLI layer.\nfunc ShouldUseColor() bool {\n\tif os.Getenv(\"NO_COLOR\") != \"\" {\n\t\treturn false\n\t}\n\tif os.Getenv(\"COG_NO_COLOR\") != \"\" {\n\t\treturn false\n\t}\n\tif os.Getenv(\"TERM\") == \"dumb\" {\n\t\treturn false\n\t}\n\tfd := os.Stderr.Fd()\n\tif !isatty.IsTerminal(fd) && !isatty.IsCygwinTerminal(fd) {\n\t\treturn false\n\t}\n\treturn true\n}\n\n// Style controls the icon/color used for a log line, independent of level.\ntype Style int\n\nconst (\n\t// StyleDefault uses the default icon for the log level.\n\tStyleDefault Style = iota\n\t// StyleSuccess uses a green ✓ icon.\n\tStyleSuccess\n)\n\n// Console represents a standardized interface for console UI. It is designed to abstract:\n// - Writing main output\n// - Giving information to user\n// - Console user interface elements (progress, interactive prompts, etc)\n// - Switching between human and machine modes for these things (e.g. don't display progress bars or colors in logs, don't prompt for input when in a script)\ntype Console struct {\n\tColor     bool\n\tIsMachine bool\n\tLevel     Level\n\tmu        sync.Mutex\n}\n\n// Debug prints a verbose debugging message, that is not displayed by default to the user.\nfunc (c *Console) Debug(msg string) {\n\tc.log(DebugLevel, msg)\n}\n\n// Info tells the user what's going on.\nfunc (c *Console) Info(msg string) {\n\tc.log(InfoLevel, msg)\n}\n\n// Success tells the user something completed successfully.\n// Displays at info level with a green ✓ prefix.\nfunc (c *Console) Success(msg string) {\n\tc.logStyled(InfoLevel, StyleSuccess, msg)\n}\n\n// Warn tells the user that something might break.\nfunc (c *Console) Warn(msg string) {\n\tc.log(WarnLevel, msg)\n}\n\n// Error tells the user that something is broken.\nfunc (c *Console) Error(msg string) {\n\tc.log(ErrorLevel, msg)\n}\n\n// Fatal level message, followed by exit\nfunc (c *Console) Fatal(msg string) {\n\tc.log(FatalLevel, msg)\n\tos.Exit(1)\n}\n\n// Debug level message\nfunc (c *Console) Debugf(msg string, v ...any) {\n\tc.log(DebugLevel, fmt.Sprintf(msg, v...))\n}\n\n// Info level message\nfunc (c *Console) Infof(msg string, v ...any) {\n\tc.log(InfoLevel, fmt.Sprintf(msg, v...))\n}\n\n// Success level message\nfunc (c *Console) Successf(msg string, v ...any) {\n\tc.logStyled(InfoLevel, StyleSuccess, fmt.Sprintf(msg, v...))\n}\n\n// Warn level message\nfunc (c *Console) Warnf(msg string, v ...any) {\n\tc.log(WarnLevel, fmt.Sprintf(msg, v...))\n}\n\n// Error level message\nfunc (c *Console) Errorf(msg string, v ...any) {\n\tc.log(ErrorLevel, fmt.Sprintf(msg, v...))\n}\n\n// Fatal level message, followed by exit\nfunc (c *Console) Fatalf(msg string, v ...any) {\n\tc.log(FatalLevel, fmt.Sprintf(msg, v...))\n\tos.Exit(1)\n}\n\n// InfoUnformatted writes a message to stderr without any prefix. Useful for conversational\n// or interactive output (e.g. login prompts) where the icon prefix would be noise.\n// Displayed at info level. Long lines are wrapped to terminal width when stderr is a TTY.\nfunc (c *Console) InfoUnformatted(msg string) {\n\tif InfoLevel < c.Level {\n\t\treturn\n\t}\n\n\ttermWidth := stderrTerminalWidth()\n\n\tc.mu.Lock()\n\tdefer c.mu.Unlock()\n\n\tfor line := range strings.SplitSeq(msg, \"\\n\") {\n\t\tif termWidth > 0 {\n\t\t\twrapped := wrapLine(line, termWidth)\n\t\t\tfor _, wl := range wrapped {\n\t\t\t\tfmt.Fprintln(os.Stderr, wl)\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\t\tfmt.Fprintln(os.Stderr, line)\n\t}\n}\n\n// InfoUnformattedf writes a formatted message to stderr without any prefix.\nfunc (c *Console) InfoUnformattedf(msg string, v ...any) {\n\tc.InfoUnformatted(fmt.Sprintf(msg, v...))\n}\n\n// Output a string to stdout. Useful for printing primary output of a command, or the output of a subcommand.\n// A newline is added to the string.\nfunc (c *Console) Output(s string) {\n\tc.mu.Lock()\n\tdefer c.mu.Unlock()\n\t_, _ = fmt.Fprintln(os.Stdout, s)\n}\n\n// Bold applies bold formatting to a string when color is enabled.\n// Use this to highlight dynamic values (image names, paths, URLs) in log messages.\nfunc (c *Console) Bold(s string) string {\n\tif c.Color {\n\t\treturn aurora.Bold(s).String()\n\t}\n\treturn s\n}\n\nfunc (c *Console) log(level Level, msg string) {\n\tc.logStyled(level, StyleDefault, msg)\n}\n\nfunc (c *Console) logStyled(level Level, style Style, msg string) {\n\tif level < c.Level {\n\t\treturn\n\t}\n\n\tprompt := \"\"\n\t// promptWidth is the visual width of the prompt (excluding ANSI codes).\n\tpromptWidth := 0\n\n\tif c.Color {\n\t\tswitch style {\n\t\tcase StyleSuccess:\n\t\t\tprompt = \" \" + aurora.Bold(aurora.Green(\"✔ \")).String()\n\t\t\tpromptWidth = 4 // \" ✔ \"\n\t\tdefault:\n\t\t\tswitch level {\n\t\t\tcase DebugLevel, InfoLevel:\n\t\t\t\tprompt = \" \" + aurora.Faint(\"⚙  \").String()\n\t\t\t\tpromptWidth = 4 // \" ⚙  \"\n\t\t\tcase WarnLevel:\n\t\t\t\tprompt = \" \" + aurora.Bold(aurora.Yellow(\"⚠ \")).String()\n\t\t\t\tpromptWidth = 4 // \" ⚠ \"\n\t\t\tcase ErrorLevel, FatalLevel:\n\t\t\t\tprompt = \" \" + aurora.Bold(aurora.Red(\"✗ \")).String()\n\t\t\t\tpromptWidth = 4 // \" ✗ \"\n\t\t\t}\n\t\t}\n\t}\n\n\ttermWidth := stderrTerminalWidth()\n\n\tc.mu.Lock()\n\tdefer c.mu.Unlock()\n\n\tfor line := range strings.SplitSeq(msg, \"\\n\") {\n\t\tif line == \"\" && (level == DebugLevel || level == InfoLevel) {\n\t\t\tfmt.Fprintln(os.Stderr)\n\t\t\tcontinue\n\t\t}\n\t\tif c.Color && level == DebugLevel {\n\t\t\tline = aurora.Faint(line).String()\n\t\t}\n\n\t\t// Wrap long lines to terminal width.\n\t\tif termWidth > 0 && promptWidth > 0 {\n\t\t\tmaxWidth := termWidth - promptWidth\n\t\t\tif maxWidth > 0 {\n\t\t\t\twrapped := wrapLine(line, maxWidth)\n\t\t\t\tfor _, wl := range wrapped {\n\t\t\t\t\tfmt.Fprintln(os.Stderr, prompt+wl)\n\t\t\t\t}\n\t\t\t\tcontinue\n\t\t\t}\n\t\t}\n\n\t\tfmt.Fprintln(os.Stderr, prompt+line)\n\t}\n}\n\n// stderrTerminalWidth returns the terminal width of stderr, or 0 if stderr\n// is not a terminal or the width cannot be determined.\nfunc stderrTerminalWidth() int {\n\tfd := os.Stderr.Fd()\n\tif !isatty.IsTerminal(fd) && !isatty.IsCygwinTerminal(fd) {\n\t\treturn 0\n\t}\n\tif fd > math.MaxInt {\n\t\treturn 0\n\t}\n\tw, _, err := term.GetSize(int(fd)) //nolint:gosec // bounded above\n\tif err != nil || w <= 0 {\n\t\treturn 0\n\t}\n\treturn w\n}\n\n// wrapLine wraps a single line of text to the given width, breaking on word\n// boundaries where possible. It operates on the visible text (which may contain\n// ANSI escape codes — these are counted as zero-width for wrapping purposes).\nfunc wrapLine(line string, maxWidth int) []string {\n\tif visibleWidth(line) <= maxWidth {\n\t\treturn []string{line}\n\t}\n\n\tvar lines []string\n\tfor len(line) > 0 {\n\t\tif visibleWidth(line) <= maxWidth {\n\t\t\tlines = append(lines, line)\n\t\t\tbreak\n\t\t}\n\n\t\t// Find the byte position where we exceed maxWidth visible chars.\n\t\tcutByte := findCutPoint(line, maxWidth)\n\n\t\t// Try to break at a space before the cut point.\n\t\tbreakAt := strings.LastIndex(line[:cutByte], \" \")\n\t\tif breakAt <= 0 {\n\t\t\t// No good break point; hard-break at cutByte.\n\t\t\tbreakAt = cutByte\n\t\t}\n\n\t\tlines = append(lines, line[:breakAt])\n\t\tline = strings.TrimLeft(line[breakAt:], \" \")\n\t}\n\treturn lines\n}\n\n// visibleWidth returns the number of visible characters in a string,\n// ignoring ANSI escape sequences.\nfunc visibleWidth(s string) int {\n\twidth := 0\n\tinEscape := false\n\tfor _, r := range s {\n\t\tif inEscape {\n\t\t\tif (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') {\n\t\t\t\tinEscape = false\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\t\tif r == '\\x1b' {\n\t\t\tinEscape = true\n\t\t\tcontinue\n\t\t}\n\t\twidth += utf8.RuneLen(r) // approximate: 1 for ASCII, may differ for wide chars\n\t\tif r > 127 {\n\t\t\twidth = width - utf8.RuneLen(r) + 1 // count non-ASCII runes as width 1\n\t\t}\n\t}\n\treturn width\n}\n\n// findCutPoint returns the byte index in s where the visible width reaches maxWidth.\nfunc findCutPoint(s string, maxWidth int) int {\n\twidth := 0\n\tinEscape := false\n\tfor i, r := range s {\n\t\tif inEscape {\n\t\t\tif (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') {\n\t\t\t\tinEscape = false\n\t\t\t}\n\t\t\tcontinue\n\t\t}\n\t\tif r == '\\x1b' {\n\t\t\tinEscape = true\n\t\t\tcontinue\n\t\t}\n\t\twidth++\n\t\tif width >= maxWidth {\n\t\t\treturn i + utf8.RuneLen(r)\n\t\t}\n\t}\n\treturn len(s)\n}\n"
  },
  {
    "path": "pkg/util/console/formatting.go",
    "content": "package console\n\nimport (\n\t\"time\"\n\n\t\"github.com/xeonx/timeago\"\n)\n\nfunc FormatTime(t time.Time) string {\n\treturn timeago.English.Format(t)\n}\n"
  },
  {
    "path": "pkg/util/console/global.go",
    "content": "package console\n\nimport (\n\t\"os\"\n\n\t\"github.com/mattn/go-isatty\"\n)\n\n// ConsoleInstance is the global instance of console, so we don't have to pass it around everywhere\nvar ConsoleInstance = &Console{\n\tColor:     ShouldUseColor(),\n\tLevel:     InfoLevel,\n\tIsMachine: false,\n}\n\n// SetLevel sets log level\nfunc SetLevel(level Level) {\n\tConsoleInstance.Level = level\n}\n\n// SetColor sets whether to print colors\nfunc SetColor(color bool) {\n\tConsoleInstance.Color = color\n}\n\n// Debug level message.\nfunc Debug(msg string) {\n\tConsoleInstance.Debug(msg)\n}\n\n// Info level message.\nfunc Info(msg string) {\n\tConsoleInstance.Info(msg)\n}\n\n// Success level message.\nfunc Success(msg string) {\n\tConsoleInstance.Success(msg)\n}\n\n// Warn level message.\nfunc Warn(msg string) {\n\tConsoleInstance.Warn(msg)\n}\n\n// Error level message.\nfunc Error(msg string) {\n\tConsoleInstance.Error(msg)\n}\n\n// Fatal level message.\nfunc Fatal(msg string) {\n\tConsoleInstance.Fatal(msg)\n}\n\n// Debug level message.\nfunc Debugf(msg string, v ...any) {\n\tConsoleInstance.Debugf(msg, v...)\n}\n\n// Info level message.\nfunc Infof(msg string, v ...any) {\n\tConsoleInstance.Infof(msg, v...)\n}\n\n// Success level message.\nfunc Successf(msg string, v ...any) {\n\tConsoleInstance.Successf(msg, v...)\n}\n\n// Warn level message.\nfunc Warnf(msg string, v ...any) {\n\tConsoleInstance.Warnf(msg, v...)\n}\n\n// Error level message.\nfunc Errorf(msg string, v ...any) {\n\tConsoleInstance.Errorf(msg, v...)\n}\n\n// Fatal level message.\nfunc Fatalf(msg string, v ...any) {\n\tConsoleInstance.Fatalf(msg, v...)\n}\n\n// InfoUnformatted writes to stderr without prefix. Useful for interactive/conversational output.\nfunc InfoUnformatted(msg string) {\n\tConsoleInstance.InfoUnformatted(msg)\n}\n\n// InfoUnformattedf writes to stderr without prefix, with formatting.\nfunc InfoUnformattedf(msg string, v ...any) {\n\tConsoleInstance.InfoUnformattedf(msg, v...)\n}\n\n// Output a line to stdout. Useful for printing primary output of a command, or the output of a subcommand.\nfunc Output(s string) {\n\tConsoleInstance.Output(s)\n}\n\n// Bold applies bold formatting to a string when color is enabled.\nfunc Bold(s string) string {\n\treturn ConsoleInstance.Bold(s)\n}\n\n// IsTTY checks if a file is a TTY or not. E.g. IsTTY(os.Stdin)\nfunc IsTTY(f *os.File) bool {\n\treturn isatty.IsTerminal(f.Fd())\n}\n"
  },
  {
    "path": "pkg/util/console/interactive.go",
    "content": "package console\n\nimport (\n\t\"bufio\"\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n\t\"slices\"\n\t\"strings\"\n)\n\ntype Interactive struct {\n\tPrompt   string\n\tDefault  string\n\tOptions  []string\n\tRequired bool\n}\n\nfunc (i Interactive) Read() (string, error) {\n\tif i.Default != \"\" && i.Options != nil && !slices.Contains(i.Options, i.Default) {\n\t\tpanic(\"Default is not an option\")\n\t}\n\n\tparens := \"\"\n\tif i.Required {\n\t\tparens += \"required\"\n\t}\n\tif i.Default != \"\" {\n\t\tif parens != \"\" {\n\t\t\tparens += \", \"\n\t\t}\n\t\tparens += \"default: \" + i.Default\n\t}\n\tif i.Options != nil {\n\t\tif parens != \"\" {\n\t\t\tparens += \", \"\n\t\t}\n\t\tparens += \"options: \" + strings.Join(i.Options, \", \")\n\t}\n\tif parens != \"\" {\n\t\tparens = \" (\" + parens + \")\"\n\t}\n\n\tfor {\n\t\tfmt.Printf(\"%s%s: \", i.Prompt, parens)\n\t\treader := bufio.NewReader(os.Stdin)\n\t\ttext, err := reader.ReadString('\\n')\n\t\tif err != nil {\n\t\t\treturn \"\", err\n\t\t}\n\t\ttext = strings.TrimSpace(text)\n\t\tif text == \"\" && i.Default != \"\" {\n\t\t\ttext = i.Default\n\t\t}\n\n\t\tif i.Required && text == \"\" {\n\t\t\tWarn(\"Please enter a value\")\n\t\t\tcontinue\n\t\t}\n\n\t\tif !i.Required && text == \"\" {\n\t\t\treturn \"\", nil\n\t\t}\n\n\t\tif i.Options != nil {\n\t\t\tif !slices.Contains(i.Options, text) {\n\t\t\t\tWarnf(\"%s is not a valid option\", text)\n\t\t\t\tcontinue\n\t\t\t}\n\t\t}\n\n\t\treturn text, nil\n\t}\n}\n\ntype InteractiveBool struct {\n\tPrompt  string\n\tDefault bool\n\t// NonDefaultFlag is the flag to suggest passing to do the thing which isn't default when running inside a script\n\tNonDefaultFlag string\n}\n\nfunc (i InteractiveBool) Read() (bool, error) {\n\tdefaults := \"y/N\"\n\tif i.Default {\n\t\tdefaults = \"Y/n\"\n\t}\n\tfor {\n\t\tfmt.Printf(\"%s (%s) \", i.Prompt, defaults)\n\t\treader := bufio.NewReader(os.Stdin)\n\t\ttext, err := reader.ReadString('\\n')\n\t\tif err != nil {\n\t\t\t// Only translate error if a flag is set\n\t\t\tif err == io.EOF && i.NonDefaultFlag != \"\" {\n\t\t\t\treturn false, fmt.Errorf(\"stdin is closed. If you're running in a script, you need to pass the '%s' option\", i.NonDefaultFlag)\n\t\t\t}\n\t\t\treturn false, err\n\t\t}\n\t\ttext = strings.ToLower(strings.TrimSpace(text))\n\t\tif text == \"yes\" || text == \"y\" {\n\t\t\treturn true, nil\n\t\t}\n\t\tif text == \"no\" || text == \"n\" {\n\t\t\treturn false, nil\n\t\t}\n\t\tif text == \"\" {\n\t\t\treturn i.Default, nil\n\t\t}\n\t\tWarn(\"Please enter 'y' or 'n'\")\n\t}\n}\n"
  },
  {
    "path": "pkg/util/console/levels.go",
    "content": "package console\n\n// Mostly lifted from https://github.com/apex/log/blob/master/levels.go\n\nimport (\n\t\"errors\"\n\t\"strings\"\n)\n\n// ErrInvalidLevel is returned if the severity level is invalid.\nvar ErrInvalidLevel = errors.New(\"invalid level\")\n\n// Level of severity.\ntype Level int\n\n// Log levels.\nconst (\n\tInvalidLevel Level = iota - 1\n\tDebugLevel\n\tInfoLevel\n\tWarnLevel\n\tErrorLevel\n\tFatalLevel\n)\n\nvar levelNames = [...]string{\n\tDebugLevel: \"debug\",\n\tInfoLevel:  \"info\",\n\tWarnLevel:  \"warn\",\n\tErrorLevel: \"error\",\n\tFatalLevel: \"fatal\",\n}\n\nvar levelStrings = map[string]Level{\n\t\"debug\":   DebugLevel,\n\t\"info\":    InfoLevel,\n\t\"warn\":    WarnLevel,\n\t\"warning\": WarnLevel,\n\t\"error\":   ErrorLevel,\n\t\"fatal\":   FatalLevel,\n}\n\n// String implementation.\nfunc (l Level) String() string {\n\treturn levelNames[l]\n}\n\n// ParseLevel parses level string.\nfunc ParseLevel(s string) (Level, error) {\n\tl, ok := levelStrings[strings.ToLower(s)]\n\tif !ok {\n\t\treturn InvalidLevel, ErrInvalidLevel\n\t}\n\n\treturn l, nil\n}\n\n// MustParseLevel parses level string or panics.\nfunc MustParseLevel(s string) Level {\n\tl, err := ParseLevel(s)\n\tif err != nil {\n\t\tpanic(\"invalid log level\")\n\t}\n\n\treturn l\n}\n"
  },
  {
    "path": "pkg/util/console/term.go",
    "content": "package console\n\nimport (\n\t\"os\"\n\n\t\"github.com/moby/term\"\n)\n\n// IsTerminal returns true if we're in a terminal and a user is interacting with us\nfunc IsTerminal() bool {\n\treturn term.IsTerminal(os.Stdin.Fd())\n}\n\n// GetWidth returns the width of the terminal (from stderr -- stdout might be piped)\n//\n// Returns 0 if we're not in a terminal\nfunc GetWidth() (uint16, error) {\n\tfd := os.Stderr.Fd()\n\tif term.IsTerminal(fd) {\n\t\tws, err := term.GetWinsize(fd)\n\t\tif err != nil {\n\t\t\treturn 0, err\n\t\t}\n\t\treturn ws.Width, nil\n\t}\n\treturn 0, nil\n}\n"
  },
  {
    "path": "pkg/util/env.go",
    "content": "package util\n\nimport (\n\t\"os\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\n// GetEnvOrDefault returns an environment variable or a default if either the environment variable\n// does not exist or fails to parse using the specified conversionFunc function\nfunc GetEnvOrDefault[T any](key string, defaultVal T, conversionFunc func(string) (T, error)) T {\n\tval, exists := os.LookupEnv(key)\n\tif exists {\n\t\tv, err := conversionFunc(val)\n\t\tif err == nil {\n\t\t\treturn v\n\t\t} else {\n\t\t\tconsole.Warnf(\"Failed to convert env var %s to expected type. Continuing with default. Error: %v\", key, err)\n\t\t}\n\t}\n\treturn defaultVal\n}\n"
  },
  {
    "path": "pkg/util/errors.go",
    "content": "package util\n\nimport \"fmt\"\n\n// WrapError is just a shortcut for using fmt.Errorf\n// to wrap an error with a message\nfunc WrapError(err error, message string) error {\n\tif err == nil {\n\t\treturn nil\n\t}\n\treturn fmt.Errorf(\"%s: %w\", message, err)\n}\n"
  },
  {
    "path": "pkg/util/files/files.go",
    "content": "package files\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n\t\"path\"\n\t\"strings\"\n\n\t\"github.com/mitchellh/go-homedir\"\n\t\"github.com/vincent-petithory/dataurl\"\n\t\"golang.org/x/sys/unix\"\n\n\tr8_path \"github.com/replicate/cog/pkg/path\"\n\t\"github.com/replicate/cog/pkg/util/mime\"\n)\n\nvar (\n\tErrorFailedToSplitDataURL = errors.New(\"Failed to split data URL into 2 parts\")\n)\n\nfunc Exists(path string) (bool, error) {\n\tif _, err := os.Stat(path); err == nil {\n\t\treturn true, nil\n\t} else if os.IsNotExist(err) {\n\t\treturn false, nil\n\t} else {\n\t\treturn false, fmt.Errorf(\"Failed to determine if %s exists: %w\", path, err)\n\t}\n}\n\nfunc IsEmpty(path string) (bool, error) {\n\tentries, err := os.ReadDir(path)\n\tif err != nil {\n\t\tif errors.Is(err, os.ErrNotExist) {\n\t\t\treturn true, nil\n\t\t}\n\t\treturn false, err\n\t}\n\treturn len(entries) == 0, nil\n}\n\nfunc IsDir(path string) (bool, error) {\n\tfile, err := os.Stat(path)\n\tif err != nil {\n\t\treturn false, err\n\t}\n\treturn file.Mode().IsDir(), nil\n}\n\nfunc IsExecutable(path string) bool {\n\treturn unix.Access(path, unix.X_OK) == nil\n}\n\nfunc CopyFile(src string, dest string) error {\n\tin, err := os.Open(src)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Failed to open %s while copying to %s: %w\", src, dest, err)\n\t}\n\tdefer in.Close()\n\n\tout, err := os.Create(dest)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Failed to create %s while copying %s: %w\", dest, src, err)\n\t}\n\tdefer out.Close()\n\n\t_, err = io.Copy(out, in)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"Failed to copy %s to %s: %w\", src, dest, err)\n\t}\n\treturn out.Close()\n}\n\nfunc WriteIfDifferent(file, content string) error {\n\tif _, err := os.Stat(file); err == nil {\n\t\tbs, err := os.ReadFile(file)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tif string(bs) == content {\n\t\t\treturn nil\n\t\t}\n\t} else if !errors.Is(err, os.ErrNotExist) {\n\t\treturn err\n\t}\n\n\t// Write out a new requirements file\n\terr := os.WriteFile(file, []byte(content), 0o644)\n\tif err != nil {\n\t\treturn err\n\t}\n\treturn nil\n}\n\nfunc WriteDataURLToFile(url string, destination string) (string, error) {\n\tif strings.HasPrefix(url, \"data:None;base64\") {\n\t\turl = strings.Replace(url, \"data:None;base64\", \"data:;base64\", 1)\n\t}\n\tdataurlObj, err := dataurl.DecodeString(url)\n\tif err != nil {\n\t\t// Attempt to fallback to binary base64 file decode.\n\t\tparts := strings.SplitN(url, \",\", 2)\n\t\tif len(parts) != 2 {\n\t\t\treturn \"\", ErrorFailedToSplitDataURL\n\t\t}\n\t\tbase64Data := parts[1]\n\t\turl = \"data:;base64,\" + base64Data\n\t\tdataurlObj, err = dataurl.DecodeString(url)\n\t\tif err != nil {\n\t\t\treturn \"\", fmt.Errorf(\"Failed to decode data URL: %w\", err)\n\t\t}\n\t}\n\toutput := dataurlObj.Data\n\n\text := path.Ext(destination)\n\tdir := path.Dir(destination)\n\tname := r8_path.TrimExt(path.Base(destination))\n\n\t// Check if ext is an integer, in which case ignore it...\n\tif r8_path.IsExtInteger(ext) {\n\t\text = \"\"\n\t\tname = path.Base(destination)\n\t}\n\n\tif ext == \"\" {\n\t\text = mime.ExtensionByType(dataurlObj.ContentType())\n\t}\n\n\tpath, err := WriteFile(output, path.Join(dir, name+ext))\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn path, nil\n}\n\nfunc WriteFile(output []byte, outputPath string) (string, error) {\n\toutputPath, err := homedir.Expand(outputPath)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\t// Write to file\n\toutFile, err := os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o755)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tif _, err := outFile.Write(output); err != nil {\n\t\treturn \"\", err\n\t}\n\tif err := outFile.Close(); err != nil {\n\t\treturn \"\", err\n\t}\n\treturn outputPath, nil\n}\n"
  },
  {
    "path": "pkg/util/files/files_test.go",
    "content": "package files\n\nimport (\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestIsExecutable(t *testing.T) {\n\tdir := t.TempDir()\n\tpath := filepath.Join(dir, \"test-file\")\n\terr := os.WriteFile(path, []byte{}, 0o644)\n\trequire.NoError(t, err)\n\n\trequire.False(t, IsExecutable(path))\n\trequire.NoError(t, os.Chmod(path, 0o744))\n\trequire.True(t, IsExecutable(path))\n}\n\nfunc TestWriteBadlyFormattedBase64DataURI(t *testing.T) {\n\tdir := t.TempDir()\n\tpath := filepath.Join(dir, \"test-file\")\n\t_, err := WriteDataURLToFile(\"data:None;base64,SGVsbG8gVGhlcmU=\", path)\n\trequire.NoError(t, err)\n}\n\nfunc TestWriteNotRecognisedBase64DataURL(t *testing.T) {\n\tdir := t.TempDir()\n\tpath := filepath.Join(dir, \"test-file\")\n\t_, err := WriteDataURLToFile(\"data:None;model/gltf-binary,SGVsbG8gVGhlcmU=\", path)\n\trequire.NoError(t, err)\n}\n"
  },
  {
    "path": "pkg/util/hash.go",
    "content": "package util\n\nimport (\n\t\"bytes\"\n\t\"crypto/sha256\"\n\t\"encoding/hex\"\n\t\"errors\"\n\t\"fmt\"\n\t\"io\"\n\t\"os\"\n)\n\nvar (\n\tErrInvalidRange = errors.New(\"Invalid byte range provided for file\")\n)\n\nfunc SHA256HashFile(path string) (string, error) {\n\thash := sha256.New()\n\n\tfile, err := os.Open(path)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tdefer file.Close()\n\n\tif _, err := io.Copy(hash, file); err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn hex.EncodeToString(hash.Sum(nil)), nil\n}\n\nfunc SHA256HashFileWithSaltAndRange(path string, start int, end int, salt string) (string, error) {\n\thash := sha256.New()\n\tlength := end - start\n\n\tif length < 0 {\n\t\treturn \"\", ErrInvalidRange\n\t}\n\n\tfile, err := os.Open(path)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tdefer file.Close()\n\n\tfileInfo, err := file.Stat()\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tif fileInfo.Size() < int64(end) {\n\t\treturn \"\", ErrInvalidRange\n\t}\n\n\t_, err = file.Seek(int64(start), 0)\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"failed to open file pointer %s: %w\", path, err)\n\t}\n\tbuf := make([]byte, length)\n\tn, err := file.Read(buf)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\tbuf = buf[:n]\n\tvar hashInput []byte\n\thashInput = append(hashInput, buf...)\n\thashInput = append(hashInput, []byte(salt)...)\n\n\tif _, err := io.Copy(hash, bytes.NewReader(hashInput)); err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn hex.EncodeToString(hash.Sum(nil)), nil\n}\n"
  },
  {
    "path": "pkg/util/hash_test.go",
    "content": "package util\n\nimport (\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestHash(t *testing.T) {\n\tdir := t.TempDir()\n\tpath := filepath.Join(dir, \"test.tmp\")\n\td1 := []byte(\"hello\\ngo\\n\")\n\terr := os.WriteFile(path, d1, 0o644)\n\trequire.NoError(t, err)\n\n\tsha256, err := SHA256HashFile(path)\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"43d250d92b5dbb47f75208de8e9a9a321d23e85eed0dc3d5dfa83bc3cc5aa68c\", sha256)\n}\n\nfunc TestHashFileWithSaltAndRange(t *testing.T) {\n\tdir := t.TempDir()\n\tpath := filepath.Join(dir, \"test.tmp\")\n\td1 := []byte(\"hello\\nreplicate\\nhello\\n\")\n\terr := os.WriteFile(path, d1, 0o644)\n\trequire.NoError(t, err)\n\n\t_, err = SHA256HashFileWithSaltAndRange(path, 0, 60, \"go\\n\")\n\trequire.Error(t, err)\n\n\t_, err = SHA256HashFileWithSaltAndRange(path, 23, 1, \"go\\n\")\n\trequire.Error(t, err)\n\n\tsha256, err := SHA256HashFileWithSaltAndRange(path, 0, 6, \"go\\n\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"43d250d92b5dbb47f75208de8e9a9a321d23e85eed0dc3d5dfa83bc3cc5aa68c\", sha256)\n\n\tsha256, err = SHA256HashFileWithSaltAndRange(path, 16, 22, \"go\\n\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"43d250d92b5dbb47f75208de8e9a9a321d23e85eed0dc3d5dfa83bc3cc5aa68c\", sha256)\n}\n"
  },
  {
    "path": "pkg/util/mime/mime.go",
    "content": "package mime\n\nimport (\n\t\"mime\"\n\t\"strings\"\n)\n\nvar typeToExtension = map[string]string{\n\t\"application/epub+zip\":                            \".epub\",\n\t\"application/gzip\":                                \".gz\",\n\t\"application/java-archive\":                        \".jar\",\n\t\"application/json\":                                \".json\",\n\t\"application/jsonl\":                               \".jsonl\",\n\t\"application/ld+json\":                             \".jsonld\",\n\t\"application/msword\":                              \".doc\",\n\t\"application/octet-stream\":                        \".bin\",\n\t\"application/ogg\":                                 \".ogx\",\n\t\"application/pdf\":                                 \".pdf\",\n\t\"application/rtf\":                                 \".rtf\",\n\t\"application/vnd.amazon.ebook\":                    \".azw\",\n\t\"application/vnd.apple.installer+xml\":             \".mpkg\",\n\t\"application/vnd.ms-excel\":                        \".xls\",\n\t\"application/vnd.ms-fontobject\":                   \".eot\",\n\t\"application/vnd.ms-powerpoint\":                   \".ppt\",\n\t\"application/vnd.oasis.opendocument.presentation\": \".odp\",\n\t\"application/vnd.oasis.opendocument.spreadsheet\":  \".ods\",\n\t\"application/vnd.oasis.opendocument.text\":         \".odt\",\n\t\"application/vnd.openxmlformats-officedocument.presentationml.presentation\": \".pptx\",\n\t\"application/vnd.openxmlformats-officedocument.wordprocessingml.document\":   \".docx\",\n\t\"application/vnd.rar\":           \".rar\",\n\t\"application/vnd.visio\":         \".vsd\",\n\t\"application/x-7z-compressed\":   \".7z\",\n\t\"application/x-abiword\":         \".abw\",\n\t\"application/x-bzip\":            \".bz\",\n\t\"application/x-bzip2\":           \".bz2\",\n\t\"application/x-cdf\":             \".cda\",\n\t\"application/x-csh\":             \".csh\",\n\t\"application/x-freearc\":         \".arc\",\n\t\"application/x-httpd-php\":       \".php\",\n\t\"application/x-ndjson\":          \".ndjson\",\n\t\"application/x-sh\":              \".sh\",\n\t\"application/x-shockwave-flash\": \".swf\",\n\t\"application/x-tar\":             \".tar\",\n\t\"application/xhtml+xml\":         \".xhtml\",\n\t\"application/xml\":               \".xml\",\n\t\"application/zip\":               \".zip\",\n\n\t\"audio/aac\":               \".aac\",\n\t\"audio/midi audio/x-midi\": \".midi\",\n\t\"audio/mpeg\":              \".mp3\",\n\t\"audio/ogg\":               \".oga\",\n\t\"audio/opus\":              \".opus\",\n\t\"audio/wav\":               \".wav\",\n\t\"audio/webm\":              \".weba\",\n\n\t\"font/otf\":   \".otf\",\n\t\"font/ttf\":   \".ttf\",\n\t\"font/woff\":  \".woff\",\n\t\"font/woff2\": \".woff2\",\n\n\t\"image/bmp\":                \".bmp\",\n\t\"image/x-ms-bmp\":           \".bmp\",\n\t\"image/gif\":                \".gif\",\n\t\"image/jpeg\":               \".jpg\",\n\t\"image/png\":                \".png\",\n\t\"image/svg+xml\":            \".svg\",\n\t\"image/tiff\":               \".tiff\",\n\t\"image/vnd.microsoft.icon\": \".ico\",\n\t\"image/webp\":               \".webp\",\n\n\t\"model/gltf-binary\": \".glb\",\n\t\"model/mtl\":         \".mtl\",\n\t\"model/obj\":         \".obj\",\n\n\t\"text/calendar\":   \".ics\",\n\t\"text/css\":        \".css\",\n\t\"text/csv\":        \".csv\",\n\t\"text/html\":       \".html\",\n\t\"text/javascript\": \".js\",\n\t\"text/markdown\":   \".md\",\n\t\"text/plain\":      \".txt\",\n\n\t\"video/3gpp\":      \".3gp\",\n\t\"video/3gpp2\":     \".3gp2\",\n\t\"video/mp2t\":      \".ts\",\n\t\"video/mp4\":       \".mp4\",\n\t\"video/mpeg\":      \".mpeg\",\n\t\"video/ogg\":       \".ogv\",\n\t\"video/webm\":      \".webm\",\n\t\"video/x-msvideo\": \".avi\",\n}\n\nvar extensionToType = map[string]string{}\n\nfunc init() {\n\tfor typ, ext := range typeToExtension {\n\t\textensionToType[ext] = typ\n\t}\n}\n\n// ExtensionByType returns the file extension associated with the media type typ.\n// When typ has no associated extension, ExtensionByType returns an empty string.\nfunc ExtensionByType(typ string) string {\n\t// Lookup extension from pre-defined map\n\text := typeToExtension[typ]\n\n\t// Fall back to mime.ExtensionsByType\n\tif ext == \"\" {\n\t\textensions, _ := mime.ExtensionsByType(typ)\n\t\tif len(extensions) > 0 {\n\t\t\text = extensions[0]\n\t\t}\n\t}\n\n\treturn ext\n}\n\n// TypeByExtension returns the media type associated with the file extension ext.\n// The extension ext should begin with a leading dot, as in \".json\"\n// When ext has no associated type, TypeByExtension returns \"application/octet-stream\"\nfunc TypeByExtension(ext string) string {\n\tif !strings.HasPrefix(ext, \".\") {\n\t\text = \".\" + ext\n\t}\n\n\t// Lookup type from pre-defined map\n\ttyp := extensionToType[ext]\n\n\t// Fall back to mime.TypeByExtension\n\tif typ == \"\" {\n\t\ttyp = mime.TypeByExtension(ext)\n\t}\n\n\t// Default to \"application/octet-stream\"\n\tif typ == \"\" {\n\t\ttyp = \"application/octet-stream\"\n\t}\n\n\treturn typ\n}\n"
  },
  {
    "path": "pkg/util/mime/mime_test.go",
    "content": "package mime\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestExtensionByType(t *testing.T) {\n\trequire.Equal(t, \".txt\", ExtensionByType(\"text/plain\"))\n\trequire.Equal(t, \".jpg\", ExtensionByType(\"image/jpeg\"))\n\trequire.Equal(t, \".png\", ExtensionByType(\"image/png\"))\n\trequire.Equal(t, \".obj\", ExtensionByType(\"model/obj\"))\n\trequire.Equal(t, \".json\", ExtensionByType(\"application/json\"))\n\trequire.Equal(t, \"\", ExtensionByType(\"asdfasdf\"))\n}\n\nfunc TestTypeByExtension(t *testing.T) {\n\trequire.Equal(t, \"text/plain\", TypeByExtension(\".txt\"))\n\trequire.Equal(t, \"image/jpeg\", TypeByExtension(\".jpg\"))\n\trequire.Equal(t, \"image/png\", TypeByExtension(\".png\"))\n\trequire.Equal(t, \"model/obj\", TypeByExtension(\".obj\"))\n\trequire.Equal(t, \"application/json\", TypeByExtension(\".json\"))\n\trequire.Equal(t, \"application/octet-stream\", TypeByExtension(\".asdfasdf\"))\n}\n"
  },
  {
    "path": "pkg/util/net.go",
    "content": "package util\n\nimport (\n\t\"fmt\"\n\t\"math/rand\"\n\t\"net\"\n\t\"time\"\n)\n\n// PickFreePort returns a TCP port in [min,max] that's not in use on the 127.0.0.1 interface.\n// Note that there's a small chance of a race condition when a port is considered free at the\n// time of the call, but not free when something tries to use it. This is good enough for dev\n// and test code though.\nfunc PickFreePort(minPort, maxPort int) (int, error) {\n\tif minPort < 1024 || maxPort > 99999 || minPort > maxPort {\n\t\treturn 0, fmt.Errorf(\"invalid port range\")\n\t}\n\n\trng := rand.New(rand.NewSource(time.Now().UnixNano())) // #nosec G404 - using math/rand is fine for test port selection\n\tfor range 20 {                                         // avoid infinite loops\n\t\tp := rng.Intn(maxPort-minPort+1) + minPort\n\t\tl, err := net.Listen(\"tcp\", fmt.Sprintf(\"127.0.0.1:%d\", p))\n\t\tif err == nil {\n\t\t\t_ = l.Close()\n\t\t\treturn p, nil // looks free\n\t\t}\n\t}\n\treturn 0, fmt.Errorf(\"could not find free port in range %d-%d\", minPort, maxPort)\n}\n"
  },
  {
    "path": "pkg/util/overwrite_yaml.go",
    "content": "package util\n\nimport (\n\t\"fmt\"\n\n\t\"go.yaml.in/yaml/v4\"\n)\n\nfunc OverwriteYAML(sourceYaml []byte, destinationYaml []byte) ([]byte, error) {\n\tvar sourceNode yaml.Node\n\terr := yaml.Unmarshal(sourceYaml, &sourceNode)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tvar destinationNode yaml.Node\n\terr = yaml.Unmarshal(destinationYaml, &destinationNode)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\terr = traverseAndCompare(sourceNode.Content[0], destinationNode.Content[0], \"\")\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn yaml.Marshal(&destinationNode)\n}\n\nfunc traverseAndCompare(sourceNode, destinationNode *yaml.Node, path string) error {\n\tif sourceNode.Kind != destinationNode.Kind {\n\t\treturn fmt.Errorf(\"Type mismatch at %s: %s vs %s\\n\", path, nodeKindToString(sourceNode.Kind), nodeKindToString(destinationNode.Kind))\n\t}\n\tsourceNode.LineComment = destinationNode.LineComment\n\tsourceNode.HeadComment = destinationNode.HeadComment\n\tsourceNode.FootComment = destinationNode.FootComment\n\n\tswitch sourceNode.Kind {\n\tcase yaml.ScalarNode:\n\t\tif sourceNode.Value != destinationNode.Value {\n\t\t\tdestinationNode.Value = sourceNode.Value\n\t\t}\n\n\tcase yaml.MappingNode:\n\t\tmap1 := mapNodeToMap(sourceNode)\n\t\tmap2 := mapNodeToMap(destinationNode)\n\n\t\tallKeys := getAllKeys(map1, map2)\n\n\t\tfor _, key := range allKeys {\n\t\t\tvar childPath string\n\t\t\tif path == \"\" {\n\t\t\t\tchildPath = key\n\t\t\t} else {\n\t\t\t\tchildPath = path + \".\" + key\n\t\t\t}\n\n\t\t\tsourceKVNodeChild, ok1 := map1[key]\n\t\t\tdestinationKVNodeChild, ok2 := map2[key]\n\n\t\t\tswitch {\n\t\t\tcase !ok1:\n\t\t\t\t// We need to remove this node\n\t\t\t\tNewContent := []*yaml.Node{}\n\t\t\t\tfor _, node := range destinationNode.Content {\n\t\t\t\t\tif node == destinationKVNodeChild[0] || node == destinationKVNodeChild[1] {\n\t\t\t\t\t\tcontinue\n\t\t\t\t\t}\n\t\t\t\t\tNewContent = append(NewContent, node)\n\t\t\t\t}\n\t\t\t\tdestinationNode.Content = NewContent\n\t\t\tcase !ok2:\n\t\t\t\t// We need to add this node\n\t\t\t\tdestinationNode.Content = append(destinationNode.Content, sourceKVNodeChild...)\n\t\t\tdefault:\n\t\t\t\terr := traverseAndCompare(sourceKVNodeChild[1], destinationKVNodeChild[1], childPath)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn err\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\tcase yaml.SequenceNode:\n\t\tsourceLen := len(sourceNode.Content)\n\t\tdestinationLen := len(destinationNode.Content)\n\n\t\tmaxLen := max(destinationLen, sourceLen)\n\n\t\tfor i := range maxLen {\n\t\t\tchildPath := fmt.Sprintf(\"%s[%d]\", path, i)\n\n\t\t\tif i >= destinationLen {\n\t\t\t\tdestinationNode.Content = append(destinationNode.Content, sourceNode.Content[i])\n\t\t\t} else if i < sourceLen {\n\t\t\t\terr := traverseAndCompare(sourceNode.Content[i], destinationNode.Content[i], childPath)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn err\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\treturn nil\n}\n\nfunc mapNodeToMap(node *yaml.Node) map[string][]*yaml.Node {\n\tresult := make(map[string][]*yaml.Node)\n\tfor i := 0; i < len(node.Content); i += 2 {\n\t\tkeyNode := node.Content[i]\n\t\tvalueNode := node.Content[i+1]\n\t\tresult[keyNode.Value] = []*yaml.Node{keyNode, valueNode}\n\t}\n\treturn result\n}\n\nfunc getAllKeys(map1, map2 map[string][]*yaml.Node) []string {\n\tkeys := make(map[string]bool)\n\tfor key := range map1 {\n\t\tkeys[key] = true\n\t}\n\tfor key := range map2 {\n\t\tkeys[key] = true\n\t}\n\n\tvar keyList []string\n\tfor key := range keys {\n\t\tkeyList = append(keyList, key)\n\t}\n\treturn keyList\n}\n\nfunc nodeKindToString(kind yaml.Kind) string {\n\tswitch kind {\n\tcase yaml.ScalarNode:\n\t\treturn \"Scalar\"\n\tcase yaml.MappingNode:\n\t\treturn \"Mapping\"\n\tcase yaml.SequenceNode:\n\t\treturn \"Sequence\"\n\tdefault:\n\t\treturn \"Unknown\"\n\t}\n}\n"
  },
  {
    "path": "pkg/util/overwrite_yaml_test.go",
    "content": "package util\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\n/*\nfunc TestOverwriteYAML(t *testing.T) {\n\tvar yamlData1 = `build:\n    command: \"build.sh\"\nimage: \"my-image\"\npredict: \"predict.py\"\ntrain: \"train.py\"\nconcurrency:\n    max: 5\nenvironment:\n    - \"VAR1=value1\"\n    - \"VAR2=value2\"\n`\n\n\tvar yamlData2 = `build:\n  command: \"build_new.sh\"\nimage: \"new-image\"\npredict: \"new_predict.py\"\nconcurrency:\n  max: 10\nenvironment:\n  - \"VAR1=new_value1\"\n  - \"VAR3=value3\"\n`\n\tcontent, err := OverwriteYAML([]byte(yamlData1), []byte(yamlData2))\n\trequire.NoError(t, err)\n\trequire.Equal(t, yamlData1, string(content))\n}\n*/\n\nfunc TestOverwriteYAMLWithComments(t *testing.T) {\n\tvar sourceYaml = `build:\n  command: \"build_new.sh\"\nimage: \"new-image\"\npredict: \"new_predict.py\"\nconcurrency:\n  max: 10\nenvironment:\n  - \"VAR1=new_value1\"\n  - \"VAR3=value3\"\n`\n\n\tvar destinationYaml = `# This here is a YAML Comment\nbuild:\n    command: \"build.sh\"\nimage: \"my-image\"\npredict: \"predict.py\"\ntrain: \"train.py\"\nconcurrency:\n    max: 5\nenvironment:\n    - \"VAR1=value1\"\n    - \"VAR2=value2\"\n`\n\n\texpected := `# This here is a YAML Comment\nbuild:\n    command: \"build_new.sh\"\nimage: \"new-image\"\npredict: \"new_predict.py\"\nconcurrency:\n    max: 10\nenvironment:\n    - \"VAR1=new_value1\"\n    - \"VAR3=value3\"\n`\n\n\tcontent, err := OverwriteYAML([]byte(sourceYaml), []byte(destinationYaml))\n\trequire.NoError(t, err)\n\trequire.Equal(t, expected, string(content))\n}\n\nfunc TestOverwriteYAMLWithLineComments(t *testing.T) {\n\tvar sourceYaml = `build:\n  command: \"build_new.sh\"\nimage: \"new-image\"\npredict: \"new_predict.py\"\nconcurrency:\n  max: 10\nenvironment:\n  - \"VAR1=new_value1\"\n  - \"VAR3=value3\"\n`\n\n\tvar destinationYaml = `# This here is a YAML Comment\nbuild:\n    # And we put this comment here for good measure\n    command: \"build.sh\"\nimage: \"my-image\"\npredict: \"predict.py\"\ntrain: \"train.py\"\nconcurrency:\n    max: 5\nenvironment:\n    - \"VAR1=value1\"\n    - \"VAR2=value2\"\n`\n\n\texpected := `# This here is a YAML Comment\nbuild:\n    # And we put this comment here for good measure\n    command: \"build_new.sh\"\nimage: \"new-image\"\npredict: \"new_predict.py\"\nconcurrency:\n    max: 10\nenvironment:\n    - \"VAR1=new_value1\"\n    - \"VAR3=value3\"\n`\n\tcontent, err := OverwriteYAML([]byte(sourceYaml), []byte(destinationYaml))\n\trequire.NoError(t, err)\n\trequire.Equal(t, expected, string(content))\n}\n\nfunc TestStep1XYaml(t *testing.T) {\n\tvar sourceYaml = `build:\n  gpu: true\n  system_packages:\n    - \"libgl1-mesa-glx\"\n    - \"libglib2.0-0\"\n  python_version: \"3.11\"\n  python_requirements: requirements.txt\npredict: \"predict.py:Predictor\"\n`\n\n\tvar destinationYaml = `# Configuration for Cog ⚙️\n# Reference: https://cog.run/yaml\n\nbuild:\n  # set to true if your model requires a GPU\n  gpu: true\n\n  # a list of ubuntu apt packages to install\n  system_packages:\n    - \"libgl1-mesa-glx\"\n    - \"libglib2.0-0\"\n\n  # python version in the form '3.11' or '3.11.4'\n  python_version: \"3.11\"\n\n  # path to a Python requirements.txt file\n  python_requirements: requirements.txt\n\n  # commands run after the environment is setup\n  run:\n  - curl -o /usr/local/bin/pget -L \"https://github.com/replicate/pget/releases/latest/download/pget_$(uname -s)_$(uname -m)\"\n  - chmod +x /usr/local/bin/pget\n\n# predict.py defines how predictions are run on your model\npredict: \"predict.py:Predictor\"`\n\n\texpected := `# Configuration for Cog ⚙️\n# Reference: https://cog.run/yaml\n\nbuild:\n    # set to true if your model requires a GPU\n    gpu: true\n    # a list of ubuntu apt packages to install\n    system_packages:\n        - \"libgl1-mesa-glx\"\n        - \"libglib2.0-0\"\n    # python version in the form '3.11' or '3.11.4'\n    python_version: \"3.11\"\n    # path to a Python requirements.txt file\n    python_requirements: requirements.txt\n# predict.py defines how predictions are run on your model\npredict: \"predict.py:Predictor\"\n`\n\tcontent, err := OverwriteYAML([]byte(sourceYaml), []byte(destinationYaml))\n\trequire.NoError(t, err)\n\trequire.Equal(t, expected, string(content))\n}\n"
  },
  {
    "path": "pkg/util/platform.go",
    "content": "package util\n\n// IsAppleSiliconMac returns whether the current machine is an Apple silicon computer, such as the MacBook Air with M1.\nfunc IsAppleSiliconMac(goos string, goarch string) bool {\n\treturn goos == \"darwin\" && goarch == \"arm64\"\n}\n"
  },
  {
    "path": "pkg/util/ringbuffer.go",
    "content": "package util\n\nimport (\n\t\"io\"\n\t\"sync\"\n)\n\n// RingBufferWriter is a writer that writes to an underlying writer and also maintains\n// a ring buffer of the last N bytes written.\ntype RingBufferWriter struct {\n\twriter io.Writer\n\tbuffer []byte\n\tsize   int\n\tpos    int\n\tmu     sync.Mutex\n}\n\n// NewRingBufferWriter creates a new RingBufferWriter that writes to w and maintains\n// a buffer of the last size bytes.\nfunc NewRingBufferWriter(w io.Writer, size int) *RingBufferWriter {\n\treturn &RingBufferWriter{\n\t\twriter: w,\n\t\tbuffer: make([]byte, size),\n\t\tsize:   size,\n\t}\n}\n\n// Write implements io.Writer interface\nfunc (w *RingBufferWriter) Write(p []byte) (n int, err error) {\n\tw.mu.Lock()\n\tdefer w.mu.Unlock()\n\n\t// Write to underlying writer\n\tn, err = w.writer.Write(p)\n\tif err != nil {\n\t\treturn n, err\n\t}\n\n\t// Update ring buffer\n\tfor _, b := range p {\n\t\tw.buffer[w.pos] = b\n\t\tw.pos = (w.pos + 1) % w.size\n\t}\n\n\treturn n, nil\n}\n\n// String returns the contents of the ring buffer as a string\nfunc (w *RingBufferWriter) String() string {\n\tw.mu.Lock()\n\tdefer w.mu.Unlock()\n\n\t// If buffer is not full, return what we have\n\tif w.pos < w.size {\n\t\treturn string(w.buffer[:w.pos])\n\t}\n\n\t// Otherwise, return the last size bytes\n\treturn string(w.buffer[w.pos:]) + string(w.buffer[:w.pos])\n}\n"
  },
  {
    "path": "pkg/util/shell/net.go",
    "content": "package shell\n\nimport (\n\t\"fmt\"\n\t\"net\"\n\t\"net/http\"\n\t\"strconv\"\n\t\"time\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc WaitForPort(port int, timeout time.Duration) error {\n\tstart := time.Now()\n\tfor {\n\t\tif PortIsOpen(port) {\n\t\t\treturn nil\n\t\t}\n\n\t\tnow := time.Now()\n\t\tif now.Sub(start) > timeout {\n\t\t\treturn fmt.Errorf(\"Timed out\")\n\t\t}\n\n\t\ttime.Sleep(100 * time.Millisecond)\n\t}\n}\n\nfunc WaitForHTTPOK(url string, timeout time.Duration) error {\n\tstart := time.Now()\n\tconsole.Debugf(\"Waiting for %s to become accessible\", url)\n\tfor {\n\t\tnow := time.Now()\n\t\tif now.Sub(start) > timeout {\n\t\t\treturn fmt.Errorf(\"Timed out\")\n\t\t}\n\n\t\ttime.Sleep(100 * time.Millisecond)\n\t\tresp, err := http.Get(url) //#nosec G107\n\t\tif err != nil {\n\t\t\tcontinue\n\t\t}\n\t\tif resp.StatusCode != http.StatusOK {\n\t\t\tcontinue\n\t\t}\n\t\tconsole.Debugf(\"Got successful response from %s\", url)\n\t\treturn nil\n\t}\n}\n\nfunc PortIsOpen(port int) bool {\n\tconn, err := net.DialTimeout(\"tcp\", net.JoinHostPort(\"\", strconv.Itoa(port)), 100*time.Millisecond)\n\tif conn != nil {\n\t\t_ = conn.Close()\n\t}\n\treturn err == nil\n}\n"
  },
  {
    "path": "pkg/util/shell/pipes.go",
    "content": "package shell\n\nimport (\n\t\"bufio\"\n\t\"io\"\n)\n\ntype PipeFunc func() (io.ReadCloser, error)\ntype LogFunc func(args ...any)\n\nfunc PipeTo(pf PipeFunc, lf LogFunc) (done chan struct{}, err error) {\n\tdone = make(chan struct{})\n\n\tpipe, err := pf()\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tscanner := bufio.NewScanner(pipe)\n\tgo func() {\n\t\tfor scanner.Scan() {\n\t\t\tline := scanner.Text()\n\t\t\tlf(line)\n\t\t}\n\t\tdone <- struct{}{}\n\t}()\n\n\treturn done, nil\n}\n"
  },
  {
    "path": "pkg/util/version/version.go",
    "content": "package version\n\nimport (\n\t\"fmt\"\n\t\"strconv\"\n\t\"strings\"\n)\n\ntype Version struct {\n\tMajor    int\n\tMinor    int\n\tPatch    *int\n\tMetadata string\n}\n\nfunc NewVersion(s string) (version *Version, err error) {\n\tplusParts := strings.SplitN(s, \"+\", 2)\n\tnumber := plusParts[0]\n\tparts := strings.Split(number, \".\")\n\tif len(parts) > 3 {\n\t\treturn nil, fmt.Errorf(\"Version must not have more than 3 parts: %s\", s)\n\t}\n\tversion = new(Version)\n\tversion.Major, err = strconv.Atoi(parts[0])\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Invalid major version %s: %w\", parts[0], err)\n\t}\n\tif len(parts) >= 2 {\n\t\tversion.Minor, err = strconv.Atoi(parts[1])\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"Invalid minor version %s: %w\", parts[1], err)\n\t\t}\n\t}\n\tif len(parts) >= 3 {\n\t\tpatch, err := strconv.Atoi(parts[2])\n\t\tif err != nil {\n\t\t\treturn nil, fmt.Errorf(\"Invalid patch version %s: %w\", parts[2], err)\n\t\t}\n\t\t// We assign a pointer here to handle cases where the patch version is not\n\t\t// explicitly assigned and we need to compare versions without patches to\n\t\t// versions with patches.\n\t\tversion.Patch = new(int)\n\t\t*version.Patch = patch\n\t}\n\n\tif len(plusParts) == 2 {\n\t\tversion.Metadata = plusParts[1]\n\t}\n\n\treturn version, nil\n}\n\nfunc MustVersion(s string) *Version {\n\tversion, err := NewVersion(s)\n\tif err != nil {\n\t\tpanic(fmt.Sprintf(\"%s\", err))\n\t}\n\treturn version\n}\n\nfunc (v *Version) Greater(other *Version) bool {\n\tswitch {\n\tcase v.Major > other.Major:\n\t\treturn true\n\tcase v.Major == other.Major && v.Minor > other.Minor:\n\t\treturn true\n\tcase v.Major == other.Major &&\n\t\tv.Minor == other.Minor &&\n\t\tv.PatchVersion() > other.PatchVersion():\n\t\treturn true\n\tdefault:\n\t\treturn false\n\t}\n}\n\nfunc (v *Version) Equal(other *Version) bool {\n\treturn v.Major == other.Major &&\n\t\tv.Minor == other.Minor &&\n\t\tv.PatchVersion() == other.PatchVersion() &&\n\t\tv.Metadata == other.Metadata\n}\n\nfunc (v *Version) GreaterOrEqual(other *Version) bool {\n\treturn v.Greater(other) || v.Equal(other)\n}\n\nfunc (v *Version) EqualMinor(other *Version) bool {\n\treturn v.Major == other.Major && v.Minor == other.Minor\n}\n\nfunc (v *Version) HasPatch() bool {\n\treturn v.Patch != nil\n}\n\nfunc (v *Version) PatchVersion() int {\n\tif v.Patch == nil {\n\t\treturn 0\n\t}\n\treturn *v.Patch\n}\n\nfunc Equal(v1 string, v2 string) bool {\n\treturn MustVersion(v1).Equal(MustVersion(v2))\n}\n\nfunc EqualMinor(v1 string, v2 string) bool {\n\treturn MustVersion(v1).EqualMinor(MustVersion(v2))\n}\n\nfunc Greater(v1 string, v2 string) bool {\n\treturn MustVersion(v1).Greater(MustVersion(v2))\n}\n\nfunc GreaterOrEqual(v1 string, v2 string) bool {\n\tleftVersion, err := NewVersion(v1)\n\tif err != nil {\n\t\treturn v1 == v2\n\t}\n\trightVersion, err := NewVersion(v2)\n\tif err != nil {\n\t\treturn v1 == v2\n\t}\n\treturn leftVersion.GreaterOrEqual(rightVersion)\n}\n\nfunc (v *Version) Matches(other *Version) bool {\n\tswitch {\n\tcase v.Major != other.Major:\n\t\treturn false\n\tcase v.Minor != other.Minor:\n\t\treturn false\n\tcase v.HasPatch() && other.HasPatch() && *v.Patch != *other.Patch:\n\t\treturn false\n\tdefault:\n\t\treturn true\n\t}\n}\n\nfunc Matches(v1 string, v2 string) bool {\n\treturn MustVersion(v1).Matches(MustVersion(v2))\n}\n\nfunc StripPatch(v string) string {\n\tver := MustVersion(v)\n\treturn fmt.Sprintf(\"%d.%d\", ver.Major, ver.Minor)\n}\n\nfunc StripModifier(v string) string {\n\tmodifierSplit := strings.Split(v, \"+\")\n\treturn modifierSplit[0]\n}\n"
  },
  {
    "path": "pkg/util/version/version_test.go",
    "content": "package version\n\nimport (\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestVersionEqual(t *testing.T) {\n\tfor _, tt := range []struct {\n\t\tv1    string\n\t\tv2    string\n\t\tequal bool\n\t}{\n\t\t{\"1\", \"1\", true},\n\t\t{\"1.0\", \"1\", true},\n\t\t{\"1\", \"1.0\", true},\n\t\t{\"1.0.0\", \"1\", true},\n\t\t{\"1.0.0\", \"1.0\", true},\n\t\t{\"1.0.0\", \"1.0.0\", true},\n\t\t{\"1.0.0+foo\", \"1.0.0\", false},\n\t\t{\"11.2\", \"11.2.0\", true},\n\t\t{\"1\", \"2\", false},\n\t\t{\"1\", \"0\", false},\n\t\t{\"1.1\", \"1\", false},\n\t\t{\"1.0.1\", \"1\", false},\n\t\t{\"1.1.0\", \"1\", false},\n\t} {\n\t\tnot := \"\"\n\t\tif tt.equal {\n\t\t\tnot = \"not \"\n\t\t}\n\t\trequire.Equal(t, tt.equal, Equal(tt.v1, tt.v2), \"%s is %sequal to %s\", tt.v1, not, tt.v2)\n\t}\n}\n\nfunc TestVersionGreater(t *testing.T) {\n\tfor _, tt := range []struct {\n\t\tv1      string\n\t\tv2      string\n\t\tgreater bool\n\t}{\n\t\t{\"1\", \"1\", false},\n\t\t{\"1.0\", \"1\", false},\n\t\t{\"1\", \"1.0\", false},\n\t\t{\"1.0.0\", \"1\", false},\n\t\t{\"1.0.0\", \"1.0\", false},\n\t\t{\"11.2\", \"11.2.0\", false},\n\t\t{\"1\", \"2\", false},\n\t\t{\"1\", \"0\", true},\n\t\t{\"1.1\", \"1\", true},\n\t\t{\"1.0.1\", \"1\", true},\n\t\t{\"1.1.0\", \"1\", true},\n\t\t{\"1.0.0+foo\", \"1\", false},\n\t} {\n\t\tnot := \"\"\n\t\tif tt.greater {\n\t\t\tnot = \"not \"\n\t\t}\n\t\trequire.Equal(t, tt.greater, Greater(tt.v1, tt.v2), \"%s is %sgreater than %s\", tt.v1, not, tt.v2)\n\t}\n}\n\nfunc TestVersionStripModifier(t *testing.T) {\n\tversion := \"2.3.1\"\n\tversionWithModifier := version + \"+cu118\"\n\tversionWithoutModifier := StripModifier(versionWithModifier)\n\trequire.Equal(t, versionWithoutModifier, version)\n}\n\nfunc TestVersionMatches(t *testing.T) {\n\tversion := \"2.3\"\n\tmatchVersion := \"2.3.2\"\n\trequire.True(t, Matches(version, matchVersion))\n}\n\nfunc TestVersionMatchesModifier(t *testing.T) {\n\tversion := \"2.3\"\n\tmatchVersion := \"2.3.2+cu118\"\n\trequire.True(t, Matches(version, matchVersion))\n}\n\nfunc TestGreaterThanOrEqualToWithInvalidPatch(t *testing.T) {\n\tleftVersion := \"1.1.0b2\"\n\trightVersion := \"1.1.0b2\"\n\trequire.True(t, GreaterOrEqual(leftVersion, rightVersion))\n}\n"
  },
  {
    "path": "pkg/web/client.go",
    "content": "package web\n\nimport (\n\t\"bytes\"\n\t\"context\"\n\t\"encoding/json\"\n\t\"errors\"\n\t\"fmt\"\n\t\"net/http\"\n\t\"net/url\"\n\t\"os\"\n\t\"strconv\"\n\t\"strings\"\n\t\"sync\"\n\t\"time\"\n\n\t\"github.com/docker/docker/api/types/image\"\n\t\"github.com/replicate/go/types\"\n\t\"golang.org/x/sync/errgroup\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/command\"\n\t\"github.com/replicate/cog/pkg/env\"\n\tr8_errors \"github.com/replicate/cog/pkg/errors\"\n\t\"github.com/replicate/cog/pkg/global\"\n\t\"github.com/replicate/cog/pkg/util\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nconst (\n\tpushStartURLPath      = \"/api/models/push-start\"\n\tstartChallengeURLPath = \"/api/models/file-challenge\"\n)\n\nvar (\n\tErrorBadResponseNewVersionEndpoint        = errors.New(\"Bad response from new version endpoint\")\n\tErrorBadResponsePushStartEndpoint         = errors.New(\"Bad response from push start endpoint\")\n\tErrorBadResponseInitiateChallengeEndpoint = errors.New(\"Bad response from start file challenge endpoint\")\n\tErrorNoSuchDigest                         = errors.New(\"No digest submitted matches the digest requested\")\n)\n\ntype Client struct {\n\tdockerCommand command.Command\n\tclient        *http.Client\n}\n\ntype File struct {\n\tPath   string `json:\"path\"`\n\tDigest string `json:\"digest\"`\n\tSize   int64  `json:\"size\"`\n}\n\ntype Env struct {\n\tCogGpu              string `json:\"COG_GPU\"`\n\tCogPredictTypeStub  string `json:\"COG_PREDICT_TYPE_STUB\"`\n\tCogTrainTypeStub    string `json:\"COG_TRAIN_TYPE_STUB\"`\n\tCogPredictCodeStrip string `json:\"COG_PREDICT_CODE_STRIP\"`\n\tCogTrainCodeStrip   string `json:\"COG_TRAIN_CODE_STRIP\"`\n\tR8CogVersion        string `json:\"R8_COG_VERSION\"`\n\tR8CudaVersion       string `json:\"R8_CUDA_VERSION\"`\n\tR8CudnnVersion      string `json:\"R8_CUDNN_VERSION\"`\n\tR8PythonVersion     string `json:\"R8_PYTHON_VERSION\"`\n\tR8TorchVersion      string `json:\"R8_TORCH_VERSION\"`\n}\n\ntype RuntimeConfig struct {\n\tWeights []File `json:\"weights\"`\n\tFiles   []File `json:\"files\"`\n\tEnv     Env    `json:\"env\"`\n}\n\ntype Version struct {\n\tAnnotations   map[string]string     `json:\"annotations\"`\n\tCogConfig     config.Config         `json:\"cog_config\"`\n\tCogVersion    string                `json:\"cog_version\"`\n\tOpenAPISchema map[string]any        `json:\"openapi_schema\"`\n\tRuntimeConfig RuntimeConfig         `json:\"runtime_config\"`\n\tVirtual       bool                  `json:\"virtual\"`\n\tPushID        string                `json:\"push_id\"`\n\tChallenges    []FileChallengeAnswer `json:\"file_challenges\"`\n}\n\ntype FileChallengeRequest struct {\n\tDigest   string `json:\"digest\"`\n\tFileType string `json:\"file_type\"`\n}\n\ntype FileChallenge struct {\n\tSalt   string `json:\"salt\"`\n\tStart  int    `json:\"byte_start\"`\n\tEnd    int    `json:\"byte_end\"`\n\tDigest string `json:\"digest\"`\n\tID     string `json:\"challenge_id\"`\n}\n\ntype FileChallengeAnswer struct {\n\tDigest      string `json:\"digest\"`\n\tHash        string `json:\"hash\"`\n\tChallengeID string `json:\"challenge_id\"`\n}\n\ntype VersionError struct {\n\tDetail  string `json:\"detail\"`\n\tPointer string `json:\"pointer\"`\n}\n\ntype VersionErrors struct {\n\tDetail string         `json:\"detail\"`\n\tErrors []VersionError `json:\"errors\"`\n\tStatus int            `json:\"status\"`\n\tTitle  string         `json:\"title\"`\n}\n\ntype VersionCreate struct {\n\tVersion string `json:\"version\"`\n}\n\ntype CogKey struct {\n\tKey       string `json:\"key\"`\n\tExpiresAt string `json:\"expires_at\"`\n}\n\ntype Keys struct {\n\tCog CogKey `json:\"cog\"`\n}\n\ntype TokenData struct {\n\tKeys Keys `json:\"keys\"`\n}\n\nfunc NewClient(dockerCommand command.Command, client *http.Client) *Client {\n\treturn &Client{\n\t\tdockerCommand: dockerCommand,\n\t\tclient:        client,\n\t}\n}\n\nfunc (c *Client) PostPushStart(ctx context.Context, pushID string, buildTime time.Duration) error {\n\tjsonBody := map[string]any{\n\t\t\"push_id\":         pushID,\n\t\t\"build_duration\":  types.Duration(buildTime).String(),\n\t\t\"push_start_time\": time.Now().UTC(),\n\t}\n\n\tjsonData, err := json.Marshal(jsonBody)\n\tif err != nil {\n\t\treturn util.WrapError(err, \"failed to marshal JSON for build start\")\n\t}\n\n\turl := webBaseURL()\n\turl.Path = pushStartURLPath\n\n\treq, err := http.NewRequestWithContext(ctx, http.MethodPost, url.String(), bytes.NewReader(jsonData))\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tresp, err := c.client.Do(req) //nolint:gosec // G704: URL from configured endpoint\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer resp.Body.Close()\n\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn util.WrapError(ErrorBadResponsePushStartEndpoint, strconv.Itoa(resp.StatusCode))\n\t}\n\n\treturn nil\n}\n\nfunc (c *Client) PostNewVersion(ctx context.Context, image string, weights []File, files []File, fileChallenges []FileChallengeAnswer) error {\n\tversion, err := c.versionFromManifest(ctx, image, weights, files, fileChallenges)\n\tif err != nil {\n\t\treturn util.WrapError(err, \"failed to build new version from manifest\")\n\t}\n\n\tjsonData, err := json.Marshal(version)\n\tif err != nil {\n\t\treturn util.WrapError(err, \"failed to marshal JSON for new version\")\n\t}\n\n\tversionUrl, err := newVersionURL(image)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\treq, err := http.NewRequestWithContext(ctx, http.MethodPost, versionUrl.String(), bytes.NewReader(jsonData))\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tresp, err := c.client.Do(req) //nolint:gosec // G704: URL from configured endpoint\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer resp.Body.Close()\n\tdecoder := json.NewDecoder(resp.Body)\n\n\tif resp.StatusCode != http.StatusCreated {\n\t\tif resp.StatusCode == http.StatusBadRequest {\n\t\t\tvar versionErrors VersionErrors\n\t\t\terr = decoder.Decode(&versionErrors)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\t\t\treturn errors.New(versionErrors.Detail)\n\t\t}\n\t\treturn util.WrapError(ErrorBadResponseNewVersionEndpoint, strconv.Itoa(resp.StatusCode))\n\t}\n\n\tvar versionCreate VersionCreate\n\terr = decoder.Decode(&versionCreate)\n\tif err != nil {\n\t\treturn err\n\t}\n\tconsole.Infof(\"New Version: %s\", versionCreate.Version)\n\n\treturn nil\n}\n\nfunc (c *Client) FetchAPIToken(ctx context.Context, entity string) (string, error) {\n\ttokenUrl := tokenURL(entity)\n\treq, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenUrl.String(), nil)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\ttokenResp, err := c.client.Do(req) //nolint:gosec // G704: URL from configured endpoint\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\tdefer tokenResp.Body.Close()\n\n\tif tokenResp.StatusCode != http.StatusOK {\n\t\treturn \"\", fmt.Errorf(\"Bad response: %s attempting to exchange tokens\", strconv.Itoa(tokenResp.StatusCode))\n\t}\n\n\tvar tokenData TokenData\n\terr = json.NewDecoder(tokenResp.Body).Decode(&tokenData)\n\tif err != nil {\n\t\treturn \"\", err\n\t}\n\n\treturn tokenData.Keys.Cog.Key, nil\n}\n\nfunc (c *Client) versionFromManifest(ctx context.Context, image string, weights []File, files []File, fileChallenges []FileChallengeAnswer) (*Version, error) {\n\tmanifest, err := c.dockerCommand.Inspect(ctx, image)\n\tif err != nil {\n\t\treturn nil, util.WrapError(err, \"failed to inspect docker image\")\n\t}\n\n\tcogConfig, err := readCogConfig(manifest)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tvar openAPISchema map[string]any\n\terr = json.Unmarshal([]byte(manifest.Config.Labels[command.CogOpenAPISchemaLabelKey]), &openAPISchema)\n\tif err != nil {\n\t\treturn nil, util.WrapError(err, \"failed to get OpenAPI schema from docker image\")\n\t}\n\n\tpredictCode, err := stripCodeFromStub(cogConfig, true)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\ttrainCode, err := stripCodeFromStub(cogConfig, false)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tvar cogGPU int\n\tif cogConfig.Build.GPU {\n\t\tcogGPU = 1\n\t}\n\n\tcogVersion := \"\"\n\ttorchVersion := \"\"\n\tcudaVersion := \"\"\n\tcudnnVersion := \"\"\n\tpythonVersion := \"\"\n\tfor _, env := range manifest.Config.Env {\n\t\tenvName, envValue, found := strings.Cut(env, \"=\")\n\t\tif !found {\n\t\t\tcontinue\n\t\t}\n\t\tswitch envName {\n\t\tcase command.R8CogVersionEnvVarName:\n\t\t\tcogVersion = envValue\n\t\tcase command.R8TorchVersionEnvVarName:\n\t\t\ttorchVersion = envValue\n\t\tcase command.R8CudaVersionEnvVarName:\n\t\t\tcudaVersion = envValue\n\t\tcase command.R8CudnnVersionEnvVarName:\n\t\t\tcudnnVersion = envValue\n\t\tcase command.R8PythonVersionEnvVarName:\n\t\t\tpythonVersion = envValue\n\t\t}\n\t}\n\n\tenv := Env{\n\t\tCogGpu:              strconv.Itoa(cogGPU),\n\t\tCogPredictTypeStub:  cogConfig.Predict,\n\t\tCogTrainTypeStub:    cogConfig.Train,\n\t\tCogPredictCodeStrip: predictCode,\n\t\tCogTrainCodeStrip:   trainCode,\n\t\tR8CogVersion:        cogVersion,\n\t\tR8CudaVersion:       cudaVersion,\n\t\tR8CudnnVersion:      cudnnVersion,\n\t\tR8PythonVersion:     pythonVersion,\n\t\tR8TorchVersion:      torchVersion,\n\t}\n\n\tprefixedFiles := make([]File, len(files))\n\n\tfor i, file := range files {\n\t\tprefixedFiles[i] = File{\n\t\t\tPath:   file.Path,\n\t\t\tDigest: \"sha256:\" + file.Digest,\n\t\t\tSize:   file.Size,\n\t\t}\n\t}\n\n\tprefixedWeights := make([]File, len(weights))\n\n\tfor i, file := range weights {\n\t\tprefixedWeights[i] = File{\n\t\t\tPath:   file.Path,\n\t\t\tDigest: \"sha256:\" + file.Digest,\n\t\t\tSize:   file.Size,\n\t\t}\n\t}\n\n\t// Digests should match whatever digest we are sending in as the\n\t// runtime config digests\n\tfor i, challenge := range fileChallenges {\n\t\tfileChallenges[i] = FileChallengeAnswer{\n\t\t\tDigest:      fmt.Sprintf(\"sha256:%s\", challenge.Digest),\n\t\t\tHash:        challenge.Hash,\n\t\t\tChallengeID: challenge.ChallengeID,\n\t\t}\n\t}\n\n\truntimeConfig := RuntimeConfig{\n\t\tWeights: prefixedWeights,\n\t\tFiles:   prefixedFiles,\n\t\tEnv:     env,\n\t}\n\n\tversion := Version{\n\t\tAnnotations:   manifest.Config.Labels,\n\t\tCogConfig:     *cogConfig,\n\t\tCogVersion:    manifest.Config.Labels[command.CogVersionLabelKey],\n\t\tOpenAPISchema: openAPISchema,\n\t\tRuntimeConfig: runtimeConfig,\n\t\tVirtual:       true,\n\t\tChallenges:    fileChallenges,\n\t}\n\n\tif pushID, ok := manifest.Config.Labels[\"run.cog.push_id\"]; ok {\n\t\tversion.PushID = pushID\n\t}\n\n\treturn &version, nil\n}\n\nfunc (c *Client) InitiateAndDoFileChallenge(ctx context.Context, weights []File, files []File) ([]FileChallengeAnswer, error) {\n\tvar challengeAnswers []FileChallengeAnswer\n\tvar mu sync.Mutex\n\n\tvar wg errgroup.Group\n\tfor _, item := range files {\n\t\twg.Go(func() error {\n\t\t\tanswer, err := c.doSingleFileChallenge(ctx, item, \"files\")\n\t\t\tif err != nil {\n\t\t\t\treturn util.WrapError(err, fmt.Sprintf(\"do file challenge for digest %s\", item.Digest))\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tchallengeAnswers = append(challengeAnswers, answer)\n\t\t\tmu.Unlock()\n\t\t\treturn nil\n\t\t})\n\t}\n\tfor _, item := range weights {\n\t\twg.Go(func() error {\n\t\t\tanswer, err := c.doSingleFileChallenge(ctx, item, \"weights\")\n\t\t\tif err != nil {\n\t\t\t\treturn util.WrapError(err, fmt.Sprintf(\"do file challenge for digest %s\", item.Digest))\n\t\t\t}\n\t\t\tmu.Lock()\n\t\t\tchallengeAnswers = append(challengeAnswers, answer)\n\t\t\tmu.Unlock()\n\t\t\treturn nil\n\t\t})\n\t}\n\tif err := wg.Wait(); err != nil {\n\t\treturn nil, util.WrapError(err, \"do file challenges\")\n\t}\n\n\treturn challengeAnswers, nil\n}\n\n// doSingleFileChallenge does a single file challenge. This is expected to be called in a goroutine.\nfunc (c *Client) doSingleFileChallenge(ctx context.Context, file File, fileType string) (FileChallengeAnswer, error) {\n\tinitiateChallengePath := webBaseURL()\n\tinitiateChallengePath.Path = startChallengeURLPath\n\n\tanswer := FileChallengeAnswer{}\n\n\tjsonData, err := json.Marshal(FileChallengeRequest{\n\t\tDigest:   file.Digest,\n\t\tFileType: fileType,\n\t})\n\n\tif err != nil {\n\t\treturn answer, util.WrapError(err, \"encode request JSON\")\n\t}\n\n\treq, err := http.NewRequestWithContext(ctx, http.MethodPost, initiateChallengePath.String(), bytes.NewReader(jsonData))\n\tif err != nil {\n\t\treturn answer, util.WrapError(err, \"build HTTP request\")\n\t}\n\tresp, err := c.client.Do(req) //nolint:gosec // G704: URL from configured endpoint\n\tif err != nil {\n\t\treturn answer, util.WrapError(err, \"do HTTP request\")\n\t}\n\tdefer resp.Body.Close()\n\n\tif resp.StatusCode != http.StatusOK {\n\t\treturn answer, util.WrapError(ErrorBadResponseInitiateChallengeEndpoint, strconv.Itoa(resp.StatusCode))\n\t}\n\n\tvar challenge FileChallenge\n\terr = json.NewDecoder(resp.Body).Decode(&challenge)\n\tif err != nil {\n\t\treturn answer, util.WrapError(err, \"decode response body\")\n\t}\n\n\tans, err := util.SHA256HashFileWithSaltAndRange(file.Path, challenge.Start, challenge.End, challenge.Salt)\n\tif err != nil {\n\t\treturn answer, util.WrapError(err, \"hash file\")\n\t}\n\treturn FileChallengeAnswer{\n\t\tDigest:      file.Digest,\n\t\tHash:        ans,\n\t\tChallengeID: challenge.ID,\n\t}, nil\n}\n\nfunc newVersionURL(image string) (url.URL, error) {\n\timageComponents := strings.Split(image, \"/\")\n\tnewVersionUrl := webBaseURL()\n\tif len(imageComponents) != 3 || imageComponents[0] != global.ReplicateRegistryHost {\n\t\treturn newVersionUrl, r8_errors.ErrorBadRegistryURL\n\t}\n\tnewVersionUrl.Path = strings.Join([]string{\"\", \"api\", \"models\", imageComponents[1], imageComponents[2], \"versions\"}, \"/\")\n\treturn newVersionUrl, nil\n}\n\nfunc tokenURL(entity string) url.URL {\n\tnewVersionUrl := webBaseURL()\n\tnewVersionUrl.Path = strings.Join([]string{\"\", \"api\", \"token\", entity}, \"/\")\n\treturn newVersionUrl\n}\n\nfunc webBaseURL() url.URL {\n\treturn url.URL{\n\t\tScheme: env.SchemeFromEnvironment(),\n\t\tHost:   env.WebHostFromEnvironment(),\n\t}\n}\n\nfunc codeFileName(cogConfig *config.Config, isPredict bool) (string, error) {\n\tvar stubComponents []string\n\tif isPredict {\n\t\tif cogConfig.Predict == \"\" {\n\t\t\treturn \"\", nil\n\t\t}\n\t\tstubComponents = strings.Split(cogConfig.Predict, \":\")\n\t} else {\n\t\tif cogConfig.Train == \"\" {\n\t\t\treturn \"\", nil\n\t\t}\n\t\tstubComponents = strings.Split(cogConfig.Train, \":\")\n\t}\n\n\tif len(stubComponents) < 2 {\n\t\treturn \"\", errors.New(\"Code stub components has less than 2 entries.\")\n\t}\n\n\treturn stubComponents[0], nil\n}\n\nfunc readCode(cogConfig *config.Config, isPredict bool) (string, string, error) {\n\tcodeFile, err := codeFileName(cogConfig, isPredict)\n\tif err != nil {\n\t\treturn \"\", codeFile, err\n\t}\n\tif codeFile == \"\" {\n\t\treturn \"\", \"\", nil\n\t}\n\n\tb, err := os.ReadFile(codeFile)\n\tif err != nil {\n\t\treturn \"\", codeFile, err\n\t}\n\n\treturn string(b), codeFile, nil\n}\n\nfunc stripCodeFromStub(cogConfig *config.Config, isPredict bool) (string, error) {\n\t// TODO: We should attempt to strip the code here, in python this is done like so:\n\t// from cog.code_xforms import strip_model_source_code\n\t// code = strip_model_source_code(\n\t//   util.read_file(os.path.join(fs, 'src', base_file)),\n\t//   [base_class],\n\t//   ['predict', 'train'],\n\t// )\n\t// Currently the behavior of the code strip attempts to strip, and if it can't it\n\t// loads the whole file in. Here we just load the whole file in.\n\t// We should figure out a way to call cog python from here to fulfill this.\n\t// It could be a good idea to do this in the layer functions where we do pip freeze\n\t// et al.\n\n\tcode, _, err := readCode(cogConfig, isPredict)\n\treturn code, err\n}\n\nfunc readCogConfig(manifest *image.InspectResponse) (*config.Config, error) {\n\tvar cogConfig config.Config\n\terr := json.Unmarshal([]byte(manifest.Config.Labels[command.CogConfigLabelKey]), &cogConfig)\n\tif err != nil {\n\t\treturn nil, util.WrapError(err, \"failed to get cog config from docker image\")\n\t}\n\n\treturn &cogConfig, nil\n}\n"
  },
  {
    "path": "pkg/web/client_test.go",
    "content": "package web\n\nimport (\n\t\"encoding/json\"\n\t\"net/http\"\n\t\"net/http/httptest\"\n\t\"net/url\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/assert\"\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/docker/dockertest\"\n\t\"github.com/replicate/cog/pkg/env\"\n)\n\nfunc TestPostNewVersion(t *testing.T) {\n\t// Setup mock http server\n\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\toutput := \"{\\\"version\\\":\\\"user/test:53c740f17ce88a61c3da5b0c20e48fd48e2da537c3a1276dec63ab11fbad6bcb\\\"}\"\n\t\tw.WriteHeader(http.StatusCreated)\n\t\tw.Write([]byte(output))\n\t}))\n\tdefer server.Close()\n\turl, err := url.Parse(server.URL)\n\trequire.NoError(t, err)\n\tt.Setenv(env.SchemeEnvVarName, url.Scheme)\n\tt.Setenv(env.WebHostEnvVarName, url.Host)\n\n\tdir := t.TempDir()\n\n\t// Create mock predict\n\tpredictPyPath := filepath.Join(dir, \"predict.py\")\n\thandle, err := os.Create(predictPyPath)\n\trequire.NoError(t, err)\n\thandle.WriteString(\"import cog\")\n\tdockertest.MockCogConfig = \"{\\\"build\\\":{\\\"python_version\\\":\\\"3.12\\\",\\\"python_packages\\\":[\\\"torch==2.5.0\\\",\\\"beautifulsoup4==4.12.3\\\"],\\\"system_packages\\\":[\\\"git\\\"]},\\\"image\\\":\\\"test\\\",\\\"predict\\\":\\\"\" + predictPyPath + \":Predictor\\\"}\"\n\n\t// Setup mock command\n\tcommand := dockertest.NewMockCommand()\n\n\tclient := NewClient(command, http.DefaultClient)\n\terr = client.PostNewVersion(t.Context(), \"r8.im/user/test\", []File{}, []File{}, nil)\n\trequire.NoError(t, err)\n}\n\nfunc TestVersionFromManifest(t *testing.T) {\n\t// Setup mock command\n\tcommand := dockertest.NewMockCommand()\n\n\t// Create mock predict\n\tdir := t.TempDir()\n\tpredictPyPath := filepath.Join(dir, \"predict.py\")\n\thandle, err := os.Create(predictPyPath)\n\trequire.NoError(t, err)\n\thandle.WriteString(\"import cog\")\n\tdockertest.MockCogConfig = \"{\\\"build\\\":{\\\"python_version\\\":\\\"3.12\\\",\\\"python_packages\\\":[\\\"torch==2.5.0\\\",\\\"beautifulsoup4==4.12.3\\\"],\\\"system_packages\\\":[\\\"git\\\"]},\\\"image\\\":\\\"test\\\",\\\"predict\\\":\\\"\" + predictPyPath + \":Predictor\\\"}\"\n\tdockertest.MockOpenAPISchema = \"{\\\"test\\\": true}\"\n\n\tclient := NewClient(command, http.DefaultClient)\n\tversion, err := client.versionFromManifest(t.Context(), \"r8.im/user/test\", []File{}, []File{}, nil)\n\trequire.NoError(t, err)\n\n\tvar openAPISchema map[string]any\n\terr = json.Unmarshal([]byte(dockertest.MockOpenAPISchema), &openAPISchema)\n\trequire.NoError(t, err)\n\n\tvar cogConfig config.Config\n\terr = json.Unmarshal([]byte(dockertest.MockCogConfig), &cogConfig)\n\trequire.NoError(t, err)\n\n\trequire.Equal(t, openAPISchema, version.OpenAPISchema)\n\trequire.Equal(t, cogConfig, version.CogConfig)\n}\n\nfunc TestVersionURLErrorWithoutR8IMPrefix(t *testing.T) {\n\t_, err := newVersionURL(\"docker.com/thing/thing\")\n\trequire.Error(t, err)\n}\n\nfunc TestVersionURLErrorWithout3Components(t *testing.T) {\n\t_, err := newVersionURL(\"username/test\")\n\trequire.Error(t, err)\n}\n\nfunc TestDoFileChallenge(t *testing.T) {\n\tdir := t.TempDir()\n\tpath := filepath.Join(dir, \"test.tmp\")\n\td1 := []byte(\"hello\\nreplicate\\nhello\\n\")\n\terr := os.WriteFile(path, d1, 0o644)\n\trequire.NoError(t, err)\n\n\tpath2 := filepath.Join(dir, \"test2.tmp\")\n\td2 := []byte(\"hello\\nreplicate\\nhello\\n\")\n\terr = os.WriteFile(path2, d2, 0o644)\n\trequire.NoError(t, err)\n\n\tfiles := []File{\n\t\t{\n\t\t\tPath:   path,\n\t\t\tDigest: \"abc\",\n\t\t\tSize:   22,\n\t\t},\n\t}\n\tweights := []File{\n\t\t{\n\t\t\tPath:   path,\n\t\t\tDigest: \"def\",\n\t\t\tSize:   22,\n\t\t},\n\t}\n\n\tabcChallenge := FileChallenge{\n\t\tID:     \"abc\",\n\t\tDigest: \"abc\",\n\t\tStart:  0,\n\t\tEnd:    6,\n\t\tSalt:   \"go\\n\",\n\t}\n\n\tdefChallenge := FileChallenge{\n\t\tID:     \"def\",\n\t\tDigest: \"def\",\n\t\tStart:  16,\n\t\tEnd:    22,\n\t\tSalt:   \"go\\n\",\n\t}\n\n\t// Setup mock http server\n\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\tw.WriteHeader(http.StatusOK)\n\t\tvar challengeRequest FileChallengeRequest\n\t\t// Ignore errors - make sure the test is set up correctly\n\t\tjson.NewDecoder(r.Body).Decode(&challengeRequest)\n\t\tif challengeRequest.Digest == \"abc\" {\n\t\t\tbody, _ := json.Marshal(abcChallenge)\n\t\t\tw.Write(body)\n\t\t} else {\n\t\t\tbody, _ := json.Marshal(defChallenge)\n\t\t\tw.Write(body)\n\t\t}\n\t}))\n\tdefer server.Close()\n\turl, err := url.Parse(server.URL)\n\trequire.NoError(t, err)\n\tt.Setenv(env.SchemeEnvVarName, url.Scheme)\n\tt.Setenv(env.WebHostEnvVarName, url.Host)\n\n\t// Setup mock command\n\tcommand := dockertest.NewMockCommand()\n\tclient := NewClient(command, http.DefaultClient)\n\tresponse, err := client.InitiateAndDoFileChallenge(t.Context(), weights, files)\n\trequire.NoError(t, err)\n\tassert.ElementsMatch(t, response, []FileChallengeAnswer{\n\t\t{\n\t\t\tChallengeID: \"abc\",\n\t\t\tDigest:      \"abc\",\n\t\t\tHash:        \"43d250d92b5dbb47f75208de8e9a9a321d23e85eed0dc3d5dfa83bc3cc5aa68c\",\n\t\t},\n\t\t{\n\t\t\tChallengeID: \"def\",\n\t\t\tDigest:      \"def\",\n\t\t\tHash:        \"43d250d92b5dbb47f75208de8e9a9a321d23e85eed0dc3d5dfa83bc3cc5aa68c\",\n\t\t},\n\t})\n}\n\nfunc TestFetchToken(t *testing.T) {\n\t// Setup mock http server\n\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\tswitch r.URL.Path {\n\t\tcase \"/api/token/user\":\n\t\t\t// Mock token exchange response\n\t\t\t//nolint:gosec\n\t\t\ttokenResponse := `{\n\t\t\t\t\"keys\": {\n\t\t\t\t\t\"cog\": {\n\t\t\t\t\t\t\"key\": \"test-api-token\",\n\t\t\t\t\t\t\"expires_at\": \"2024-12-31T23:59:59Z\"\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}`\n\t\t\tw.WriteHeader(http.StatusOK)\n\t\t\tw.Write([]byte(tokenResponse))\n\t\tdefault:\n\t\t\tw.WriteHeader(http.StatusNotFound)\n\t\t}\n\t}))\n\tdefer server.Close()\n\turl, err := url.Parse(server.URL)\n\trequire.NoError(t, err)\n\tt.Setenv(env.SchemeEnvVarName, url.Scheme)\n\tt.Setenv(env.WebHostEnvVarName, url.Host)\n\n\t// Setup mock command\n\tcommand := dockertest.NewMockCommand()\n\n\tclient := NewClient(command, http.DefaultClient)\n\ttoken, err := client.FetchAPIToken(t.Context(), \"user\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"test-api-token\", token)\n}\n"
  },
  {
    "path": "pkg/weights/manifest.go",
    "content": "package weights\n\nimport (\n\t\"encoding/binary\"\n\t\"encoding/hex\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"hash/crc32\"\n\t\"io\"\n\t\"os\"\n\t\"path\"\n)\n\n// Manifest contains metadata about weights files in a model\ntype Manifest struct {\n\tFiles map[string]Metadata `json:\"files\"`\n}\n\n// Metadata contains information about a file\ntype Metadata struct {\n\t// CRC32 is the CRC32 checksum of the file encoded as a hexadecimal string\n\tCRC32 string `json:\"crc32\"`\n}\n\n// NewManifest creates a new manifest\nfunc NewManifest() *Manifest {\n\treturn &Manifest{}\n}\n\n// LoadManifest loads a manifest from a file\nfunc LoadManifest(filename string) (*Manifest, error) {\n\tif _, err := os.Stat(filename); err != nil {\n\t\treturn nil, err\n\t}\n\tfile, err := os.Open(filename)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tdefer file.Close()\n\n\tm := &Manifest{}\n\tdecoder := json.NewDecoder(file)\n\tif err := decoder.Decode(m); err != nil {\n\t\treturn nil, err\n\t}\n\treturn m, nil\n}\n\n// Save saves a manifest to a file\nfunc (m *Manifest) Save(filename string) error {\n\tif err := os.MkdirAll(path.Dir(filename), 0o755); err != nil {\n\t\treturn err\n\t}\n\n\tfile, err := os.Create(filename)\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer file.Close()\n\tencoder := json.NewEncoder(file)\n\treturn encoder.Encode(m)\n}\n\n// Equal compares the files in two manifests for strict equality\nfunc (m *Manifest) Equal(other *Manifest) bool {\n\tif len(m.Files) != len(other.Files) {\n\t\treturn false\n\t}\n\n\tfor path, crc32 := range m.Files {\n\t\tif otherCrc32, ok := other.Files[path]; !ok || otherCrc32 != crc32 {\n\t\t\treturn false\n\t\t}\n\t}\n\n\treturn true\n}\n\n// AddFile adds a file to the manifest, calculating its CRC32 checksum\nfunc (m *Manifest) AddFile(path string) error {\n\tcrc32Algo := crc32.NewIEEE()\n\t// generate checksum of file\n\tfile, err := os.Open(path)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to open file %s: %w\", path, err)\n\t}\n\tdefer file.Close()\n\t_, err = io.Copy(crc32Algo, file)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to generate checksum of file %s: %w\", path, err)\n\t}\n\tchecksum := crc32Algo.Sum32()\n\n\t// encode checksum as hexadecimal string\n\tbytes := make([]byte, 4)\n\tbinary.LittleEndian.PutUint32(bytes, checksum)\n\tencoded := hex.EncodeToString(bytes)\n\n\tif m.Files == nil {\n\t\tm.Files = make(map[string]Metadata)\n\t}\n\tm.Files[path] = Metadata{\n\t\tCRC32: encoded,\n\t}\n\n\treturn nil\n}\n"
  },
  {
    "path": "pkg/weights/weights.go",
    "content": "package weights\n\nimport (\n\t\"os\"\n\t\"path/filepath\"\n\t\"slices\"\n\t\"sort\"\n\t\"strings\"\n)\n\nvar prefixesToIgnore = []string{\".cog\", \".git\", \"__pycache__\"}\n\nvar suffixesToIgnore = []string{\n\t\".py\", \".ipynb\", \".whl\", // Python projects\n\t\".jpg\", \".jpeg\", \".png\", \".webp\", \".svg\", \".gif\", \".avif\", \".heic\", // images\n\t\".mp4\", \".mov\", \".avi\", \".wmv\", \".mkv\", \".webm\", // videos\n\t\".mp3\", \".wav\", \".ogg\", \".flac\", \".aac\", \".m4a\", // audio files\n\t\".log\", // logs\n}\n\n// FileWalker is a function type that walks the file tree rooted at root, calling walkFn for each file or directory in the tree, including root.\ntype FileWalker func(root string, walkFn filepath.WalkFunc) error\n\nfunc FindWeights(fw FileWalker) ([]string, []string, error) {\n\tvar files []string\n\tvar codeFiles []string\n\terr := fw(\".\", func(path string, info os.FileInfo, err error) error {\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\t\tif info.IsDir() {\n\t\t\treturn nil\n\t\t}\n\t\tif isGitFile(path) {\n\t\t\treturn nil\n\t\t}\n\t\tif isCodeFile(path) {\n\t\t\tcodeFiles = append(codeFiles, path)\n\t\t\treturn nil\n\t\t}\n\n\t\tif info.Size() < sizeThreshold {\n\t\t\treturn nil\n\t\t}\n\t\tif isNonModelFiles(path) {\n\t\t\treturn nil\n\t\t}\n\n\t\tfiles = append(files, path)\n\t\treturn nil\n\t})\n\tif err != nil {\n\t\treturn nil, nil, err\n\t}\n\n\t// by sorting the files by levels, we can filter out directories that are prefixes of other directories\n\t// e.g. /a/b/ is a prefix of /a/b/c/, so we can filter out /a/b/c/\n\tsortFilesByLevels(files)\n\n\tdirs, rootFiles := getDirsAndRootfiles(files)\n\tdirs = filterDirsContainingCode(dirs, codeFiles)\n\n\treturn dirs, rootFiles, nil\n}\n\nfunc isNonModelFiles(path string) bool {\n\tfor _, prefix := range prefixesToIgnore {\n\t\tif strings.HasPrefix(path, prefix) {\n\t\t\treturn true\n\t\t}\n\t}\n\tfor _, suffix := range suffixesToIgnore {\n\t\tif strings.HasSuffix(path, suffix) {\n\t\t\treturn true\n\t\t}\n\t}\n\treturn false\n}\n\nconst sizeThreshold = 10 * 1024 * 1024 // 10MB\n\nfunc sortFilesByLevels(files []string) {\n\tsort.Slice(files, func(i, j int) bool {\n\t\tlist1 := strings.Split(files[i], \"/\")\n\t\tlist2 := strings.Split(files[j], \"/\")\n\t\tif len(list1) != len(list2) {\n\t\t\treturn len(list1) < len(list2)\n\t\t}\n\t\tfor k := range list1 {\n\t\t\tif list1[k] != list2[k] {\n\t\t\t\treturn list1[k] < list2[k]\n\t\t\t}\n\t\t}\n\t\treturn false\n\t})\n}\n\n// isCodeFile detects if a given path is a code file based on whether the file path ends with \".py\" or \".ipynb\"\nfunc isCodeFile(path string) bool {\n\text := filepath.Ext(path)\n\treturn ext == \".py\" || ext == \".ipynb\"\n}\n\nfunc isGitFile(path string) bool {\n\tdir, _ := filepath.Split(path)\n\tfolders := strings.Split(filepath.Clean(dir), string(filepath.Separator))\n\treturn slices.Contains(folders, \".git\")\n}\n\n// filterDirsContainingCode filters out directories that contain code files.\n// If a dir is a prefix for any given codeFiles, it will be filtered out.\nfunc filterDirsContainingCode(dirs []string, codeFiles []string) []string {\n\tfilteredDirs := make([]string, 0, len(dirs))\n\n\t// Filter out directories that are prefixes of code directories\n\tfor _, dir := range dirs {\n\t\tisPrefix := false\n\t\tfor _, codeFile := range codeFiles {\n\t\t\tif strings.HasPrefix(codeFile, dir) {\n\t\t\t\tisPrefix = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tif !isPrefix {\n\t\t\tfilteredDirs = append(filteredDirs, dir)\n\t\t}\n\t}\n\n\treturn filteredDirs\n}\n\nfunc getDirsAndRootfiles(files []string) ([]string, []string) {\n\t// get all the directories that contain model weights files\n\t// remove sub-directories if their parent directory is already in the list\n\tvar dirs []string\n\n\t// for large model files in root directory, we should not add the \".\" to dirs\n\tvar rootFiles []string\n\tfor _, f := range files {\n\t\tdir := filepath.Dir(f)\n\t\tif dir == \".\" || dir == \"/\" {\n\t\t\trootFiles = append(rootFiles, f)\n\t\t\tcontinue\n\t\t}\n\n\t\tif hasParent(dir, dirs) {\n\t\t\tcontinue\n\t\t}\n\t\tdirs = append(dirs, dir)\n\t}\n\treturn dirs, rootFiles\n}\n\nfunc hasParent(dir string, dirs []string) bool {\n\tfor _, d := range dirs {\n\t\tparent := d + string(filepath.Separator)\n\t\tchild := dir + string(filepath.Separator)\n\t\tif strings.HasPrefix(child, parent) {\n\t\t\treturn true\n\t\t}\n\n\t}\n\treturn false\n}\n"
  },
  {
    "path": "pkg/weights/weights_test.go",
    "content": "package weights\n\nimport (\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\t\"time\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\n// mockFileInfo is a test type to mock os.FileInfo\ntype mockFileInfo struct {\n\tsize int64\n}\n\nfunc (mfi mockFileInfo) Size() int64 {\n\treturn mfi.size\n}\nfunc (mfi mockFileInfo) Name() string {\n\treturn \"\"\n}\nfunc (mfi mockFileInfo) Mode() os.FileMode {\n\treturn 0\n}\nfunc (mfi mockFileInfo) ModTime() time.Time {\n\treturn time.Time{}\n}\nfunc (mfi mockFileInfo) IsDir() bool {\n\treturn false\n}\nfunc (mfi mockFileInfo) Sys() any {\n\treturn nil\n}\n\n// Test case for root directory with large and small model files\nfunc TestRootDirModelFiles(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, sizeThreshold, sizeThreshold - 1}\n\t\tfor i, path := range []string{\"large-a\", \"large-b\", \"small\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"large-a\", \"large-b\"}, rootFiles)\n\trequire.Empty(t, dirs)\n}\n\n// Test case for sub directory with large and small model files\nfunc TestSubDirModelFiles(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, sizeThreshold, sizeThreshold - 1}\n\t\tfor i, path := range []string{\"models/large-a\", \"models/large-b\", \"models/small\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Empty(t, rootFiles)\n\trequire.Equal(t, []string{\"models\"}, dirs)\n}\n\n// Test case for both root and sub directory with large model files\nfunc TestRootAndSubDirModelFiles(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, sizeThreshold}\n\t\tfor i, path := range []string{\"root-large\", \"models/large-a\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"root-large\"}, rootFiles)\n\trequire.Equal(t, []string{\"models\"}, dirs)\n}\n\n// Test case for root directory with both large model and code files\nfunc TestRootDirLargeModelAndCodeFiles(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, 1024}\n\t\tfor i, path := range []string{\"root-large\", \"predict.py\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"root-large\"}, rootFiles)\n\trequire.Empty(t, dirs)\n}\n\n// Test case for sub directory with both large model and code files\nfunc TestSubDirLargeModelAndCodeFiles(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, 1024}\n\t\tfor i, path := range []string{\"models/root-large\", \"models/predict.py\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Empty(t, rootFiles)\n\trequire.Empty(t, dirs)\n}\n\n// Test case for sub-directory with code files under large model directory\nfunc TestSubDirLargeModelDirWithCodeFiles(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, 1024}\n\t\tfor i, path := range []string{\"models/root-large\", \"models/code/predict.py\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Empty(t, rootFiles)\n\trequire.Empty(t, dirs)\n}\n\n// Test case for sorting for model directories\nfunc TestDirSorting(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, sizeThreshold, sizeThreshold}\n\t\tfor i, path := range []string{\"models2/b/large\", \"models2/a/large\", \"models/large\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Empty(t, rootFiles)\n\trequire.Equal(t, []string{\"models\", \"models2/a\", \"models2/b\"}, dirs)\n}\n\n// Test case for merging sub-directories with large models\nfunc TestSubDirMerge(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, sizeThreshold, sizeThreshold}\n\t\tfor i, path := range []string{\"models/b/large\", \"models/a/large\", \"models/large\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Empty(t, rootFiles)\n\trequire.Equal(t, []string{\"models\"}, dirs)\n}\n\n// Test case for ignoring files within a .git directory\nfunc TestIgnoreGitFiles(t *testing.T) {\n\tmockFileWalker := func(root string, walkFn filepath.WalkFunc) error {\n\t\tsizes := []int64{sizeThreshold, sizeThreshold, 1024}\n\t\tfor i, path := range []string{\".git/root-large\", \"root-large\", \"predict.py\"} {\n\t\t\twalkFn(path, mockFileInfo{size: sizes[i]}, nil)\n\t\t}\n\t\treturn nil\n\t}\n\n\tdirs, rootFiles, err := FindWeights(mockFileWalker)\n\trequire.NoError(t, err)\n\trequire.Equal(t, []string{\"root-large\"}, rootFiles)\n\trequire.Empty(t, dirs)\n}\n"
  },
  {
    "path": "pkg/wheels/wheels.go",
    "content": "// Package wheels provides configuration for sourcing cog and coglet wheels.\npackage wheels\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"regexp\"\n\t\"sort\"\n\t\"strings\"\n\n\t\"github.com/replicate/cog/pkg/global\"\n\tcogversion \"github.com/replicate/cog/pkg/util/version\"\n)\n\nvar semverPreReleaseRe = regexp.MustCompile(`-alpha(\\d+)|-beta(\\d+)|-rc(\\d+)|-dev(\\d*)`)\n\n// pep440PreReleaseRe matches PEP 440 pre-release identifiers (a1, b2, rc1, .dev1)\nvar pep440PreReleaseRe = regexp.MustCompile(`\\d(a|b|rc|\\.dev)\\d`)\n\n// IsPreRelease returns true if the version string contains a pre-release identifier\n// in either semver (-alpha1, -beta2, -rc1, -dev1) or PEP 440 (a1, b2, rc1, .dev1) format.\nfunc IsPreRelease(version string) bool {\n\treturn semverPreReleaseRe.MatchString(version) || pep440PreReleaseRe.MatchString(version)\n}\n\n// MinimumSDKVersion is the minimum cog SDK version that can be explicitly requested.\n// Versions older than this lack features required by the current CLI.\nconst MinimumSDKVersion = \"0.16.0\"\n\n// BaseVersionRe extracts the MAJOR.MINOR.PATCH prefix, ignoring pre-release suffixes.\nvar BaseVersionRe = regexp.MustCompile(`^(\\d+\\.\\d+\\.\\d+)`)\n\n// ValidateSDKVersion checks that a PyPI WheelConfig does not request a version\n// older than MinimumSDKVersion. Non-PyPI sources, unpinned versions, and nil\n// configs are always valid.\nfunc ValidateSDKVersion(config *WheelConfig, label string) error {\n\tif config == nil || config.Source != WheelSourcePyPI || config.Version == \"\" {\n\t\treturn nil\n\t}\n\tbase := config.Version\n\tif m := BaseVersionRe.FindString(base); m != \"\" {\n\t\tbase = m\n\t}\n\treqVer, err := cogversion.NewVersion(base)\n\tif err != nil {\n\t\treturn nil // unparseable — let pip catch real problems\n\t}\n\tminVer := cogversion.MustVersion(MinimumSDKVersion)\n\tif reqVer.GreaterOrEqual(minVer) {\n\t\treturn nil\n\t}\n\treturn fmt.Errorf(\"%s version %s is below the minimum required version %s\", label, config.Version, MinimumSDKVersion)\n}\n\n// WheelSource represents the source type for the wheel to install\ntype WheelSource int\n\nconst (\n\t// WheelSourcePyPI installs from PyPI (default for released builds)\n\tWheelSourcePyPI WheelSource = iota\n\t// WheelSourceURL uses a custom URL\n\tWheelSourceURL\n\t// WheelSourceFile uses a local file path\n\tWheelSourceFile\n)\n\n// String returns the string representation of the WheelSource\nfunc (s WheelSource) String() string {\n\tswitch s {\n\tcase WheelSourcePyPI:\n\t\treturn \"pypi\"\n\tcase WheelSourceURL:\n\t\treturn \"url\"\n\tcase WheelSourceFile:\n\t\treturn \"file\"\n\tdefault:\n\t\treturn \"unknown\"\n\t}\n}\n\n// WheelConfig represents the configuration for which wheel to install\ntype WheelConfig struct {\n\t// Source indicates where the wheel comes from\n\tSource WheelSource\n\t// URL is set when Source is WheelSourceURL\n\tURL string\n\t// Path is set when Source is WheelSourceFile (absolute path)\n\tPath string\n\t// Version is set when Source is WheelSourcePyPI (optional, empty = latest)\n\tVersion string\n}\n\n// CogSDKWheelEnvVar is the environment variable name for cog SDK wheel selection\nconst CogSDKWheelEnvVar = \"COG_SDK_WHEEL\"\n\n// CogletWheelEnvVar is the environment variable name for coglet wheel selection\nconst CogletWheelEnvVar = \"COGLET_WHEEL\"\n\n// ParseWheelValue parses a wheel env var value and returns the appropriate WheelConfig.\n// Supported values:\n//   - \"pypi\" - Install from PyPI (latest version)\n//   - \"pypi:0.12.0\" - Install specific version from PyPI\n//   - \"https://...\" or \"http://...\" - Direct wheel URL\n//   - \"/path/to/file.whl\" or \"relative/path\" - Local file or directory (resolved to abspath)\n//\n// Paths that point to directories are resolved later by the Resolve functions,\n// which glob for the appropriate wheel inside the directory.\n//\n// Returns nil if the value is empty (caller should use auto-detection).\nfunc ParseWheelValue(value string) *WheelConfig {\n\tvalue = strings.TrimSpace(value)\n\tif value == \"\" {\n\t\treturn nil\n\t}\n\n\t// \"pypi\" or \"pypi:version\" requests PyPI\n\tif strings.EqualFold(value, \"pypi\") {\n\t\treturn &WheelConfig{Source: WheelSourcePyPI}\n\t}\n\tif strings.HasPrefix(strings.ToLower(value), \"pypi:\") {\n\t\t// Extract version after \"pypi:\" prefix, preserving original case\n\t\treturn &WheelConfig{Source: WheelSourcePyPI, Version: value[5:]}\n\t}\n\n\t// Check for URL (http:// or https://)\n\tif strings.HasPrefix(value, \"https://\") || strings.HasPrefix(value, \"http://\") {\n\t\treturn &WheelConfig{Source: WheelSourceURL, URL: value}\n\t}\n\n\t// Treat everything else as a file/directory path - resolve to absolute\n\tabsPath, err := filepath.Abs(value)\n\tif err != nil {\n\t\tabsPath = value\n\t}\n\treturn &WheelConfig{Source: WheelSourceFile, Path: absPath}\n}\n\nvar executablePath = os.Executable\nvar evalSymlinks = filepath.EvalSymlinks\n\n// goarchToWheelPlatform maps GOARCH values to wheel filename platform substrings.\nfunc goarchToWheelPlatform(goarch string) string {\n\tswitch goarch {\n\tcase \"amd64\":\n\t\treturn \"x86_64\"\n\tcase \"arm64\":\n\t\treturn \"aarch64\"\n\tdefault:\n\t\treturn \"\"\n\t}\n}\n\nfunc bestWheelMatch(matches []string, platform string) string {\n\tif len(matches) == 0 {\n\t\treturn \"\"\n\t}\n\tif platform != \"\" {\n\t\tplatStr := goarchToWheelPlatform(platform)\n\t\tif platStr != \"\" {\n\t\t\tvar filtered []string\n\t\t\tfor _, match := range matches {\n\t\t\t\tbase := filepath.Base(match)\n\t\t\t\tif strings.Contains(base, platStr) || strings.Contains(base, \"-none-any\") {\n\t\t\t\t\tfiltered = append(filtered, match)\n\t\t\t\t}\n\t\t\t}\n\t\t\tmatches = filtered\n\t\t}\n\t}\n\tif len(matches) == 0 {\n\t\treturn \"\"\n\t}\n\tsort.Strings(matches)\n\treturn matches[len(matches)-1]\n}\n\n// distFromExecutable returns the dist/ directory relative to the running cog\n// binary, if it appears to be in a goreleaser output layout (dist/go//cog).\n// Returns empty string if the path cannot be determined.\nfunc distFromExecutable() string {\n\texePath, err := executablePath()\n\tif err != nil {\n\t\treturn \"\"\n\t}\n\texePath, err = evalSymlinks(exePath)\n\tif err != nil {\n\t\treturn \"\"\n\t}\n\n\tdistDir := filepath.Clean(filepath.Join(filepath.Dir(exePath), \"..\", \"..\"))\n\tif info, err := os.Stat(distDir); err == nil && info.IsDir() {\n\t\treturn distDir\n\t}\n\treturn \"\"\n}\n\n// findWheelInAutoDetectDist checks ./dist and dist relative to the cog executable.\n// Returns the absolute path if found, empty string otherwise.\nfunc findWheelInAutoDetectDist(pattern string, platform string) string {\n\tmatches, _ := filepath.Glob(filepath.Join(\"dist\", pattern))\n\tif best := bestWheelMatch(matches, platform); best != \"\" {\n\t\tabsPath, _ := filepath.Abs(best)\n\t\tif absPath != \"\" {\n\t\t\treturn absPath\n\t\t}\n\t\treturn best\n\t}\n\n\tif distDir := distFromExecutable(); distDir != \"\" {\n\t\tmatches, _ = filepath.Glob(filepath.Join(distDir, pattern))\n\t\tif best := bestWheelMatch(matches, platform); best != \"\" {\n\t\t\treturn best\n\t\t}\n\t}\n\n\treturn \"\"\n}\n\n// DetectLocalSDKVersion checks dist/ (CWD and executable-relative) for a cog\n// SDK wheel and extracts the version from its filename. Returns empty string if\n// no local wheel is found.\nfunc DetectLocalSDKVersion() string {\n\tpath := findWheelInAutoDetectDist(\"cog-*.whl\", \"\")\n\tif path == \"\" {\n\t\treturn \"\"\n\t}\n\t// Wheel filename format: cog----.whl\n\tbase := filepath.Base(path)\n\tif !strings.HasPrefix(base, \"cog-\") {\n\t\treturn \"\"\n\t}\n\trest := strings.TrimPrefix(base, \"cog-\")\n\tif idx := strings.Index(rest, \"-\"); idx > 0 {\n\t\treturn rest[:idx]\n\t}\n\treturn \"\"\n}\n\n// resolveWheelPath resolves a wheel path that may be a file or directory.\n// If path is a directory, globs for pattern inside it, filtering by platform if non-empty.\n// If path is a file, returns it directly.\nfunc resolveWheelPath(path string, pattern string, platform string, envVar string) (string, error) {\n\tinfo, err := os.Stat(path) //nolint:gosec // G703: path from build config, not user input\n\tif err != nil {\n\t\treturn \"\", fmt.Errorf(\"%s: path not found: %s\", envVar, path)\n\t}\n\n\tif !info.IsDir() {\n\t\treturn path, nil\n\t}\n\n\tmatches, _ := filepath.Glob(filepath.Join(path, pattern))\n\tif len(matches) == 0 {\n\t\treturn \"\", fmt.Errorf(\"%s: no wheel matching '%s' found in %s\\n\\nTo build the wheel, run: mise run build:sdk (for cog) or mise run build:coglet (for coglet)\", envVar, pattern, path)\n\t}\n\n\t// Filter by platform if specified\n\tplatStr := goarchToWheelPlatform(platform)\n\tif platStr != \"\" {\n\t\tvar filtered []string\n\t\tfor _, m := range matches {\n\t\t\tbase := filepath.Base(m)\n\t\t\tif strings.Contains(base, platStr) || strings.Contains(base, \"-none-any\") {\n\t\t\t\tfiltered = append(filtered, m)\n\t\t\t}\n\t\t}\n\t\tif len(filtered) == 0 {\n\t\t\treturn \"\", fmt.Errorf(\"%s: no wheel for platform %s found in %s (found %d for other platforms)\", envVar, platform, path, len(matches))\n\t\t}\n\t\tmatches = filtered\n\t}\n\n\tif len(matches) > 1 {\n\t\treturn \"\", fmt.Errorf(\"%s: multiple wheels matching '%s' in %s — specify the exact file path\", envVar, pattern, path)\n\t}\n\n\treturn matches[0], nil\n}\n\n// ResolveCogWheel resolves the WheelConfig for the cog SDK.\n//\n// Parameters:\n//   - envValue: value of COG_SDK_WHEEL env var (empty string if not set)\n//   - version: the CLI version (e.g. \"dev\", \"0.17.0\", \"0.17.0-alpha1\")\n//\n// Resolution order:\n//  1. envValue (if non-empty, explicit override)\n//  2. Auto-detect: check dist/cog-*.whl (for development builds only)\n//  3. Default: PyPI latest (use build.sdk_version in cog.yaml to pin)\nfunc ResolveCogWheel(envValue string, version string) (*WheelConfig, error) {\n\t// Check explicit env var first\n\tif config := ParseWheelValue(envValue); config != nil {\n\t\tif config.Source == WheelSourceFile {\n\t\t\t// cog SDK is pure Python (py3-none-any), no platform filtering needed\n\t\t\tresolved, err := resolveWheelPath(config.Path, \"cog-*.whl\", \"\", CogSDKWheelEnvVar)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tconfig.Path = resolved\n\t\t}\n\t\treturn config, nil\n\t}\n\n\tisDev := version == \"dev\" || strings.Contains(version, \"-dev\") || strings.Contains(version, \"+\")\n\n\t// Auto-detect for dev builds: check ./dist or executable-relative dist\n\tif isDev {\n\t\tif path := findWheelInAutoDetectDist(\"cog-*.whl\", \"\"); path != \"\" {\n\t\t\treturn &WheelConfig{Source: WheelSourceFile, Path: path}, nil\n\t\t}\n\t}\n\n\t// Default: PyPI (always latest; use sdk_version in cog.yaml to pin)\n\treturn &WheelConfig{Source: WheelSourcePyPI}, nil\n}\n\n// GetCogWheelConfig is a convenience wrapper that reads COG_SDK_WHEEL from the environment\n// and version from global.Version.\nfunc GetCogWheelConfig() (*WheelConfig, error) {\n\treturn ResolveCogWheel(os.Getenv(CogSDKWheelEnvVar), global.Version)\n}\n\n// ResolveCogletWheel resolves the WheelConfig for coglet.\n//\n// targetArch is the GOARCH of the Docker build target (e.g. \"amd64\", \"arm64\").\n// It is used to select the correct platform-specific wheel from dist/.\n//\n// Resolution order:\n//  1. envValue (COGLET_WHEEL) if non-empty — explicit override\n//  2. Auto-detect: check ./dist for coglet-*.whl (development builds only)\n//  3. Default: PyPI latest (use COGLET_WHEEL=pypi:x.y.z to pin)\n//\n// Coglet is always required. Returns a valid config or an error.\n// The platform parameter is a GOARCH value (e.g. \"amd64\", \"arm64\") used to select\n// the correct platform-specific wheel from a directory. Pass \"\" to skip filtering.\nfunc ResolveCogletWheel(envValue string, version string, platform string) (*WheelConfig, error) {\n\t// Check explicit env var first\n\tif config := ParseWheelValue(envValue); config != nil {\n\t\tif config.Source == WheelSourceFile {\n\t\t\tresolved, err := resolveWheelPath(config.Path, \"coglet-*.whl\", platform, CogletWheelEnvVar)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tconfig.Path = resolved\n\t\t}\n\t\treturn config, nil\n\t}\n\n\tisDev := version == \"dev\" || strings.Contains(version, \"-dev\") || strings.Contains(version, \"+\")\n\n\t// Auto-detect for dev builds: check ./dist or executable-relative dist\n\tif isDev {\n\t\tif path := findWheelInAutoDetectDist(\"coglet-*.whl\", platform); path != \"\" {\n\t\t\treturn &WheelConfig{Source: WheelSourceFile, Path: path}, nil\n\t\t}\n\t}\n\n\t// Default: PyPI (always latest; use COGLET_WHEEL=pypi:x.y.z to pin)\n\treturn &WheelConfig{Source: WheelSourcePyPI}, nil\n}\n\n// GetCogletWheelConfig is a convenience wrapper that reads COGLET_WHEEL from the environment\n// and version from global.Version. targetArch is the GOARCH of the Docker build target\n// (e.g. \"amd64\", \"arm64\") used to select the correct platform-specific wheel.\nfunc GetCogletWheelConfig(targetArch string) (*WheelConfig, error) {\n\treturn ResolveCogletWheel(os.Getenv(CogletWheelEnvVar), global.Version, targetArch)\n}\n\n// SemverToPEP440 converts a semver pre-release version to PEP 440 format.\n// e.g. \"0.17.0-alpha1\" -> \"0.17.0a1\", \"0.17.0-beta2\" -> \"0.17.0b2\",\n// \"0.17.0-rc1\" -> \"0.17.0rc1\", \"0.17.0-dev1\" -> \"0.17.0.dev1\"\n// Stable versions pass through unchanged: \"0.17.0\" -> \"0.17.0\"\nfunc SemverToPEP440(version string) string {\n\treturn semverPreReleaseRe.ReplaceAllStringFunc(version, func(match string) string {\n\t\tmatch = strings.TrimPrefix(match, \"-\")\n\t\tmatch = strings.Replace(match, \"alpha\", \"a\", 1)\n\t\tmatch = strings.Replace(match, \"beta\", \"b\", 1)\n\t\t// rc stays as rc in PEP 440\n\t\t// dev -> .dev (PEP 440 uses dot separator)\n\t\tif strings.HasPrefix(match, \"dev\") {\n\t\t\treturn \".\" + match\n\t\t}\n\t\treturn match\n\t})\n}\n\n// PyPIPackageURL returns the pip install specifier for a PyPI package.\n// If version is empty, returns just the package name (latest).\n// Otherwise returns \"package==version\" with the version converted to PEP 440.\nfunc (c *WheelConfig) PyPIPackageURL(packageName string) string {\n\tif c.Version == \"\" {\n\t\treturn packageName\n\t}\n\treturn packageName + \"==\" + SemverToPEP440(c.Version)\n}\n"
  },
  {
    "path": "pkg/wheels/wheels_test.go",
    "content": "package wheels\n\nimport (\n\t\"os\"\n\t\"path/filepath\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n)\n\nfunc TestCogSDKWheelEnvVarName(t *testing.T) {\n\trequire.Equal(t, \"COG_SDK_WHEEL\", CogSDKWheelEnvVar)\n}\n\nfunc TestWheelSourceString(t *testing.T) {\n\ttests := []struct {\n\t\tsource   WheelSource\n\t\texpected string\n\t}{\n\t\t{WheelSourcePyPI, \"pypi\"},\n\t\t{WheelSourceURL, \"url\"},\n\t\t{WheelSourceFile, \"file\"},\n\t\t{WheelSource(99), \"unknown\"},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.expected, func(t *testing.T) {\n\t\t\trequire.Equal(t, tt.expected, tt.source.String())\n\t\t})\n\t}\n}\n\nfunc TestParseWheelValue(t *testing.T) {\n\ttests := []struct {\n\t\tname     string\n\t\tinput    string\n\t\texpected *WheelConfig\n\t}{\n\t\t// Empty/nil cases\n\t\t{\n\t\t\tname:     \"empty string returns nil\",\n\t\t\tinput:    \"\",\n\t\t\texpected: nil,\n\t\t},\n\t\t{\n\t\t\tname:     \"whitespace only returns nil\",\n\t\t\tinput:    \"   \",\n\t\t\texpected: nil,\n\t\t},\n\n\t\t// PyPI values\n\t\t{\n\t\t\tname:     \"pypi keyword\",\n\t\t\tinput:    \"pypi\",\n\t\t\texpected: &WheelConfig{Source: WheelSourcePyPI},\n\t\t},\n\t\t{\n\t\t\tname:     \"pypi uppercase\",\n\t\t\tinput:    \"PYPI\",\n\t\t\texpected: &WheelConfig{Source: WheelSourcePyPI},\n\t\t},\n\t\t{\n\t\t\tname:     \"pypi with version\",\n\t\t\tinput:    \"pypi:0.12.0\",\n\t\t\texpected: &WheelConfig{Source: WheelSourcePyPI, Version: \"0.12.0\"},\n\t\t},\n\t\t{\n\t\t\tname:     \"pypi with version uppercase\",\n\t\t\tinput:    \"PYPI:1.0.0\",\n\t\t\texpected: &WheelConfig{Source: WheelSourcePyPI, Version: \"1.0.0\"},\n\t\t},\n\n\t\t// relative directory paths (e.g. \"dist\") are resolved to absolute\n\t\t{\n\t\t\tname:  \"dist as relative path\",\n\t\t\tinput: \"dist\",\n\t\t\texpected: &WheelConfig{\n\t\t\t\tSource: WheelSourceFile,\n\t\t\t\t// Path will be converted to absolute\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname:  \"dist uppercase as relative path\",\n\t\t\tinput: \"DIST\",\n\t\t\texpected: &WheelConfig{\n\t\t\t\tSource: WheelSourceFile,\n\t\t\t\t// Path will be converted to absolute\n\t\t\t},\n\t\t},\n\n\t\t// URLs\n\t\t{\n\t\t\tname:  \"https URL\",\n\t\t\tinput: \"https://example.com/wheel.whl\",\n\t\t\texpected: &WheelConfig{\n\t\t\t\tSource: WheelSourceURL,\n\t\t\t\tURL:    \"https://example.com/wheel.whl\",\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname:  \"http URL\",\n\t\t\tinput: \"http://example.com/wheel.whl\",\n\t\t\texpected: &WheelConfig{\n\t\t\t\tSource: WheelSourceURL,\n\t\t\t\tURL:    \"http://example.com/wheel.whl\",\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname:  \"github release URL\",\n\t\t\tinput: \"https://github.com/replicate/cog/releases/download/v0.1.0/cog-0.1.0-py3-none-any.whl\",\n\t\t\texpected: &WheelConfig{\n\t\t\t\tSource: WheelSourceURL,\n\t\t\t\tURL:    \"https://github.com/replicate/cog/releases/download/v0.1.0/cog-0.1.0-py3-none-any.whl\",\n\t\t\t},\n\t\t},\n\n\t\t// File paths\n\t\t{\n\t\t\tname:  \"absolute path\",\n\t\t\tinput: \"/path/to/wheel.whl\",\n\t\t\texpected: &WheelConfig{\n\t\t\t\tSource: WheelSourceFile,\n\t\t\t\tPath:   \"/path/to/wheel.whl\",\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname:  \"relative path with ./\",\n\t\t\tinput: \"./dist/wheel.whl\",\n\t\t\texpected: &WheelConfig{\n\t\t\t\tSource: WheelSourceFile,\n\t\t\t\t// Path will be converted to absolute\n\t\t\t},\n\t\t},\n\t\t{\n\t\t\tname:  \"relative path without ./\",\n\t\t\tinput: \"path/to/wheel.whl\",\n\t\t\texpected: &WheelConfig{\n\t\t\t\tSource: WheelSourceFile,\n\t\t\t\t// Path will be converted to absolute\n\t\t\t},\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := ParseWheelValue(tt.input)\n\t\t\tif tt.expected == nil {\n\t\t\t\trequire.Nil(t, result)\n\t\t\t} else {\n\t\t\t\trequire.NotNil(t, result)\n\t\t\t\trequire.Equal(t, tt.expected.Source, result.Source)\n\t\t\t\trequire.Equal(t, tt.expected.URL, result.URL)\n\t\t\t\t// For relative paths, just verify they're converted to absolute\n\t\t\t\tif tt.expected.Path == \"\" && result.Source == WheelSourceFile {\n\t\t\t\t\trequire.True(t, filepath.IsAbs(result.Path), \"path should be absolute: %s\", result.Path)\n\t\t\t\t} else {\n\t\t\t\t\trequire.Equal(t, tt.expected.Path, result.Path)\n\t\t\t\t}\n\t\t\t\trequire.Equal(t, tt.expected.Version, result.Version)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestGetCogWheelConfig(t *testing.T) {\n\t// Create temp dir for file path tests and to avoid auto-detect from repo root\n\ttmpDir := t.TempDir()\n\twheelFile := filepath.Join(tmpDir, \"custom.whl\")\n\trequire.NoError(t, os.WriteFile(wheelFile, []byte(\"fake wheel\"), 0o600))\n\n\t// Change to temp dir and clear REPO_ROOT to prevent auto-detection from repo dist/\n\torigDir, err := os.Getwd()\n\trequire.NoError(t, err)\n\trequire.NoError(t, os.Chdir(tmpDir))\n\tdefer func() { require.NoError(t, os.Chdir(origDir)) }()\n\tt.Setenv(\"REPO_ROOT\", \"\")\n\n\ttests := []struct {\n\t\tname           string\n\t\tenvValue       string\n\t\tglobalVersion  string\n\t\texpectedSource WheelSource\n\t\texpectedPath   string\n\t\texpectedURL    string\n\t\texpectedVer    string\n\t}{\n\t\t// Release build defaults to PyPI latest (no version pin)\n\t\t{\n\t\t\tname:           \"release build defaults to PyPI latest\",\n\t\t\tenvValue:       \"\",\n\t\t\tglobalVersion:  \"0.12.0\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"\",\n\t\t},\n\t\t// Dev build with explicit pypi (auto-detection tested separately in TestGetCogWheelConfigAutoDetect)\n\t\t{\n\t\t\tname:           \"dev build defaults to PyPI without version\",\n\t\t\tenvValue:       \"pypi\",\n\t\t\tglobalVersion:  \"dev\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"\",\n\t\t},\n\t\t// Snapshot build with explicit pypi\n\t\t{\n\t\t\tname:           \"snapshot build defaults to PyPI without version\",\n\t\t\tenvValue:       \"pypi\",\n\t\t\tglobalVersion:  \"0.16.12-dev+g6793b492\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"\",\n\t\t},\n\t\t// Explicit pypi override\n\t\t{\n\t\t\tname:           \"explicit pypi\",\n\t\t\tenvValue:       \"pypi\",\n\t\t\tglobalVersion:  \"dev\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"\",\n\t\t},\n\t\t{\n\t\t\tname:           \"explicit pypi with version\",\n\t\t\tenvValue:       \"pypi:0.11.0\",\n\t\t\tglobalVersion:  \"0.12.0\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"0.11.0\",\n\t\t},\n\t\t// URL override\n\t\t{\n\t\t\tname:           \"URL override\",\n\t\t\tenvValue:       \"https://example.com/custom.whl\",\n\t\t\tglobalVersion:  \"0.12.0\",\n\t\t\texpectedSource: WheelSourceURL,\n\t\t\texpectedURL:    \"https://example.com/custom.whl\",\n\t\t},\n\t\t// File path override (use the real temp file)\n\t\t{\n\t\t\tname:           \"file path override\",\n\t\t\tenvValue:       wheelFile,\n\t\t\tglobalVersion:  \"0.12.0\",\n\t\t\texpectedSource: WheelSourceFile,\n\t\t\texpectedPath:   wheelFile,\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := ResolveCogWheel(tt.envValue, tt.globalVersion)\n\t\t\trequire.NoError(t, err)\n\t\t\trequire.NotNil(t, result)\n\t\t\trequire.Equal(t, tt.expectedSource, result.Source)\n\t\t\trequire.Equal(t, tt.expectedURL, result.URL)\n\t\t\trequire.Equal(t, tt.expectedPath, result.Path)\n\t\t\trequire.Equal(t, tt.expectedVer, result.Version)\n\t\t})\n\t}\n}\n\nfunc TestGetCogWheelConfigErrors(t *testing.T) {\n\t// Test error cases for wheel config\n\tt.Run(\"file not found\", func(t *testing.T) {\n\t\tt.Setenv(CogSDKWheelEnvVar, \"/nonexistent/path/wheel.whl\")\n\t\t_, err := GetCogWheelConfig()\n\t\trequire.Error(t, err)\n\t\trequire.Contains(t, err.Error(), \"path not found\")\n\t})\n}\n\nfunc TestGetCogWheelConfigAutoDetect(t *testing.T) {\n\t// Create a temp directory with a wheel file\n\ttmpDir := t.TempDir()\n\tdistDir := filepath.Join(tmpDir, \"dist\")\n\trequire.NoError(t, os.MkdirAll(distDir, 0o750))\n\n\twheelPath := filepath.Join(distDir, \"cog-0.1.0-py3-none-any.whl\")\n\trequire.NoError(t, os.WriteFile(wheelPath, []byte(\"fake wheel content\"), 0o600))\n\n\t// Change to temp dir\n\torigDir, err := os.Getwd()\n\trequire.NoError(t, err)\n\trequire.NoError(t, os.Chdir(tmpDir))\n\tdefer func() { require.NoError(t, os.Chdir(origDir)) }()\n\n\t// Test auto-detection in dev mode\n\tresult, err := ResolveCogWheel(\"\", \"dev\")\n\trequire.NoError(t, err)\n\trequire.NotNil(t, result)\n\trequire.Equal(t, WheelSourceFile, result.Source)\n\trequire.Contains(t, result.Path, \"cog-0.1.0-py3-none-any.whl\")\n\n\t// Test that release mode does NOT auto-detect (and has no version pin)\n\tresult, err = ResolveCogWheel(\"\", \"0.12.0\")\n\trequire.NoError(t, err)\n\trequire.NotNil(t, result)\n\trequire.Equal(t, WheelSourcePyPI, result.Source)\n\trequire.Equal(t, \"\", result.Version)\n}\n\nfunc TestResolveCogWheelUsesExecutableDist(t *testing.T) {\n\trootDir := t.TempDir()\n\tdistDir := filepath.Join(rootDir, \"dist\")\n\trequire.NoError(t, os.MkdirAll(distDir, 0o750))\n\n\twheelPath := filepath.Join(distDir, \"cog-0.1.0-py3-none-any.whl\")\n\trequire.NoError(t, os.WriteFile(wheelPath, []byte(\"fake wheel content\"), 0o600))\n\n\tfakeExe := filepath.Join(distDir, \"go\", \"linux_amd64\", \"cog\")\n\trequire.NoError(t, os.MkdirAll(filepath.Dir(fakeExe), 0o750))\n\trequire.NoError(t, os.WriteFile(fakeExe, []byte(\"binary\"), 0o600))\n\n\torigExecutablePath := executablePath\n\torigEvalSymlinks := evalSymlinks\n\tdefer func() {\n\t\texecutablePath = origExecutablePath\n\t\tevalSymlinks = origEvalSymlinks\n\t}()\n\n\texecutablePath = func() (string, error) {\n\t\treturn fakeExe, nil\n\t}\n\tevalSymlinks = func(path string) (string, error) {\n\t\treturn path, nil\n\t}\n\n\tcwd := t.TempDir()\n\torigDir, err := os.Getwd()\n\trequire.NoError(t, err)\n\trequire.NoError(t, os.Chdir(cwd))\n\tdefer func() { require.NoError(t, os.Chdir(origDir)) }()\n\n\tresult, err := ResolveCogWheel(\"\", \"dev\")\n\trequire.NoError(t, err)\n\trequire.NotNil(t, result)\n\trequire.Equal(t, WheelSourceFile, result.Source)\n\trequire.Contains(t, result.Path, \"cog-0.1.0-py3-none-any.whl\")\n}\n\nfunc TestGetCogletWheelConfig(t *testing.T) {\n\t// Change to temp dir to prevent auto-detection from repo dist/\n\ttmpDir := t.TempDir()\n\torigDir, err := os.Getwd()\n\trequire.NoError(t, err)\n\trequire.NoError(t, os.Chdir(tmpDir))\n\tdefer func() { require.NoError(t, os.Chdir(origDir)) }()\n\n\ttests := []struct {\n\t\tname           string\n\t\tenvValue       string\n\t\tglobalVersion  string\n\t\texpectedSource WheelSource\n\t\texpectedPath   string\n\t\texpectedURL    string\n\t\texpectedVer    string\n\t}{\n\t\t// Default: coglet from PyPI latest (release build, no version pin)\n\t\t{\n\t\t\tname:           \"release default uses PyPI latest\",\n\t\t\tenvValue:       \"\",\n\t\t\tglobalVersion:  \"0.12.0\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"\",\n\t\t},\n\t\t{\n\t\t\tname:           \"dev default falls back to PyPI without version\",\n\t\t\tenvValue:       \"\",\n\t\t\tglobalVersion:  \"dev\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"\",\n\t\t},\n\t\t// Explicit pypi\n\t\t{\n\t\t\tname:           \"explicit pypi\",\n\t\t\tenvValue:       \"pypi\",\n\t\t\tglobalVersion:  \"0.12.0\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"\",\n\t\t},\n\t\t{\n\t\t\tname:           \"explicit pypi with version\",\n\t\t\tenvValue:       \"pypi:0.11.0\",\n\t\t\tglobalVersion:  \"0.12.0\",\n\t\t\texpectedSource: WheelSourcePyPI,\n\t\t\texpectedVer:    \"0.11.0\",\n\t\t},\n\t\t// URL override\n\t\t{\n\t\t\tname:           \"URL override\",\n\t\t\tenvValue:       \"https://example.com/coglet.whl\",\n\t\t\tglobalVersion:  \"0.12.0\",\n\t\t\texpectedSource: WheelSourceURL,\n\t\t\texpectedURL:    \"https://example.com/coglet.whl\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult, err := ResolveCogletWheel(tt.envValue, tt.globalVersion, \"amd64\")\n\t\t\trequire.NoError(t, err)\n\t\t\trequire.NotNil(t, result)\n\t\t\trequire.Equal(t, tt.expectedSource, result.Source)\n\t\t\trequire.Equal(t, tt.expectedURL, result.URL)\n\t\t\trequire.Equal(t, tt.expectedPath, result.Path)\n\t\t\trequire.Equal(t, tt.expectedVer, result.Version)\n\t\t})\n\t}\n}\n\nfunc TestPyPIPackageURL(t *testing.T) {\n\ttests := []struct {\n\t\tname        string\n\t\tconfig      *WheelConfig\n\t\tpackageName string\n\t\texpected    string\n\t}{\n\t\t{\n\t\t\tname:        \"no version\",\n\t\t\tconfig:      &WheelConfig{Source: WheelSourcePyPI},\n\t\t\tpackageName: \"cog\",\n\t\t\texpected:    \"cog\",\n\t\t},\n\t\t{\n\t\t\tname:        \"with version\",\n\t\t\tconfig:      &WheelConfig{Source: WheelSourcePyPI, Version: \"0.12.0\"},\n\t\t\tpackageName: \"cog\",\n\t\t\texpected:    \"cog==0.12.0\",\n\t\t},\n\t\t{\n\t\t\tname:        \"coglet with version\",\n\t\t\tconfig:      &WheelConfig{Source: WheelSourcePyPI, Version: \"0.1.0\"},\n\t\t\tpackageName: \"coglet\",\n\t\t\texpected:    \"coglet==0.1.0\",\n\t\t},\n\t\t{\n\t\t\tname:        \"alpha pre-release converted to PEP 440\",\n\t\t\tconfig:      &WheelConfig{Source: WheelSourcePyPI, Version: \"0.17.0-alpha1\"},\n\t\t\tpackageName: \"cog\",\n\t\t\texpected:    \"cog==0.17.0a1\",\n\t\t},\n\t\t{\n\t\t\tname:        \"beta pre-release converted to PEP 440\",\n\t\t\tconfig:      &WheelConfig{Source: WheelSourcePyPI, Version: \"0.17.0-beta2\"},\n\t\t\tpackageName: \"cog\",\n\t\t\texpected:    \"cog==0.17.0b2\",\n\t\t},\n\t\t{\n\t\t\tname:        \"rc pre-release converted to PEP 440\",\n\t\t\tconfig:      &WheelConfig{Source: WheelSourcePyPI, Version: \"1.0.0-rc1\"},\n\t\t\tpackageName: \"cog\",\n\t\t\texpected:    \"cog==1.0.0rc1\",\n\t\t},\n\t\t{\n\t\t\tname:        \"dev pre-release converted to PEP 440\",\n\t\t\tconfig:      &WheelConfig{Source: WheelSourcePyPI, Version: \"0.17.0-dev1\"},\n\t\t\tpackageName: \"cog\",\n\t\t\texpected:    \"cog==0.17.0.dev1\",\n\t\t},\n\t}\n\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\tresult := tt.config.PyPIPackageURL(tt.packageName)\n\t\t\trequire.Equal(t, tt.expected, result)\n\t\t})\n\t}\n}\n\nfunc TestIsPreRelease(t *testing.T) {\n\ttests := []struct {\n\t\tversion  string\n\t\texpected bool\n\t}{\n\t\t// semver format\n\t\t{\"0.17.0-alpha1\", true},\n\t\t{\"0.17.0-beta2\", true},\n\t\t{\"0.17.0-rc1\", true},\n\t\t{\"0.17.0-dev1\", true},\n\t\t// PEP 440 format\n\t\t{\"0.17.0a1\", true},\n\t\t{\"0.17.0b2\", true},\n\t\t{\"0.17.0rc1\", true},\n\t\t{\"0.17.0.dev1\", true},\n\t\t// stable\n\t\t{\"0.17.0\", false},\n\t\t{\"1.0.0\", false},\n\t}\n\tfor _, tt := range tests {\n\t\tt.Run(tt.version, func(t *testing.T) {\n\t\t\trequire.Equal(t, tt.expected, IsPreRelease(tt.version))\n\t\t})\n\t}\n}\n\nfunc TestValidateSDKVersion(t *testing.T) {\n\ttests := []struct {\n\t\tname      string\n\t\tconfig    *WheelConfig\n\t\tlabel     string\n\t\texpectErr bool\n\t\terrMsg    string\n\t}{\n\t\t{name: \"exact minimum valid\", config: &WheelConfig{Source: WheelSourcePyPI, Version: \"0.16.0\"}, label: \"cog\"},\n\t\t{name: \"above minimum valid\", config: &WheelConfig{Source: WheelSourcePyPI, Version: \"0.17.0\"}, label: \"cog\"},\n\t\t{name: \"nil config valid\", config: nil, label: \"cog\"},\n\t\t{name: \"no version pin valid\", config: &WheelConfig{Source: WheelSourcePyPI, Version: \"\"}, label: \"cog\"},\n\t\t{name: \"URL source not checked\", config: &WheelConfig{Source: WheelSourceURL, URL: \"https://example.com/old.whl\"}, label: \"cog\"},\n\t\t{name: \"file source not checked\", config: &WheelConfig{Source: WheelSourceFile, Path: \"/tmp/old.whl\"}, label: \"cog\"},\n\t\t{\n\t\t\tname: \"below minimum errors\", config: &WheelConfig{Source: WheelSourcePyPI, Version: \"0.15.0\"},\n\t\t\tlabel: \"cog\", expectErr: true,\n\t\t\terrMsg: \"cog version 0.15.0 is below the minimum required version 0.16.0\",\n\t\t},\n\t\t{\n\t\t\tname: \"pre-release of old version errors\", config: &WheelConfig{Source: WheelSourcePyPI, Version: \"0.15.0-rc1\"},\n\t\t\tlabel: \"cog\", expectErr: true,\n\t\t\terrMsg: \"cog version 0.15.0-rc1 is below the minimum required version 0.16.0\",\n\t\t},\n\t}\n\tfor _, tt := range tests {\n\t\tt.Run(tt.name, func(t *testing.T) {\n\t\t\terr := ValidateSDKVersion(tt.config, tt.label)\n\t\t\tif tt.expectErr {\n\t\t\t\trequire.Error(t, err)\n\t\t\t\trequire.Equal(t, tt.errMsg, err.Error())\n\t\t\t} else {\n\t\t\t\trequire.NoError(t, err)\n\t\t\t}\n\t\t})\n\t}\n}\n\nfunc TestPyPIPackageURLPreRelease(t *testing.T) {\n\tcfg := &WheelConfig{Source: WheelSourcePyPI, Version: \"0.17.0-alpha1\"}\n\trequire.Equal(t, \"cog==0.17.0a1\", cfg.PyPIPackageURL(\"cog\"))\n}\n"
  },
  {
    "path": "pyproject.toml",
    "content": "[build-system]\nrequires = [\"setuptools\", \"setuptools_scm[toml]\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[project]\nname = \"cog\"\ndescription = \"Containers for machine learning\"\nreadme = \"README.md\"\nauthors = [{ name = \"Replicate\", email = \"team@replicate.com\" }]\nlicense.file = \"LICENSE\"\nurls.\"Source\" = \"https://github.com/replicate/cog\"\n\nclassifiers = [\n  \"Programming Language :: Python :: 3.10\",\n  \"Programming Language :: Python :: 3.11\",\n  \"Programming Language :: Python :: 3.12\",\n  \"Programming Language :: Python :: 3.13\",\n]\nrequires-python = \">=3.10\"\n\ndependencies = [\n  \"typing_extensions>=4.0\",\n  \"pyyaml>=6.0\",\n  \"structlog>=21.0.0\",\n  \"requests>=2.25.0\",\n  \"coglet>=0.1.0,<1.0\",\n]\n\ndynamic = [\"version\"]\n\n[dependency-groups]\ndev = [\n    \"build>=1.2.2.post1\",\n    \"ruff\",\n    \"setuptools-scm>=8.2.0\",\n]\n\ntest = [\n    \"pytest\",\n    \"pytest-timeout\",\n    \"pytest-xdist\",\n    \"pytest-cov\",\n]\n\n[tool.setuptools_scm]\nwrite_to = \"python/cog/_version.py\"\n\n[tool.pyright]\n# TODO: remove this and bring the codebase inline with the current default\nstrictParameterNoneValue = false\n# legacy behavior, fixed in PEP688\ndisableBytesTypePromotions = true\ninclude = [\"python\"]\nexclude = [\"python/tests\"]\nreportMissingParameterType = \"error\"\nreportUnknownLambdaType = \"error\"\nreportUnnecessaryIsInstance = \"warning\"\nreportUnnecessaryComparison = \"warning\"\nreportUnnecessaryContains = \"warning\"\nreportMissingTypeArgument = \"error\"\nreportUnusedExpression = \"warning\"\n\n[tool.pytest.ini_options]\nasyncio_default_fixture_loop_scope = \"function\"\n\n[tool.setuptools]\ninclude-package-data = false\n\n[tool.setuptools.packages.find]\nwhere = [\"python\"]\ninclude = [\"cog*\"]\nexclude = [\"tests*\"]\n\n[tool.pylint.main]\ndisable = [\n  \"C0114\", # Missing module docstring\n  \"C0115\", # Missing class docstring\n  \"C0116\", # Missing function or method docstring\n  \"C0301\", # Line too long\n  \"C0413\", # Import should be placed at the top of the module\n  \"R0903\", # Too few public methods\n  \"W0622\", # Redefining built-in\n]\ngood-names = [\"id\", \"input\"]\n\nignore-paths = [\"python/cog/_version.py\", \"python/tests\"]\n\n[tool.ruff]\nexclude = [\"python/cog/_version.py\"]\nlint.select = [\n  \"E\",   # pycodestyle error\n  \"F\",   # Pyflakes\n  \"I\",   # isort\n  \"W\",   # pycodestyle warning\n  \"S\",   # flake8-bandit\n  \"B\",   # flake8-bugbear\n  \"ANN\", # flake8-annotations\n]\nlint.ignore = [\n  \"E501\",   # Line too long\n  \"S101\",   # Use of `assert` detected\"\n  \"S113\",   # Probable use of requests call without timeout\n  \"B008\",   # Do not perform function call in argument defaults\n  \"ANN001\", # Missing type annotation for function argument\n  \"ANN002\", # Missing type annotation for `*args`\n  \"ANN003\", # Missing type annotation for `**kwargs`\n  \"ANN401\", # Dynamically typed expressions are disallowed\n]\nextend-exclude = [\n  \"python/tests/server/fixtures/*\",\n  \"crates/coglet-python/**/*.pyi\",\n  \"crates/coglet-python/scripts/*\",\n]\nsrc = [\"python\"]\n\n[tool.ruff.lint.per-file-ignores]\n\"python/cog/server/http.py\" = [\n  \"S104\", # Possible binding to all interfaces\n]\n\"python/tests/*\" = [\n  \"S101\", # Use of assert\n  \"S104\", # Possible binding to all interfaces\n  \"S105\", # Possible hardcoded password\n  \"S106\", # Possible hardcoded password in argument\n  \"S108\", # Probable insecure usage of temp file\n  \"S110\", # try-except-pass\n  \"S301\", # pickle can be unsafe when used to deserialize untrusted data\n  \"S603\", # subprocess call — tests use subprocess for isolation\n  \"ANN\",  # Type annotations not required in tests\n  \"B011\", # Do not assert False\n]\n\"crates/coglet-python/tests/*\" = [\n  \"S101\", # Use of assert\n  \"S104\", # Possible binding to all interfaces\n  \"S110\", # try-except-pass\n  \"S603\", # subprocess call\n  \"ANN\",  # Type annotations not required in tests\n  \"B904\", # raise from\n]\n\"tools/test-harness/*\" = [\n  \"S310\", # URL open — URLs are hardcoded to https://\n  \"S603\", # subprocess call — harness invokes cog/docker/git by design\n  \"S607\", # partial executable path — cog/docker/git resolved via PATH\n  \"ANN\",  # Type annotations not required in tooling\n]\n"
  },
  {
    "path": "python/cog/.gitignore",
    "content": "/_version.py\n"
  },
  {
    "path": "python/cog/__init__.py",
    "content": "\"\"\"\nCog SDK: Define machine learning models with standard Python.\n\nThis package provides the core types and classes for building Cog predictors.\n\nExample:\n    from cog import BasePredictor, Input, Path\n\n    class Predictor(BasePredictor):\n        def setup(self):\n            # Load model weights\n            self.model = load_model()\n\n        def predict(\n            self,\n            prompt: str = Input(description=\"Input prompt\"),\n            image: Path = Input(description=\"Input image\"),\n        ) -> str:\n            return self.model.generate(prompt, image)\n\"\"\"\n\nimport sys as _sys\n\nfrom coglet import CancelationException as CancelationException\n\nfrom ._version import __version__\nfrom .input import FieldInfo, Input\nfrom .model import BaseModel\nfrom .predictor import BasePredictor\nfrom .types import (\n    AsyncConcatenateIterator,\n    ConcatenateIterator,\n    File,\n    Path,\n    Secret,\n    URLFile,\n    URLPath,\n)\n\n\n# ---------------------------------------------------------------------------\n# Backwards-compatibility shim: ExperimentalFeatureWarning\n#\n# This class was removed when the Python HTTP server was replaced by coglet.\n# Existing models import it to suppress warnings, e.g.:\n#\n#     from cog import ExperimentalFeatureWarning\n#     warnings.filterwarnings(\"ignore\", category=ExperimentalFeatureWarning)\n#\n# The shim keeps those models working. The stderr message is printed\n# directly so it cannot be swallowed by warnings.filterwarnings(\"ignore\").\n# ---------------------------------------------------------------------------\nclass _ExperimentalFeatureWarning(FutureWarning):\n    \"\"\"Deprecated: ExperimentalFeatureWarning is no longer used by Cog.\n\n    This class exists only for backwards compatibility. Remove the import\n    and any associated ``warnings.filterwarnings(...)`` calls from your code.\n    \"\"\"\n\n    pass\n\n\ndef __getattr__(name: str) -> object:\n    if name == \"ExperimentalFeatureWarning\":\n        print(\n            \"cog: ExperimentalFeatureWarning is deprecated and will be removed in a \"\n            \"future release. Remove `ExperimentalFeatureWarning` from your imports \"\n            \"and any associated `warnings.filterwarnings(...)` calls.\",\n            file=_sys.stderr,\n        )\n        # Cache in module namespace so __getattr__ is not called again and\n        # the deprecation message prints at most once.\n        globals()[\"ExperimentalFeatureWarning\"] = _ExperimentalFeatureWarning\n        return _ExperimentalFeatureWarning\n    if name == \"emit_metric\":\n        print(\n            \"cog: emit_metric() is deprecated and will be removed in a future release. \"\n            \"Use current_scope().record_metric(name, value) instead.\",\n            file=_sys.stderr,\n        )\n\n        def emit_metric(name: str, value: float) -> None:  # noqa: A002 — name is the metric name here, not the module attr\n            current_scope().record_metric(name, value)  # type: ignore[attr-defined]\n\n        # Cache so __getattr__ is not called again — the deprecation message\n        # prints at most once (on first import), not on every call.\n        globals()[\"emit_metric\"] = emit_metric\n        return emit_metric\n    raise AttributeError(f\"module 'cog' has no attribute {name!r}\")\n\n\ndef current_scope() -> object:\n    \"\"\"Get the current prediction scope for recording metrics.\n\n    Returns a Scope object with a ``metrics`` attribute for recording\n    prediction metrics. Outside a prediction context, returns a no-op scope\n    that silently ignores all operations (never ``None``).\n\n    Example::\n\n        from cog import current_scope\n\n        scope = current_scope()\n        scope.record_metric(\"temperature\", 0.7)\n        scope.metrics[\"token_count\"] = 42\n        scope.metrics.record(\"logprobs\", -1.2, mode=\"append\")\n    \"\"\"\n    import coglet\n\n    return coglet._sdk.current_scope()  # type: ignore[attr-defined]  # PyO3 native submodule\n\n\n__all__ = [\n    # Version\n    \"__version__\",\n    # Core classes\n    \"BasePredictor\",\n    \"BaseModel\",\n    # Input\n    \"Input\",\n    \"FieldInfo\",\n    # Types\n    \"Path\",\n    \"Secret\",\n    \"File\",\n    \"URLFile\",\n    \"URLPath\",\n    \"ConcatenateIterator\",\n    \"AsyncConcatenateIterator\",\n    # Exceptions\n    \"CancelationException\",\n    # Metrics\n    \"current_scope\",\n    # Deprecated compat shims\n    \"ExperimentalFeatureWarning\",\n    \"emit_metric\",\n]\n"
  },
  {
    "path": "python/cog/_adt.py",
    "content": "\"\"\"\nInternal ADT (Abstract Data Types) for predictor introspection.\n\nThis module defines the type system used internally for introspecting\npredictor inputs and outputs, generating OpenAPI schemas, and validating\ninput values.\n\"\"\"\n\nimport dataclasses\nimport os\nimport typing\nfrom dataclasses import dataclass\nfrom enum import Enum, auto\nfrom typing import Any, Callable, Dict, List, Optional, Set, Union\n\nfrom .coder import Coder\nfrom .types import File, Path, Secret\n\n\ndef _type_name(tpe: Any) -> str:\n    \"\"\"Get a human-readable name for a type.\"\"\"\n    try:\n        return tpe.__name__\n    except AttributeError:\n        return str(tpe)\n\n\ndef _is_union(tpe: type) -> bool:\n    \"\"\"Check if a type is a Union type.\"\"\"\n    if typing.get_origin(tpe) is Union:\n        return True\n    # Python 3.10+ has UnionType for X | Y syntax\n    from types import UnionType\n\n    if typing.get_origin(tpe) is UnionType:\n        return True\n    return False\n\n\nclass PrimitiveType(Enum):\n    \"\"\"Primitive types supported by Cog.\"\"\"\n\n    BOOL = auto()\n    FLOAT = auto()\n    INTEGER = auto()\n    STRING = auto()\n    PATH = auto()\n    FILE = auto()  # Deprecated, use PATH\n    SECRET = auto()\n    ANY = auto()\n    CUSTOM = auto()\n\n    @staticmethod\n    def _python_type() -> Dict[\"PrimitiveType\", type | Any]:\n        return {\n            PrimitiveType.BOOL: bool,\n            PrimitiveType.FLOAT: float,\n            PrimitiveType.INTEGER: int,\n            PrimitiveType.STRING: str,\n            PrimitiveType.PATH: Path,\n            PrimitiveType.FILE: File,\n            PrimitiveType.SECRET: Secret,\n            PrimitiveType.ANY: Any,\n            PrimitiveType.CUSTOM: Any,\n        }\n\n    @staticmethod\n    def _json_type() -> Dict[\"PrimitiveType\", str]:\n        return {\n            PrimitiveType.BOOL: \"boolean\",\n            PrimitiveType.FLOAT: \"number\",\n            PrimitiveType.INTEGER: \"integer\",\n            PrimitiveType.STRING: \"string\",\n            PrimitiveType.PATH: \"string\",\n            PrimitiveType.FILE: \"string\",\n            PrimitiveType.SECRET: \"string\",\n            PrimitiveType.ANY: \"object\",\n            PrimitiveType.CUSTOM: \"object\",\n        }\n\n    @staticmethod\n    def _adt_type() -> Dict[type | Any, \"PrimitiveType\"]:\n        return {\n            bool: PrimitiveType.BOOL,\n            float: PrimitiveType.FLOAT,\n            int: PrimitiveType.INTEGER,\n            str: PrimitiveType.STRING,\n            Path: PrimitiveType.PATH,\n            File: PrimitiveType.FILE,\n            Secret: PrimitiveType.SECRET,\n            Any: PrimitiveType.ANY,\n        }\n\n    @staticmethod\n    def from_type(tpe: type | Any) -> \"PrimitiveType\":\n        \"\"\"Determine the PrimitiveType for a given Python type.\"\"\"\n        if match := PrimitiveType._adt_type().get(tpe):\n            return match\n\n        try:\n            if tpe is os.PathLike or (\n                isinstance(tpe, type) and issubclass(tpe, os.PathLike)  # type: ignore[arg-type]\n            ):\n                return PrimitiveType.PATH\n        except TypeError:\n            # issubclass raises TypeError for non-class types\n            pass\n\n        return PrimitiveType.CUSTOM\n\n    def normalize(self, value: Any) -> Any:\n        \"\"\"Normalize a value to this primitive type.\"\"\"\n        pt = PrimitiveType._python_type()[self]\n        tpe = type(value)\n\n        if self is PrimitiveType.CUSTOM:\n            return value\n        elif self is PrimitiveType.ANY:\n            return value\n        elif self is PrimitiveType.FILE:\n            # For File inputs, convert URL strings to file-like objects immediately\n            # using File.validate() - the worker won't need to do any conversion\n            import io\n\n            if isinstance(value, io.IOBase):\n                return value\n            # URL string or data URI - validate to file-like object\n            return File.validate(value)\n        elif self is PrimitiveType.PATH:\n            # Convert strings/URLs to Path or URLPath objects\n            if isinstance(value, Path):\n                return value\n            return Path.validate(value)\n        elif self is PrimitiveType.SECRET:\n            # Convert strings to Secret objects\n            if isinstance(value, Secret):\n                return value\n            return Secret(value)\n        else:\n            # Handle enums by extracting their value\n            if issubclass(tpe, Enum):\n                if not issubclass(tpe, pt):\n                    raise ValueError(\n                        f\"enum {_type_name(tpe)} is used as {_type_name(pt)} \"\n                        \"but does not extend it\"\n                    )\n                value = value.value\n            v = pt(value)\n            # For numeric types, allow string coercion (e.g., \"3\" -> 3)\n            # but verify the conversion is valid (not lossy for floats)\n            if v != value:\n                # Allow string to numeric conversion\n                if isinstance(value, str) and pt in (int, float):\n                    return v\n                # Allow int to float conversion\n                if isinstance(value, int) and pt is float:\n                    return v\n                raise ValueError(f\"failed to normalize value as {_type_name(pt)}\")\n            return v\n\n    def python_type_name(self) -> str:\n        \"\"\"Get the Python type name for this primitive.\"\"\"\n        return _type_name(PrimitiveType._python_type()[self])\n\n    def json_type(self) -> Dict[str, Any]:\n        \"\"\"Get the JSON Schema type for this primitive.\"\"\"\n        jt: Dict[str, Any] = {\"type\": self._json_type()[self]}\n        if self in {PrimitiveType.PATH, PrimitiveType.FILE}:\n            jt[\"format\"] = \"uri\"\n        elif self is PrimitiveType.SECRET:\n            jt[\"format\"] = \"password\"\n            jt[\"writeOnly\"] = True\n            jt[\"x-cog-secret\"] = True\n        return jt\n\n    def json_encode(self, value: Any) -> Any:\n        \"\"\"Encode a value for JSON serialization.\"\"\"\n        if self is PrimitiveType.FLOAT:\n            return float(value)\n        elif self in {PrimitiveType.PATH, PrimitiveType.FILE}:\n            return value\n        elif self is PrimitiveType.SECRET:\n            # Secret objects need to be unwrapped for JSON serialization\n            if isinstance(value, Secret):\n                return value.get_secret_value()\n            return value\n        elif self is PrimitiveType.ANY:\n            return value\n        return value\n\n\nclass Repetition(Enum):\n    \"\"\"Field repetition/optionality.\"\"\"\n\n    REQUIRED = 1\n    OPTIONAL = 2\n    REPEATED = 3\n\n\n@dataclass(frozen=True)\nclass FieldType:\n    \"\"\"Type information for an input/output field.\"\"\"\n\n    primitive: PrimitiveType\n    repetition: Repetition\n    coder: Optional[Coder]\n\n    @staticmethod\n    def from_type(tpe: type) -> \"FieldType\":\n        \"\"\"Create a FieldType from a Python type annotation.\"\"\"\n        origin = typing.get_origin(tpe)\n\n        # Handle bare collection types\n        if tpe is list:\n            tpe = List[Any]\n            origin = typing.get_origin(tpe)\n        elif tpe is dict:\n            tpe = Dict[str, Any]\n            origin = typing.get_origin(tpe)\n        elif tpe is set:\n            tpe = Set[Any]\n            origin = typing.get_origin(tpe)\n\n        if origin is dict:\n            # dict / Dict[K, V] → opaque JSON object, consistent with the\n            # static Go schema generator's SchemaAnyType().\n            return FieldType(\n                primitive=PrimitiveType.ANY,\n                repetition=Repetition.REQUIRED,\n                coder=None,\n            )\n\n        if origin in (list, List):\n            t_args = typing.get_args(tpe)\n            if t_args:\n                if len(t_args) != 1:\n                    raise ValueError(\"List must have one type argument\")\n                elem_t = t_args[0]\n                nested_t = typing.get_origin(elem_t)\n                if nested_t is not None:\n                    raise ValueError(\n                        f\"List cannot have nested type {_type_name(nested_t)}\"\n                    )\n            else:\n                elem_t = Any\n            repetition = Repetition.REPEATED\n\n        elif _is_union(tpe):\n            t_args = typing.get_args(tpe)\n            if not (len(t_args) == 2 and type(None) in t_args):\n                raise ValueError(f\"unsupported union type {tpe}\")\n            elem_t = t_args[0] if t_args[1] is type(None) else t_args[1]\n            nested_t = typing.get_origin(elem_t)\n            if nested_t is not None:\n                raise ValueError(\n                    f\"Optional cannot have nested type {_type_name(nested_t)}\"\n                )\n            repetition = Repetition.OPTIONAL\n\n        else:\n            elem_t = tpe\n            repetition = Repetition.REQUIRED\n\n        cog_t = PrimitiveType.from_type(elem_t)\n        coder = None\n        if cog_t is PrimitiveType.CUSTOM:\n            coder = Coder.lookup(elem_t)\n            if coder is None:\n                raise ValueError(f\"unsupported Cog type {_type_name(elem_t)}\")\n\n        return FieldType(primitive=cog_t, repetition=repetition, coder=coder)\n\n    def normalize(self, value: Any) -> Any:\n        \"\"\"Normalize a value according to this field type.\"\"\"\n        if self.repetition is Repetition.REQUIRED:\n            return self.primitive.normalize(value)\n        elif self.repetition is Repetition.OPTIONAL:\n            return None if value is None else self.primitive.normalize(value)\n        elif self.repetition is Repetition.REPEATED:\n            return [self.primitive.normalize(v) for v in value]\n        return value\n\n    def json_type(self) -> Dict[str, Any]:\n        \"\"\"Get the JSON Schema type for this field.\"\"\"\n        if self.repetition is Repetition.REPEATED:\n            return {\"type\": \"array\", \"items\": self.primitive.json_type()}\n        return self.primitive.json_type()\n\n    def json_encode(self, value: Any) -> Any:\n        \"\"\"Encode a value for JSON serialization.\"\"\"\n        f: Callable[[Any], Any] = self.primitive.json_encode\n        if self.primitive is PrimitiveType.CUSTOM:\n            assert self.coder is not None\n            f = self.coder.encode\n        if self.repetition is Repetition.REPEATED:\n            return [f(x) for x in value]\n        return f(value)\n\n    def json_decode(self, value: Any) -> Any:\n        \"\"\"Decode a value from JSON.\"\"\"\n        if self.primitive is not PrimitiveType.CUSTOM:\n            return value\n        assert self.coder is not None\n        f = self.coder.decode\n        if self.repetition is Repetition.REPEATED:\n            return [f(x) for x in value]\n        return f(value)\n\n    def python_type_name(self) -> str:\n        \"\"\"Get the Python type name for this field.\"\"\"\n        if self.repetition is Repetition.REQUIRED:\n            return self.primitive.python_type_name()\n        elif self.repetition is Repetition.OPTIONAL:\n            return f\"Optional[{self.primitive.python_type_name()}]\"\n        elif self.repetition is Repetition.REPEATED:\n            return f\"List[{self.primitive.python_type_name()}]\"\n        return self.primitive.python_type_name()\n\n\n@dataclass(frozen=True)\nclass InputField:\n    \"\"\"Metadata for a predictor input field.\"\"\"\n\n    name: str\n    order: int\n    type: FieldType\n    default: Any = None\n    description: Optional[str] = None\n    ge: Optional[Union[int, float]] = None\n    le: Optional[Union[int, float]] = None\n    min_length: Optional[int] = None\n    max_length: Optional[int] = None\n    regex: Optional[str] = None\n    choices: Optional[List[Union[str, int]]] = None\n    deprecated: Optional[bool] = None\n\n\nclass OutputKind(Enum):\n    \"\"\"Kind of output a predictor produces.\"\"\"\n\n    SINGLE = 1\n    LIST = 2\n    ITERATOR = 3\n    CONCAT_ITERATOR = 4\n    OBJECT = 5\n\n\n@dataclass(frozen=True)\nclass OutputType:\n    \"\"\"Type information for predictor output.\"\"\"\n\n    kind: OutputKind\n    type: Optional[PrimitiveType] = None\n    fields: Optional[Dict[str, FieldType]] = None\n    coder: Optional[Coder] = None\n\n    def json_type(self) -> Dict[str, Any]:\n        \"\"\"Get the JSON Schema type for this output.\"\"\"\n        jt: Dict[str, Any] = {\"title\": \"Output\"}\n\n        if self.kind is OutputKind.SINGLE:\n            assert self.type is not None\n            jt.update(self.type.json_type())\n\n        elif self.kind is OutputKind.LIST:\n            assert self.type is not None\n            jt.update({\"type\": \"array\", \"items\": self.type.json_type()})\n\n        elif self.kind is OutputKind.ITERATOR:\n            assert self.type is not None\n            jt.update(\n                {\n                    \"type\": \"array\",\n                    \"items\": self.type.json_type(),\n                    \"x-cog-array-type\": \"iterator\",\n                }\n            )\n\n        elif self.kind is OutputKind.CONCAT_ITERATOR:\n            assert self.type is not None\n            jt.update(\n                {\n                    \"type\": \"array\",\n                    \"items\": self.type.json_type(),\n                    \"x-cog-array-type\": \"iterator\",\n                    \"x-cog-array-display\": \"concatenate\",\n                }\n            )\n\n        elif self.kind is OutputKind.OBJECT:\n            assert self.fields is not None\n            props = {}\n            for name, field_type in self.fields.items():\n                props[name] = field_type.primitive.json_type()\n                props[name][\"title\"] = name.replace(\"_\", \" \").title()\n            jt.update(\n                {\n                    \"type\": \"object\",\n                    \"properties\": props,\n                    \"required\": list(self.fields.keys()),\n                }\n            )\n\n        return jt\n\n    def normalize(self, value: Any) -> Any:\n        \"\"\"Normalize an output value.\"\"\"\n        return self._transform(value, json=False)\n\n    def json_encode(self, value: Any) -> Any:\n        \"\"\"Encode an output value for JSON serialization.\"\"\"\n        if self.coder is not None:\n            if self.kind is OutputKind.LIST:\n                return [self.coder.encode(x) for x in value]\n            return self.coder.encode(value)\n\n        o = self._transform(value, json=True)\n        if self.kind is OutputKind.OBJECT:\n            # Expand dataclass to dict\n            tpe = type(o)\n            if not dataclasses.is_dataclass(tpe):\n                raise ValueError(f\"{tpe} is not a dataclass\")\n            return {f.name: getattr(o, f.name) for f in dataclasses.fields(o)}\n        return o\n\n    def _transform(self, value: Any, json: bool) -> Any:\n        \"\"\"Transform an output value (normalize or encode).\"\"\"\n        if self.kind in {\n            OutputKind.SINGLE,\n            OutputKind.ITERATOR,\n            OutputKind.CONCAT_ITERATOR,\n        }:\n            assert self.type is not None\n            f: Callable[[Any], Any] = (\n                self.type.json_encode if json else self.type.normalize\n            )\n            return f(value)\n\n        elif self.kind is OutputKind.LIST:\n            assert self.type is not None\n            f = self.type.json_encode if json else self.type.normalize\n            return [f(x) for x in value]\n\n        elif self.kind is OutputKind.OBJECT:\n            assert self.fields is not None\n            for name, ft in self.fields.items():\n                f = ft.json_encode if json else ft.normalize\n                if not hasattr(value, name):\n                    raise ValueError(f\"missing output field: {name}\")\n                v = getattr(value, name)\n                if v is None:\n                    if ft.repetition is not Repetition.OPTIONAL:\n                        raise ValueError(f\"missing value for output field: {name}\")\n                setattr(value, name, f(v))\n            return value\n\n        raise RuntimeError(f\"unsupported output kind {self.kind}\")\n\n\n@dataclass(frozen=True)\nclass PredictorInfo:\n    \"\"\"Complete type information for a predictor.\"\"\"\n\n    module_name: str\n    predictor_name: str\n    inputs: Dict[str, InputField]\n    output: OutputType\n"
  },
  {
    "path": "python/cog/_inspector.py",
    "content": "\"\"\"\nInternal inspector for predictor introspection.\n\nThis module provides functions to inspect predictor classes and functions,\nextract input/output type information, and validate inputs.\n\"\"\"\n\nimport importlib\nimport inspect\nimport re\nimport sys\nimport typing\nfrom dataclasses import MISSING, Field\nfrom enum import Enum\nfrom types import ModuleType\nfrom typing import Any, AsyncIterator, Callable, Dict, Iterator, Type\n\nfrom . import _adt as adt\nfrom .coder import Coder\nfrom .input import FieldInfo\nfrom .model import BaseModel\nfrom .types import AsyncConcatenateIterator, ConcatenateIterator\n\ntry:\n    from pydantic import (  # pyright: ignore[reportMissingImports]\n        BaseModel as PydanticBaseModel,\n    )\nexcept ImportError:\n    PydanticBaseModel = None  # type: ignore[assignment,misc]\n\n\ndef _check_parent(child: type, parent: type) -> bool:\n    \"\"\"Check if a type has a parent in its MRO.\"\"\"\n    return any(c is parent for c in inspect.getmro(child))\n\n\ndef _type_name(tpe: Any) -> str:\n    \"\"\"Get a human-readable name for a type.\"\"\"\n    try:\n        return tpe.__name__\n    except AttributeError:\n        return str(tpe)\n\n\ndef _validate_setup(f: Callable[..., Any]) -> None:\n    \"\"\"Validate a predictor's setup method.\"\"\"\n    if not inspect.isfunction(f):\n        raise ValueError(\"setup is not a function\")\n\n    spec = inspect.getfullargspec(f)\n\n    if spec.args[:1] != [\"self\"]:\n        raise ValueError(\"setup() must have 'self' as first argument\")\n\n    non_default_args = spec.args\n    if spec.defaults is not None:\n        non_default_args = non_default_args[: -len(spec.defaults)]\n\n    extra_args = [a for a in non_default_args if a not in {\"self\", \"weights\"}]\n    if extra_args:\n        raise ValueError(f\"unexpected setup() arguments: {', '.join(extra_args)}\")\n\n    if spec.varargs is not None:\n        raise ValueError(\"setup() must not have *args\")\n    if spec.varkw is not None:\n        raise ValueError(\"setup() must not have **kwargs\")\n    if spec.kwonlyargs:\n        raise ValueError(\"setup() must not have keyword-only args\")\n    if spec.kwonlydefaults:\n        raise ValueError(\"setup() must not have keyword-only defaults\")\n    if spec.annotations.get(\"return\") is not None:\n        raise ValueError(\"setup() must return None\")\n\n\ndef _validate_predict(f: Callable[..., Any], f_name: str, is_class_fn: bool) -> None:\n    \"\"\"Validate a predictor's predict method.\"\"\"\n    if not inspect.isfunction(f):\n        raise ValueError(f\"{f_name} is not a function\")\n\n    spec = inspect.getfullargspec(f)\n\n    if is_class_fn and spec.args[:1] != [\"self\"]:\n        raise ValueError(f\"{f_name}() must have 'self' as first argument\")\n    if spec.varargs is not None:\n        raise ValueError(f\"{f_name}() must not have *args\")\n    if spec.varkw is not None:\n        raise ValueError(f\"{f_name}() must not have **kwargs\")\n    if spec.kwonlyargs:\n        raise ValueError(f\"{f_name}() must not have keyword-only args\")\n    if spec.kwonlydefaults:\n        raise ValueError(f\"{f_name}() must not have keyword-only defaults\")\n    if spec.annotations.get(\"return\") is None:\n        raise ValueError(f\"{f_name}() must have a return type annotation\")\n\n\ndef _validate_input_constraints(\n    name: str, ft: adt.FieldType, field_info: FieldInfo\n) -> None:\n    \"\"\"Validate that FieldInfo constraints are compatible with the field type.\"\"\"\n    cog_t = ft.primitive\n    in_repr = f\"{name}: {ft.python_type_name()}\"\n\n    # Extract actual default for validation\n    defaults = []\n    if field_info.default is not None:\n        if isinstance(field_info.default, Field):\n            if field_info.default.default_factory is not MISSING:\n                actual_default = field_info.default.default_factory()\n            elif field_info.default.default is not MISSING:\n                actual_default = field_info.default.default\n            else:\n                actual_default = None\n        else:\n            actual_default = field_info.default\n\n        if actual_default is not None:\n            if ft.repetition is adt.Repetition.REPEATED:\n                defaults = ft.normalize(actual_default)\n            else:\n                defaults = [ft.normalize(actual_default)]\n\n    numeric_types = {adt.PrimitiveType.FLOAT, adt.PrimitiveType.INTEGER}\n\n    # Validate ge/le constraints\n    if field_info.ge is not None or field_info.le is not None:\n        if cog_t not in numeric_types:\n            raise ValueError(f\"incompatible input type for ge/le: {in_repr}\")\n        if defaults:\n            if field_info.ge is not None and not all(\n                x >= field_info.ge for x in defaults\n            ):\n                raise ValueError(\n                    f\"invalid default for {in_repr}: must be at minimum {field_info.ge}\"\n                )\n            if field_info.le is not None and not all(\n                x <= field_info.le for x in defaults\n            ):\n                raise ValueError(\n                    f\"invalid default for {in_repr}: must be at maximum {field_info.le}\"\n                )\n\n    # Validate min_length/max_length constraints\n    if field_info.min_length is not None or field_info.max_length is not None:\n        if cog_t is not adt.PrimitiveType.STRING:\n            raise ValueError(\n                f\"incompatible input type for min_length/max_length: {in_repr}\"\n            )\n        if defaults:\n            if field_info.min_length is not None and not all(\n                len(x) >= field_info.min_length for x in defaults\n            ):\n                raise ValueError(\n                    f\"default conflicts with min_length={field_info.min_length} for input: {in_repr}\"\n                )\n            if field_info.max_length is not None and not all(\n                len(x) <= field_info.max_length for x in defaults\n            ):\n                raise ValueError(\n                    f\"default conflicts with max_length={field_info.max_length} for input: {in_repr}\"\n                )\n\n    # Validate regex constraint\n    if field_info.regex is not None:\n        if cog_t is not adt.PrimitiveType.STRING:\n            raise ValueError(f\"incompatible input type for regex: {in_repr}\")\n        if defaults:\n            regex = re.compile(field_info.regex)\n            if not all(regex.match(x) for x in defaults):\n                raise ValueError(f\"default not a regex match for input: {in_repr}\")\n\n    # Validate choices constraint\n    if field_info.choices is not None:\n        choice_types = {adt.PrimitiveType.INTEGER, adt.PrimitiveType.STRING}\n        if cog_t not in choice_types:\n            raise ValueError(f\"incompatible input type for choices: {in_repr}\")\n        if len(field_info.choices) < 2:\n            raise ValueError(\n                f\"choices={field_info.choices!r} must have >= 2 elements: {in_repr}\"\n            )\n        if field_info.ge is not None or field_info.le is not None:\n            raise ValueError(f\"choices and ge/le are mutually exclusive: {in_repr}\")\n        if field_info.min_length is not None or field_info.max_length is not None:\n            raise ValueError(\n                f\"choices and min_length/max_length are mutually exclusive: {in_repr}\"\n            )\n        # Normalize enum values in choices\n        choices = [\n            cog_t.normalize(c) if isinstance(c, Enum) else c for c in field_info.choices\n        ]\n        if not all(adt.PrimitiveType.from_type(type(x)) is cog_t for x in choices):\n            raise ValueError(f\"not all choices have the same type as input: {in_repr}\")\n\n\ndef _create_input_field(\n    order: int, name: str, tpe: type, field_info: Any\n) -> adt.InputField:\n    \"\"\"Create an InputField from type annotation and optional FieldInfo or raw default.\"\"\"\n    try:\n        ft = adt.FieldType.from_type(tpe)\n    except (ValueError, AssertionError) as e:\n        raise ValueError(f\"invalid input field {name}: {e}\") from e\n\n    if field_info is None:\n        return adt.InputField(name=name, order=order, type=ft)\n\n    # Handle raw default values (not FieldInfo)\n    if not isinstance(field_info, FieldInfo):\n        # It's a raw default value like \"world\" or 42\n        default = ft.normalize(field_info) if field_info is not None else None\n        return adt.InputField(name=name, order=order, type=ft, default=default)\n\n    _validate_input_constraints(name, ft, field_info)\n\n    # Extract default value\n    if isinstance(field_info.default, Field):\n        if field_info.default.default_factory is not MISSING:\n            default = field_info.default\n        elif field_info.default.default is not MISSING:\n            default = ft.normalize(field_info.default.default)\n        else:\n            default = None\n    else:\n        default = (\n            None if field_info.default is None else ft.normalize(field_info.default)\n        )\n\n    # Normalize choices\n    choices = (\n        None\n        if field_info.choices is None\n        else [ft.primitive.normalize(c) for c in field_info.choices]\n    )\n\n    return adt.InputField(\n        name=name,\n        order=order,\n        type=ft,\n        default=default,\n        description=field_info.description,\n        ge=float(field_info.ge) if field_info.ge is not None else None,\n        le=float(field_info.le) if field_info.le is not None else None,\n        min_length=field_info.min_length,\n        max_length=field_info.max_length,\n        regex=field_info.regex,\n        choices=choices,\n        deprecated=field_info.deprecated,\n    )\n\n\nclass _AnyType:\n    \"\"\"Placeholder type for Any output (for compatibility).\"\"\"\n\n    @staticmethod\n    def normalize(value: Any) -> Any:\n        return value\n\n    @staticmethod\n    def json_type() -> Dict[str, Any]:\n        return {}\n\n    @staticmethod\n    def json_encode(value: Any) -> Any:\n        return value\n\n\n_any_type = _AnyType()\n\n\ndef _create_output_type(tpe: type) -> adt.OutputType:\n    \"\"\"Create an OutputType from a return type annotation.\"\"\"\n    if tpe is Any:\n        print(\n            \"Warning: use of Any as output type is error-prone and highly discouraged\"\n        )\n        return adt.OutputType(kind=adt.OutputKind.SINGLE, type=_any_type)  # type: ignore[arg-type]\n\n    if inspect.isclass(tpe) and _check_parent(tpe, BaseModel):\n        fields = {}\n        for name, t in tpe.__annotations__.items():\n            ft = adt.FieldType.from_type(t)\n            fields[name] = ft\n        return adt.OutputType(kind=adt.OutputKind.OBJECT, fields=fields)\n\n    if (\n        PydanticBaseModel is not None\n        and inspect.isclass(tpe)\n        and _check_parent(tpe, PydanticBaseModel)\n    ):\n        fields = {}\n        for name, field_info in tpe.model_fields.items():\n            ft = adt.FieldType.from_type(field_info.annotation)\n            fields[name] = ft\n        return adt.OutputType(kind=adt.OutputKind.OBJECT, fields=fields)\n\n    origin = typing.get_origin(tpe)\n    concat_iters = {ConcatenateIterator, AsyncConcatenateIterator}\n\n    if origin in {typing.get_origin(Iterator), typing.get_origin(AsyncIterator)}:\n        kind = adt.OutputKind.ITERATOR\n        t_args = typing.get_args(tpe)\n        if len(t_args) != 1:\n            raise ValueError(\"iterator type must have a type argument\")\n        ft = adt.FieldType.from_type(t_args[0])\n        if ft.repetition is not adt.Repetition.REQUIRED:\n            raise ValueError(\"iterator element type must not be Optional or List\")\n\n    elif origin in concat_iters or tpe in concat_iters:\n        kind = adt.OutputKind.CONCAT_ITERATOR\n        t_args = typing.get_args(tpe)\n        if len(t_args) != 1:\n            raise ValueError(\"iterator type must have a type argument\")\n        ft = adt.FieldType.from_type(t_args[0])\n        if ft.repetition is not adt.Repetition.REQUIRED:\n            raise ValueError(\"iterator element type must not be Optional or List\")\n        if ft.primitive is not adt.PrimitiveType.STRING:\n            raise ValueError(f\"{_type_name(tpe)} must have str element\")\n\n    else:\n        ft = adt.FieldType.from_type(tpe)\n        if ft.repetition is adt.Repetition.OPTIONAL:\n            raise ValueError(\"output must not be Optional\")\n        if ft.repetition == adt.Repetition.REQUIRED:\n            kind = adt.OutputKind.SINGLE\n        elif ft.repetition == adt.Repetition.REPEATED:\n            kind = adt.OutputKind.LIST\n        else:\n            raise RuntimeError(f\"unexpected repetition: {ft.repetition}\")\n\n    return adt.OutputType(kind=kind, type=ft.primitive, coder=ft.coder)\n\n\ndef _create_predictor_info(\n    module_name: str,\n    predictor_name: str,\n    f: Callable[..., Any],\n    f_name: str,\n    is_class_fn: bool,\n) -> adt.PredictorInfo:\n    \"\"\"Create PredictorInfo from a predict function.\"\"\"\n    _validate_predict(f, f_name, is_class_fn)\n    spec = inspect.getfullargspec(f)\n\n    # Use get_type_hints to resolve string annotations (from __future__ import annotations)\n    try:\n        type_hints = typing.get_type_hints(f)\n    except Exception:\n        # Fall back to raw annotations if get_type_hints fails\n        type_hints = spec.annotations\n\n    # Skip 'self' for class methods\n    names = spec.args[1:] if is_class_fn else spec.args\n    defaults = list(spec.defaults) if spec.defaults else []\n    field_infos = [None] * (len(names) - len(defaults)) + defaults\n\n    inputs: Dict[str, adt.InputField] = {}\n    for i, (name, field_info) in enumerate(zip(names, field_infos, strict=False)):\n        tpe = type_hints.get(name)\n        if tpe is None:\n            raise ValueError(f\"missing type annotation for input: {name}\")\n        inputs[name] = _create_input_field(i, name, tpe, field_info)\n\n    return_type = type_hints.get(\"return\", spec.annotations.get(\"return\"))\n    if return_type is None:\n        raise ValueError(\"missing return type annotation for predict method\")\n    output = _create_output_type(return_type)\n    return adt.PredictorInfo(module_name, predictor_name, inputs, output)\n\n\ndef _unwrap(f: Callable[..., Any]) -> Callable[..., Any]:\n    \"\"\"Unwrap decorated functions to get the original function.\"\"\"\n    g = f\n    while hasattr(g, \"__closure__\") and g.__closure__ is not None:\n        cs = [\n            c.cell_contents\n            for c in g.__closure__\n            if inspect.isfunction(c.cell_contents)\n        ]\n        if len(cs) > 1:\n            raise ValueError(f\"unable to inspect function decorator: {f}\")\n        if len(cs) == 0:\n            return g\n        g = cs[0]\n    return g\n\n\ndef _is_coder(cls: Type[Any]) -> bool:\n    \"\"\"Check if a class is a Coder subclass.\"\"\"\n    return inspect.isclass(cls) and cls is not Coder and _check_parent(cls, Coder)\n\n\ndef _find_coders(module: ModuleType) -> None:\n    \"\"\"Find and register coders defined in a module.\"\"\"\n    # Direct imports: from cog.coders.some_coder import SomeCoder\n    for _, c in inspect.getmembers(module, _is_coder):\n        Coder.register(c)\n\n    # Module imports: from cog.coders import some_coders\n    for _, m in inspect.getmembers(module, inspect.ismodule):\n        for _, c in inspect.getmembers(m, _is_coder):\n            Coder.register(c)\n\n\ndef create_predictor(module_name: str, predictor_name: str) -> adt.PredictorInfo:\n    \"\"\"\n    Create PredictorInfo by inspecting a predictor class or function.\n\n    Args:\n        module_name: The module containing the predictor\n        predictor_name: The name of the predictor class or function\n\n    Returns:\n        PredictorInfo with input/output type information\n    \"\"\"\n    try:\n        module = importlib.import_module(module_name)\n    except (ImportError, ModuleNotFoundError) as e:\n        raise ImportError(f\"failed to import predictor module: {e}\") from e\n\n    fullname = f\"{module_name}.{predictor_name}\"\n    if not hasattr(module, predictor_name):\n        # Check if module is partially loaded (common with import errors)\n        if module_name in sys.modules:\n            raise ImportError(\n                f\"predictor {predictor_name} not found in {module_name} \"\n                \"(module may have import errors)\"\n            )\n        raise ValueError(f\"predictor not found: {fullname}\")\n\n    p = getattr(module, predictor_name)\n\n    if inspect.isclass(p):\n        if not hasattr(p, \"predict\"):\n            raise ValueError(f\"predict method not found: {fullname}\")\n\n        if hasattr(p, \"setup\"):\n            _validate_setup(_unwrap(p.setup))\n\n        predict_fn_name = \"predict\"\n        predict_fn = _unwrap(getattr(p, predict_fn_name))\n        is_class_fn = True\n\n    elif inspect.isfunction(p):\n        predict_fn_name = predictor_name\n        predict_fn = _unwrap(p)\n        is_class_fn = False\n\n    else:\n        raise ValueError(f\"invalid predictor {fullname}\")\n\n    # Find coders before validating predict function\n    _find_coders(module)\n\n    return _create_predictor_info(\n        module_name, predictor_name, predict_fn, predict_fn_name, is_class_fn\n    )\n\n\ndef check_input(\n    inputs: Dict[str, adt.InputField], values: Dict[str, Any]\n) -> Dict[str, Any]:\n    \"\"\"\n    Validate and normalize input values against InputField definitions.\n\n    Args:\n        inputs: Dictionary of InputField definitions\n        values: Dictionary of input values to validate\n\n    Returns:\n        Dictionary of normalized input values\n    \"\"\"\n    kwargs: Dict[str, Any] = {}\n\n    # Process provided values\n    for name, value in values.items():\n        input_field = inputs.get(name)\n        if input_field is None:\n            print(f\"WARNING unknown input field ignored: {name}\")\n        else:\n            try:\n                kwargs[name] = input_field.type.normalize(value)\n            except ValueError as e:\n                # Reformat normalize errors to use \"field: message\" format\n                # and avoid leaking user input values\n                msg = str(e)\n                if \"failed to normalize value\" in msg:\n                    # Extract just the type name without the value\n                    if \" as \" in msg:\n                        type_name = msg.split(\" as \", 1)[1]\n                        raise ValueError(\n                            f\"{name}: Invalid value for type {type_name}\"\n                        ) from None\n                    raise ValueError(f\"{name}: Invalid value\") from None\n                # For other normalize errors, prepend field name\n                raise ValueError(f\"{name}: {msg}\") from None\n\n    # Apply defaults for missing values\n    for name, input_field in inputs.items():\n        if name not in kwargs:\n            if isinstance(input_field.default, Field):\n                if input_field.default.default_factory is not MISSING:\n                    kwargs[name] = input_field.default.default_factory()\n                elif input_field.default.default is not MISSING:\n                    kwargs[name] = input_field.default.default\n                else:\n                    if input_field.type.repetition is not adt.Repetition.OPTIONAL:\n                        raise ValueError(f\"{name}: Field required\")\n                    kwargs[name] = None\n            elif input_field.default is not None:\n                kwargs[name] = input_field.default\n            else:\n                if input_field.type.repetition is not adt.Repetition.OPTIONAL:\n                    raise ValueError(f\"{name}: Field required\")\n                kwargs[name] = None\n\n        # Validate constraints\n        v = kwargs[name]\n        values_to_check = []\n        if input_field.type.repetition is adt.Repetition.REQUIRED:\n            values_to_check = [v]\n        elif input_field.type.repetition is adt.Repetition.OPTIONAL:\n            values_to_check = [] if v is None else [v]\n        elif input_field.type.repetition is adt.Repetition.REPEATED:\n            values_to_check = v\n\n        if input_field.ge is not None:\n            if not all(x >= input_field.ge for x in values_to_check):\n                raise ValueError(\n                    f\"{name} fails constraint >= {int(input_field.ge) if input_field.ge == int(input_field.ge) else input_field.ge}\"\n                )\n\n        if input_field.le is not None:\n            if not all(x <= input_field.le for x in values_to_check):\n                raise ValueError(\n                    f\"{name} fails constraint <= {int(input_field.le) if input_field.le == int(input_field.le) else input_field.le}\"\n                )\n\n        if input_field.min_length is not None:\n            if not all(len(x) >= input_field.min_length for x in values_to_check):\n                raise ValueError(\n                    f\"{name} fails constraint len() >= {input_field.min_length}\"\n                )\n\n        if input_field.max_length is not None:\n            if not all(len(x) <= input_field.max_length for x in values_to_check):\n                raise ValueError(\n                    f\"{name} fails constraint len() <= {input_field.max_length}\"\n                )\n\n        if input_field.regex is not None:\n            p = re.compile(input_field.regex)\n            if not all(p.match(x) is not None for x in values_to_check):\n                raise ValueError(f\"{name} does not match regex {input_field.regex!r}\")\n\n        if input_field.choices is not None:\n            if not all(x in input_field.choices for x in values_to_check):\n                raise ValueError(\n                    f\"{name} does not match choices {input_field.choices!r}\"\n                )\n\n    return kwargs\n"
  },
  {
    "path": "python/cog/_schemas.py",
    "content": "\"\"\"\nInternal schema generation for OpenAPI.\n\nThis module provides functions to generate OpenAPI JSON schemas from\nPredictorInfo.\n\"\"\"\n\nfrom dataclasses import MISSING, Field\nfrom typing import Any, Dict\n\nfrom . import _adt as adt\nfrom .mode import Mode\n\n\ndef to_json_input(predictor: adt.PredictorInfo) -> Dict[str, Any]:\n    \"\"\"Generate OpenAPI schema for predictor inputs.\"\"\"\n    schema: Dict[str, Any] = {\n        \"properties\": {},\n        \"type\": \"object\",\n        \"title\": \"Input\",\n    }\n    required = []\n\n    for name, input_field in predictor.inputs.items():\n        prop: Dict[str, Any] = {\"x-order\": input_field.order}\n\n        if input_field.choices is not None:\n            prop[\"allOf\"] = [{\"$ref\": f\"#/components/schemas/{name}\"}]\n        else:\n            prop[\"title\"] = name.replace(\"_\", \" \").title()\n            prop.update(input_field.type.json_type())\n\n        # Determine required status and default value:\n        # - name: type = Input() -> required\n        # - name: type = Input(default=value) -> not required, has default\n        # - name: Optional[type] = Input() -> not required, default None\n        # - name: Optional[type] = Input(default=value) -> not required, has default\n        # - name: list[type] = Input() -> required\n        # - name: list[type] = Input(default=[...]) -> not required, has default\n\n        if input_field.default is None:\n            if input_field.type.repetition in {\n                adt.Repetition.REQUIRED,\n                adt.Repetition.REPEATED,\n            }:\n                required.append(name)\n        else:\n            # Extract actual default for schema\n            if isinstance(input_field.default, Field):\n                if input_field.default.default_factory is not MISSING:\n                    actual_default = input_field.default.default_factory()\n                elif input_field.default.default is not MISSING:\n                    actual_default = input_field.default.default\n                else:\n                    actual_default = None\n            else:\n                actual_default = input_field.default\n\n            if actual_default is not None:\n                normalized = input_field.type.normalize(actual_default)\n                prop[\"default\"] = input_field.type.json_encode(normalized)\n\n        # Optional types are nullable\n        if input_field.type.repetition is adt.Repetition.OPTIONAL:\n            prop[\"nullable\"] = True\n\n        # Add constraints\n        if input_field.description is not None:\n            prop[\"description\"] = input_field.description\n        if input_field.ge is not None:\n            prop[\"minimum\"] = input_field.ge\n        if input_field.le is not None:\n            prop[\"maximum\"] = input_field.le\n        if input_field.min_length is not None:\n            prop[\"minLength\"] = input_field.min_length\n        if input_field.max_length is not None:\n            prop[\"maxLength\"] = input_field.max_length\n        if input_field.regex is not None:\n            prop[\"pattern\"] = input_field.regex\n        if input_field.deprecated is not None:\n            prop[\"deprecated\"] = input_field.deprecated\n\n        schema[\"properties\"][name] = prop\n\n    if required:\n        schema[\"required\"] = required\n\n    return schema\n\n\ndef to_json_enums(predictor: adt.PredictorInfo) -> Dict[str, Any]:\n    \"\"\"Generate OpenAPI schema for enum inputs (choices).\"\"\"\n    enums = {}\n\n    for name, input_field in predictor.inputs.items():\n        if input_field.choices is None:\n            continue\n\n        enum_schema = {\n            \"title\": name,\n            \"description\": \"An enumeration.\",\n            \"enum\": input_field.choices,\n        }\n        enum_schema.update(input_field.type.primitive.json_type())\n        enums[name] = enum_schema\n\n    return enums\n\n\ndef to_json_output(predictor: adt.PredictorInfo) -> Dict[str, Any]:\n    \"\"\"Generate OpenAPI schema for predictor output.\"\"\"\n    return predictor.output.json_type()\n\n\ndef to_json_schema(\n    predictor: adt.PredictorInfo, mode: Mode = Mode.PREDICT\n) -> Dict[str, Any]:\n    \"\"\"\n    Generate a complete OpenAPI schema for a predictor.\n\n    This creates the full OpenAPI specification with Input, Output,\n    and enum schemas populated from the predictor info.\n\n    Args:\n        predictor: The predictor info to generate schema from\n        mode: The prediction mode (Mode.PREDICT or Mode.TRAIN)\n    \"\"\"\n    # Determine routes and schema names based on mode\n    if mode == Mode.TRAIN:\n        main_route = \"/trainings\"\n        cancel_route = \"/trainings/{training_id}/cancel\"\n        request_schema = \"TrainingRequest\"\n        response_schema = \"TrainingResponse\"\n        id_param_name = \"training_id\"\n        id_param_title = \"Training Id\"\n        summary = \"Train\"\n        description = \"Run a training session\"\n        operation_id = \"train_trainings_post\"\n        cancel_operation_id = \"cancel_trainings__training_id__cancel_post\"\n    else:\n        main_route = \"/predictions\"\n        cancel_route = \"/predictions/{prediction_id}/cancel\"\n        request_schema = \"PredictionRequest\"\n        response_schema = \"PredictionResponse\"\n        id_param_name = \"prediction_id\"\n        id_param_title = \"Prediction Id\"\n        summary = \"Predict\"\n        description = \"Run a single prediction on the model\"\n        operation_id = \"predict_predictions_post\"\n        cancel_operation_id = \"cancel_predictions__prediction_id__cancel_post\"\n\n    # Base OpenAPI schema structure\n    schema: Dict[str, Any] = {\n        \"openapi\": \"3.0.2\",\n        \"info\": {\"title\": \"Cog\", \"version\": \"0.1.0\"},\n        \"paths\": {\n            \"/\": {\n                \"get\": {\n                    \"summary\": \"Root\",\n                    \"operationId\": \"root__get\",\n                    \"responses\": {\n                        \"200\": {\n                            \"description\": \"Successful Response\",\n                            \"content\": {\"application/json\": {\"schema\": {}}},\n                        }\n                    },\n                }\n            },\n            \"/health-check\": {\n                \"get\": {\n                    \"summary\": \"Healthcheck\",\n                    \"operationId\": \"healthcheck_health_check_get\",\n                    \"responses\": {\n                        \"200\": {\n                            \"description\": \"Successful Response\",\n                            \"content\": {\"application/json\": {\"schema\": {}}},\n                        }\n                    },\n                }\n            },\n            main_route: {\n                \"post\": {\n                    \"summary\": summary,\n                    \"description\": description,\n                    \"operationId\": operation_id,\n                    \"requestBody\": {\n                        \"content\": {\n                            \"application/json\": {\n                                \"schema\": {\n                                    \"$ref\": f\"#/components/schemas/{request_schema}\"\n                                }\n                            }\n                        }\n                    },\n                    \"responses\": {\n                        \"200\": {\n                            \"description\": \"Successful Response\",\n                            \"content\": {\n                                \"application/json\": {\n                                    \"schema\": {\n                                        \"$ref\": f\"#/components/schemas/{response_schema}\"\n                                    }\n                                }\n                            },\n                        },\n                        \"422\": {\n                            \"description\": \"Validation Error\",\n                            \"content\": {\n                                \"application/json\": {\n                                    \"schema\": {\n                                        \"$ref\": \"#/components/schemas/HTTPValidationError\"\n                                    }\n                                }\n                            },\n                        },\n                    },\n                }\n            },\n            cancel_route: {\n                \"post\": {\n                    \"summary\": \"Cancel\",\n                    \"operationId\": cancel_operation_id,\n                    \"parameters\": [\n                        {\n                            \"required\": True,\n                            \"schema\": {\"title\": id_param_title, \"type\": \"string\"},\n                            \"name\": id_param_name,\n                            \"in\": \"path\",\n                        }\n                    ],\n                    \"responses\": {\n                        \"200\": {\n                            \"description\": \"Successful Response\",\n                            \"content\": {\"application/json\": {\"schema\": {}}},\n                        },\n                        \"422\": {\n                            \"description\": \"Validation Error\",\n                            \"content\": {\n                                \"application/json\": {\n                                    \"schema\": {\n                                        \"$ref\": \"#/components/schemas/HTTPValidationError\"\n                                    }\n                                }\n                            },\n                        },\n                    },\n                }\n            },\n        },\n        \"components\": {\n            \"schemas\": {\n                \"HTTPValidationError\": {\n                    \"title\": \"HTTPValidationError\",\n                    \"type\": \"object\",\n                    \"properties\": {\n                        \"detail\": {\n                            \"title\": \"Detail\",\n                            \"type\": \"array\",\n                            \"items\": {\"$ref\": \"#/components/schemas/ValidationError\"},\n                        }\n                    },\n                },\n                request_schema: {\n                    \"title\": request_schema,\n                    \"type\": \"object\",\n                    \"properties\": {\n                        \"id\": {\"title\": \"Id\", \"type\": \"string\"},\n                        \"input\": {\"$ref\": \"#/components/schemas/Input\"},\n                    },\n                },\n                response_schema: {\n                    \"title\": response_schema,\n                    \"type\": \"object\",\n                    \"properties\": {\n                        \"input\": {\"$ref\": \"#/components/schemas/Input\"},\n                        \"output\": {\"$ref\": \"#/components/schemas/Output\"},\n                        \"id\": {\"title\": \"Id\", \"type\": \"string\"},\n                        \"version\": {\"title\": \"Version\", \"type\": \"string\"},\n                        \"created_at\": {\n                            \"title\": \"Created At\",\n                            \"type\": \"string\",\n                            \"format\": \"date-time\",\n                        },\n                        \"started_at\": {\n                            \"title\": \"Started At\",\n                            \"type\": \"string\",\n                            \"format\": \"date-time\",\n                        },\n                        \"completed_at\": {\n                            \"title\": \"Completed At\",\n                            \"type\": \"string\",\n                            \"format\": \"date-time\",\n                        },\n                        \"status\": {\"title\": \"Status\", \"type\": \"string\"},\n                        \"error\": {\"title\": \"Error\", \"type\": \"string\"},\n                        \"logs\": {\"title\": \"Logs\", \"type\": \"string\"},\n                        \"metrics\": {\"title\": \"Metrics\", \"type\": \"object\"},\n                    },\n                },\n                \"Status\": {\n                    \"title\": \"Status\",\n                    \"description\": \"An enumeration.\",\n                    \"enum\": [\n                        \"starting\",\n                        \"processing\",\n                        \"succeeded\",\n                        \"canceled\",\n                        \"failed\",\n                    ],\n                    \"type\": \"string\",\n                },\n                \"ValidationError\": {\n                    \"title\": \"ValidationError\",\n                    \"required\": [\"loc\", \"msg\", \"type\"],\n                    \"type\": \"object\",\n                    \"properties\": {\n                        \"loc\": {\n                            \"title\": \"Location\",\n                            \"type\": \"array\",\n                            \"items\": {\n                                \"anyOf\": [{\"type\": \"string\"}, {\"type\": \"integer\"}]\n                            },\n                        },\n                        \"msg\": {\"title\": \"Message\", \"type\": \"string\"},\n                        \"type\": {\"title\": \"Error Type\", \"type\": \"string\"},\n                    },\n                },\n            }\n        },\n    }\n\n    # Add Input, Output, and enum schemas\n    schema[\"components\"][\"schemas\"][\"Input\"] = to_json_input(predictor)\n    schema[\"components\"][\"schemas\"][\"Output\"] = to_json_output(predictor)\n    schema[\"components\"][\"schemas\"].update(to_json_enums(predictor))\n\n    return schema\n"
  },
  {
    "path": "python/cog/coder.py",
    "content": "\"\"\"\nCog SDK Coder system for custom type encoding/decoding.\n\nThis module provides the Coder base class for defining custom type\nserialization between Python types and JSON.\n\"\"\"\n\nfrom abc import abstractmethod\nfrom typing import Any, Dict, Optional, Set, Type\n\n\nclass Coder:\n    \"\"\"\n    Base class for custom type encoders/decoders.\n\n    Implement this to add support for custom types in predictor inputs/outputs.\n    Register your coder with Coder.register() to make it available.\n\n    Example:\n        from cog import Coder\n        from myapp import MyCustomType\n\n        class MyCustomCoder(Coder):\n            @staticmethod\n            def factory(tpe: Type) -> Optional[\"MyCustomCoder\"]:\n                if tpe is MyCustomType:\n                    return MyCustomCoder()\n                return None\n\n            def encode(self, value: MyCustomType) -> dict:\n                return {\"data\": value.to_dict()}\n\n            def decode(self, value: dict) -> MyCustomType:\n                return MyCustomType.from_dict(value[\"data\"])\n\n        # Register the coder\n        Coder.register(MyCustomCoder)\n    \"\"\"\n\n    _coders: Set[Type[\"Coder\"]] = set()\n\n    @staticmethod\n    def register(coder: Type[\"Coder\"]) -> None:\n        \"\"\"\n        Register a coder class for custom type handling.\n\n        Args:\n            coder: A Coder subclass to register.\n        \"\"\"\n        Coder._coders.add(coder)\n\n    @staticmethod\n    def lookup(tpe: type | Any) -> Optional[\"Coder\"]:\n        \"\"\"\n        Find a coder that can handle the given type.\n\n        Args:\n            tpe: The type to find a coder for.\n\n        Returns:\n            A Coder instance if one is found, None otherwise.\n        \"\"\"\n        for cls in Coder._coders:\n            c = cls.factory(tpe)\n            if c is not None:\n                return c\n        return None\n\n    @staticmethod\n    @abstractmethod\n    def factory(tpe: Type[Any]) -> Optional[\"Coder\"]:\n        \"\"\"\n        Factory method to create a coder for a given type.\n\n        Override this to check if your coder can handle the type and\n        return an instance if so.\n\n        Args:\n            tpe: The type to potentially handle.\n\n        Returns:\n            A Coder instance if this coder can handle the type, None otherwise.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def encode(self, x: Any) -> Dict[str, Any]:\n        \"\"\"\n        Encode a value to a JSON-serializable dictionary.\n\n        Args:\n            x: The value to encode.\n\n        Returns:\n            A dictionary representation of the value.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def decode(self, x: Dict[str, Any]) -> Any:\n        \"\"\"\n        Decode a dictionary back to the original type.\n\n        Args:\n            x: The dictionary to decode.\n\n        Returns:\n            The decoded value.\n        \"\"\"\n        pass\n"
  },
  {
    "path": "python/cog/command/__init__.py",
    "content": "\"\"\"Cog CLI command modules.\"\"\"\n"
  },
  {
    "path": "python/cog/command/openapi_schema.py",
    "content": "\"\"\"\npython -m cog.command.openapi_schema\n\nThis prints a JSON object describing the OpenAPI schema of the model.\nSchema is generated by introspecting the predictor's type annotations\nwithout starting the HTTP server.\n\"\"\"\n\nimport importlib.util\nimport json\nimport os\nimport sys\nfrom typing import Any, Dict\n\nfrom .._inspector import create_predictor\nfrom .._schemas import to_json_schema\nfrom ..config import Config\nfrom ..errors import ConfigDoesNotExist\nfrom ..mode import Mode\nfrom ..suppress_output import suppress_output\n\n\ndef _load_module_from_ref(ref: str) -> tuple[str, str]:\n    \"\"\"Load a predictor module from a ref like 'predict.py:Predictor' or 'my-subdir/predict.py:Predictor'.\n\n    Uses spec_from_file_location to load the module by file path, which handles\n    subdirectory predictors correctly (unlike import_module which requires the\n    module to be on sys.path).\n\n    Returns (module_name, class_name) with the module pre-loaded in sys.modules.\n    \"\"\"\n    module_path, class_name = ref.rsplit(\":\", 1) if \":\" in ref else (ref, \"Predictor\")\n    module_name = os.path.basename(module_path).removesuffix(\".py\")\n\n    # Load module from file path so subdirectory predictors work\n    spec = importlib.util.spec_from_file_location(module_name, module_path)\n    if spec is not None and spec.loader is not None:\n        module = importlib.util.module_from_spec(spec)\n        sys.modules[module_name] = module\n        spec.loader.exec_module(module)\n\n    return module_name, class_name\n\n\ndef remove_title_next_to_ref(\n    schema_node: Any,\n) -> Any:\n    \"\"\"\n    Recursively remove 'title' from schema components that have a '$ref'.\n    This function addresses a non-compliance issue in FastAPI's OpenAPI schema generation.\n    \"\"\"\n    if isinstance(schema_node, dict):\n        if \"$ref\" in schema_node and \"title\" in schema_node:\n            del schema_node[\"title\"]\n        for _key, value in schema_node.items():\n            remove_title_next_to_ref(value)\n    elif isinstance(schema_node, list):\n        for i, item in enumerate(schema_node):\n            schema_node[i] = remove_title_next_to_ref(item)\n    return schema_node\n\n\ndef fix_nullable_anyof(schema_node: Any) -> None:\n    \"\"\"\n    Convert anyOf with null type to nullable: true for OpenAPI 3.0 compatibility.\n\n    FastAPI generates: {\"anyOf\": [{\"type\": \"string\"}, {\"type\": \"null\"}]}\n    OpenAPI 3.0 wants: {\"type\": \"string\", \"nullable\": true}\n    \"\"\"\n    if isinstance(schema_node, dict):\n        if \"anyOf\" in schema_node:\n            anyof = schema_node[\"anyOf\"]\n            if isinstance(anyof, list) and len(anyof) == 2:\n                # Check if one is {\"type\": \"null\"}\n                null_idx = None\n                other_idx = None\n                for i, item in enumerate(anyof):\n                    if isinstance(item, dict) and item.get(\"type\") == \"null\":\n                        null_idx = i\n                    else:\n                        other_idx = i\n\n                if null_idx is not None and other_idx is not None:\n                    other = anyof[other_idx]\n                    if isinstance(other, dict):\n                        # Replace anyOf with the non-null type + nullable\n                        del schema_node[\"anyOf\"]\n                        schema_node.update(other)\n                        schema_node[\"nullable\"] = True\n\n        for value in schema_node.values():\n            fix_nullable_anyof(value)\n    elif isinstance(schema_node, list):\n        for item in schema_node:\n            fix_nullable_anyof(item)\n\n\nif __name__ == \"__main__\":\n    schema: Dict[str, Any] = {}\n    try:\n        config = Config()\n        # Determine mode: prefer predict, fall back to train\n        try:\n            ref = config.get_predictor_ref(Mode.PREDICT)\n            mode = Mode.PREDICT\n        except ValueError:\n            ref = config.get_predictor_ref(Mode.TRAIN)\n            mode = Mode.TRAIN\n\n        module_name, class_name = _load_module_from_ref(ref)\n\n        with suppress_output():\n            predictor_info = create_predictor(module_name, class_name)\n\n        schema = to_json_schema(predictor_info, mode)\n        remove_title_next_to_ref(schema)\n        fix_nullable_anyof(schema)\n    except FileNotFoundError:\n        raise ConfigDoesNotExist(\"cog.yaml not found\") from None\n\n    print(json.dumps(schema, indent=2))\n"
  },
  {
    "path": "python/cog/config.py",
    "content": "\"\"\"\nConfiguration from cog.yaml.\n\nThis module is restored for the legacy runtime schema generation path\n(python -m cog.command.openapi_schema). It reads cog.yaml to determine\nthe predictor reference.\n\"\"\"\n\nimport os\nfrom typing import Any, Optional\n\nimport yaml\n\nfrom .errors import ConfigDoesNotExist\nfrom .mode import Mode\n\nCOG_YAML_FILE = \"cog.yaml\"\nCOG_PREDICT_TYPE_STUB_ENV_VAR = \"COG_PREDICT_TYPE_STUB\"\nCOG_TRAIN_TYPE_STUB_ENV_VAR = \"COG_TRAIN_TYPE_STUB\"\n\n\nclass Config:\n    \"\"\"A class for reading the cog.yaml properties.\"\"\"\n\n    def __init__(self, config: Optional[dict[str, Any]] = None) -> None:\n        self._config = config\n\n    @property\n    def _cog_config(self) -> dict[str, Any]:\n        config = self._config\n        if config is None:\n            config_path = os.path.abspath(COG_YAML_FILE)\n            try:\n                with open(config_path, encoding=\"utf-8\") as handle:\n                    config = yaml.safe_load(handle)\n            except FileNotFoundError as e:\n                raise ConfigDoesNotExist(\n                    f\"Could not find {config_path}\",\n                ) from e\n            self._config = config\n        return config\n\n    @property\n    def predictor_predict_ref(self) -> Optional[str]:\n        env_val = os.environ.get(COG_PREDICT_TYPE_STUB_ENV_VAR)\n        if env_val:\n            return env_val\n        return self._cog_config.get(str(Mode.PREDICT))\n\n    @property\n    def predictor_train_ref(self) -> Optional[str]:\n        env_val = os.environ.get(COG_TRAIN_TYPE_STUB_ENV_VAR)\n        if env_val:\n            return env_val\n        return self._cog_config.get(str(Mode.TRAIN))\n\n    def get_predictor_ref(self, mode: Mode) -> str:\n        predictor_ref = None\n        if mode == Mode.PREDICT:\n            predictor_ref = self.predictor_predict_ref\n        elif mode == Mode.TRAIN:\n            predictor_ref = self.predictor_train_ref\n        if predictor_ref is None:\n            raise ValueError(\n                f\"Can't run predictions: '{mode}' option not found in cog.yaml\"\n            )\n        return predictor_ref\n"
  },
  {
    "path": "python/cog/errors.py",
    "content": "class CogError(Exception):\n    \"\"\"Base class for all Cog errors.\"\"\"\n\n\nclass ConfigDoesNotExist(CogError):\n    \"\"\"Exception raised when a cog.yaml does not exist.\"\"\"\n\n\nclass PredictorNotSet(CogError):\n    \"\"\"Exception raised when 'predict' is not set in cog.yaml when it needs to be.\"\"\"\n"
  },
  {
    "path": "python/cog/input.py",
    "content": "\"\"\"\nCog SDK Input definition.\n\nThis module provides the Input() function and FieldInfo class for defining\npredictor input parameters with constraints and metadata.\n\"\"\"\n\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, List, Optional, Union\n\n\n@dataclass(frozen=True)\nclass FieldInfo:\n    \"\"\"\n    Internal dataclass to hold Input metadata.\n\n    This stores the constraints and metadata for a predictor input parameter.\n    Users don't typically create this directly - use Input() instead.\n    \"\"\"\n\n    default: Any = None\n    description: Optional[str] = None\n    ge: Optional[Union[int, float]] = None\n    le: Optional[Union[int, float]] = None\n    min_length: Optional[int] = None\n    max_length: Optional[int] = None\n    regex: Optional[str] = None\n    choices: Optional[List[Union[str, int]]] = None\n    deprecated: Optional[bool] = None\n\n\ndef Input(\n    default: Any = None,\n    *,\n    default_factory: Optional[Callable[..., Any]] = None,\n    description: Optional[str] = None,\n    ge: Optional[Union[int, float]] = None,\n    le: Optional[Union[int, float]] = None,\n    min_length: Optional[int] = None,\n    max_length: Optional[int] = None,\n    regex: Optional[str] = None,\n    choices: Optional[List[Union[str, int]]] = None,\n    deprecated: Optional[bool] = None,\n) -> Any:\n    \"\"\"\n    Create an input field specification for a predictor parameter.\n\n    Use this to add metadata and constraints to predictor inputs.\n\n    Example::\n\n        from cog import BasePredictor, Input\n\n        class Predictor(BasePredictor):\n            def predict(\n                self,\n                prompt: str = Input(description=\"The input prompt\"),\n                temperature: float = Input(default=0.7, ge=0.0, le=2.0),\n                max_tokens: int = Input(default=100, ge=1, le=4096),\n            ) -> str:\n                ...\n\n    Args:\n        default: Default value for the field. Must be an immutable literal value.\n        description: Human-readable description of the input.\n        ge: Minimum value (greater than or equal) for numeric inputs.\n        le: Maximum value (less than or equal) for numeric inputs.\n        min_length: Minimum length for string inputs.\n        max_length: Maximum length for string inputs.\n        regex: Regular expression pattern for string inputs.\n        choices: List of allowed values.\n        deprecated: Whether the input is deprecated.\n\n    Returns:\n        A FieldInfo instance containing the field metadata.\n    \"\"\"\n    if default_factory is not None:\n        raise TypeError(\n            \"default_factory is not supported in Input(). \"\n            \"Use a literal default value instead: Input(default=...). \"\n            \"Mutable defaults like lists should use immutable alternatives \"\n            \"(e.g. a comma-separated string) or be constructed in predict().\"\n        )\n\n    return FieldInfo(\n        default=default,\n        description=description,\n        ge=ge,\n        le=le,\n        min_length=min_length,\n        max_length=max_length,\n        regex=regex,\n        choices=choices,\n        deprecated=deprecated,\n    )\n"
  },
  {
    "path": "python/cog/mode.py",
    "content": "from enum import Enum\n\n\nclass Mode(Enum):\n    \"\"\"Enumeration over the different prediction modes.\"\"\"\n\n    PREDICT = \"predict\"\n    TRAIN = \"train\"\n\n    def __str__(self) -> str:\n        return str(self.value)\n"
  },
  {
    "path": "python/cog/model.py",
    "content": "\"\"\"\nCog SDK BaseModel definition.\n\nThis module provides the BaseModel class that users can subclass to define\nstructured output types. BaseModel automatically converts subclasses into\ndataclasses.\n\"\"\"\n\nfrom dataclasses import dataclass, is_dataclass\n\n\nclass BaseModel:\n    \"\"\"\n    Base class for structured output types.\n\n    Subclasses are automatically converted to dataclasses. This provides\n    a clean way to define output schemas without explicit dataclass decorators.\n\n    Example:\n        from cog import BaseModel\n\n        class Output(BaseModel):\n            text: str\n            confidence: float\n\n        # Use as return type\n        def predict(self, prompt: str) -> Output:\n            return Output(text=\"hello\", confidence=0.9)\n\n    By default, auto_dataclass=True, which means all subclasses are\n    automatically wrapped with @dataclass. You can disable this with\n    auto_dataclass=False if you need manual control:\n\n        class ManualModel(BaseModel, auto_dataclass=False):\n            # You must apply @dataclass yourself or handle initialization\n            pass\n    \"\"\"\n\n    def __init_subclass__(\n        cls,\n        *,\n        auto_dataclass: bool = True,\n        init: bool = True,\n        **kwargs: object,\n    ) -> None:\n        \"\"\"\n        Hook called when BaseModel is subclassed.\n\n        This automatically wraps subclasses with @dataclass unless\n        auto_dataclass=False is specified.\n\n        Args:\n            auto_dataclass: If True, automatically apply @dataclass to the class.\n            init: If True (and auto_dataclass=True), generate __init__.\n            **kwargs: Additional arguments passed to @dataclass.\n        \"\"\"\n        # BaseModel is parented to `object` so we have nothing to pass up to it,\n        # we pass the kwargs to dataclass() only.\n        super().__init_subclass__()\n\n        # For sanity, the primary base class must inherit from BaseModel\n        if not issubclass(cls.__bases__[0], BaseModel):\n            raise TypeError(\n                f'Primary base class of \"{cls.__name__}\" must inherit from BaseModel'\n            )\n        elif not auto_dataclass:\n            try:\n                if (\n                    cls.__bases__[0] != BaseModel\n                    and cls.__bases__[0].__auto_dataclass is True  # type: ignore[attr-defined]\n                ):\n                    raise ValueError(\n                        f'Primary base class of \"{cls.__name__}\" '\n                        f'(\"{cls.__bases__[0].__name__}\") has auto_dataclass=True, '\n                        f'but \"{cls.__name__}\" has auto_dataclass=False. '\n                        \"This creates broken field inheritance.\"\n                    )\n            except AttributeError:\n                raise RuntimeError(\n                    f'Primary base class of \"{cls.__name__}\" is a child of a child '\n                    \"of `BaseModel`, but `auto_dataclass` tracking does not exist. \"\n                    \"This is likely a bug or other programming error.\"\n                ) from None\n\n        for base in cls.__bases__[1:]:\n            if is_dataclass(base):\n                raise TypeError(\n                    f'Cannot mixin dataclass \"{base.__name__}\" while inheriting '\n                    \"from `BaseModel`\"\n                )\n\n        # Once manual dataclass handling is enabled, we never apply the auto\n        # dataclass logic again. It becomes the responsibility of the user to\n        # ensure that all dataclass semantics are handled.\n        if not auto_dataclass:\n            cls.__auto_dataclass = False  # type: ignore[attr-defined]\n            return\n\n        # All children should be dataclass'd. This is the only way to ensure\n        # that the dataclass inheritance is handled properly.\n        dataclass(init=init, **kwargs)(cls)  # type: ignore[call-overload]\n        cls.__auto_dataclass = True  # type: ignore[attr-defined]\n"
  },
  {
    "path": "python/cog/predictor.py",
    "content": "\"\"\"\nCog SDK BasePredictor definition.\n\nThis module provides the BasePredictor class that users subclass to define\ntheir model's prediction interface.\n\"\"\"\n\nimport importlib\nimport importlib.util\nimport inspect\nimport os\nimport sys\nfrom typing import Any, Optional, Union\n\nfrom .types import Path\n\n\nclass BasePredictor:\n    \"\"\"\n    Base class for Cog predictors.\n\n    Subclass this to define your model's prediction interface. Override\n    the `setup` method to load your model, and the `predict` method to\n    run predictions.\n\n    Example:\n        from cog import BasePredictor, Input, Path\n\n        class Predictor(BasePredictor):\n            def setup(self):\n                self.model = load_model()\n\n            def predict(self, prompt: str = Input(description=\"Input text\")) -> str:\n                self.record_metric(\"temperature\", 0.7)\n                return self.model.generate(prompt)\n    \"\"\"\n\n    def setup(\n        self,\n        weights: Optional[Union[Path, str]] = None,\n    ) -> None:\n        \"\"\"\n        Prepare the model for predictions.\n\n        This method is called once when the predictor is initialized. Use it\n        to load model weights and do any other one-time setup.\n\n        Args:\n            weights: Optional path to model weights. Can be a local path or URL.\n        \"\"\"\n        pass\n\n    def predict(self, **kwargs: Any) -> Any:\n        \"\"\"\n        Run a single prediction.\n\n        Override this method to implement your model's prediction logic.\n        Input parameters should be annotated with types and optionally\n        use Input() for additional metadata.\n\n        Args:\n            **kwargs: Prediction inputs as defined by the method signature.\n\n        Returns:\n            The prediction output.\n\n        Raises:\n            NotImplementedError: If predict is not implemented.\n        \"\"\"\n        raise NotImplementedError(\"predict has not been implemented by parent class.\")\n\n    @property\n    def scope(self) -> Any:\n        \"\"\"The current prediction scope.\n\n        Provides access to the full scope API for advanced metric operations\n        like dict-style access and deletion::\n\n            self.scope.metrics[\"token_count\"] = 42\n            del self.scope.metrics[\"token_count\"]\n\n        Outside an active prediction this returns a no-op scope.\n        \"\"\"\n        import coglet\n\n        return coglet._sdk.current_scope()  # type: ignore[attr-defined]\n\n    def record_metric(self, key: str, value: Any, mode: str = \"replace\") -> None:\n        \"\"\"Record a prediction metric.\n\n        Convenience method for recording metrics on the current prediction\n        scope. Outside an active prediction this is a silent no-op.\n\n        Args:\n            key: Metric name. Use dot-separated keys (e.g. ``\"timing.inference\"``)\n                to create nested objects in the metrics output.\n            value: Metric value. Supported types: bool, int, float, str, list, dict.\n                Setting a value to ``None`` deletes the metric.\n            mode: Accumulation mode. One of:\n                - ``\"replace\"`` (default): overwrite any previous value.\n                - ``\"incr\"``: add to the existing numeric value.\n                - ``\"append\"``: append to an array.\n\n        Example::\n\n            class Predictor(BasePredictor):\n                def predict(self, prompt: str) -> str:\n                    self.record_metric(\"temperature\", 0.7)\n                    self.record_metric(\"token_count\", 1, mode=\"incr\")\n                    return self.model.generate(prompt)\n        \"\"\"\n        self.scope.record_metric(key, value, mode=mode)\n\n\ndef load_predictor_from_ref(ref: str) -> BasePredictor:\n    \"\"\"Load a predictor from a module:class reference (e.g. 'predict.py:Predictor').\"\"\"\n    module_path, class_name = ref.split(\":\", 1) if \":\" in ref else (ref, \"Predictor\")\n    module_name = os.path.basename(module_path).replace(\".py\", \"\")\n\n    # Use spec_from_file_location to load from file path\n    spec = importlib.util.spec_from_file_location(module_name, module_path)\n    if spec is None or spec.loader is None:\n        raise ImportError(f\"Cannot load module from {module_path}\")\n    module = importlib.util.module_from_spec(spec)\n    # Add module to sys.modules so pickle can find it\n    sys.modules[module_name] = module\n    spec.loader.exec_module(module)\n\n    predictor = getattr(module, class_name)\n    # It could be a class or a function (for training)\n    if inspect.isclass(predictor):\n        return predictor()\n    return predictor\n\n\ndef has_setup_weights(predictor: BasePredictor) -> bool:\n    \"\"\"Check if predictor's setup accepts a weights parameter.\"\"\"\n    if not hasattr(predictor, \"setup\"):\n        return False\n    sig = inspect.signature(predictor.setup)\n    return \"weights\" in sig.parameters\n\n\ndef extract_setup_weights(predictor: BasePredictor) -> Optional[Union[Path, str]]:\n    \"\"\"Extract weights from environment for setup.\"\"\"\n    weights = os.environ.get(\"COG_WEIGHTS\")\n    if weights:\n        return weights\n    return None\n"
  },
  {
    "path": "python/cog/server/__init__.py",
    "content": ""
  },
  {
    "path": "python/cog/server/http.py",
    "content": "import argparse\nimport os\nimport sys\nfrom enum import Enum\n\nimport coglet\n\n\nclass Mode(Enum):\n    PREDICT = \"predict\"\n    TRAIN = \"train\"\n\n    def __str__(self) -> str:\n        return str(self.value)\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description=\"Cog HTTP server\")\n    parser.add_argument(\n        \"-v\", \"--version\", action=\"store_true\", help=\"Show version and exit\"\n    )\n    parser.add_argument(\n        \"--host\",\n        dest=\"host\",\n        type=str,\n        default=\"0.0.0.0\",\n        help=\"Host to bind to\",\n    )\n    parser.add_argument(\n        \"--await-explicit-shutdown\",\n        dest=\"await_explicit_shutdown\",\n        type=bool,\n        default=False,\n        help=\"Ignore SIGTERM and wait for a request to /shutdown (or a SIGINT) before exiting\",\n    )\n    parser.add_argument(\n        \"--x-mode\",\n        dest=\"mode\",\n        type=Mode,\n        default=Mode.PREDICT,\n        choices=list(Mode),\n        help=\"Experimental: Run in 'predict' or 'train' mode\",\n    )\n    # Accept but ignore other args for compatibility\n    parser.add_argument(\"--threads\", dest=\"threads\", type=int, default=None)\n    parser.add_argument(\"--upload-url\", dest=\"upload_url\", type=str, default=None)\n    args = parser.parse_args()\n\n    if args.version:\n        print(f\"coglet (Rust) {coglet.__version__}\")  # type: ignore[attr-defined]\n        sys.exit(0)\n\n    port = int(os.getenv(\"PORT\", \"5000\"))\n    is_train = args.mode == Mode.TRAIN\n\n    # Resolve predictor ref from env vars (set by Dockerfile at build time)\n    if is_train:\n        predictor_ref = os.environ.get(\"COG_TRAIN_TYPE_STUB\")\n    else:\n        predictor_ref = os.environ.get(\"COG_PREDICT_TYPE_STUB\")\n\n    if not predictor_ref:\n        env_var = \"COG_TRAIN_TYPE_STUB\" if is_train else \"COG_PREDICT_TYPE_STUB\"\n        print(\n            f\"ERROR: {env_var} environment variable is not set.\\n\"\n            f\"This should be set automatically by 'cog build'. If running manually,\\n\"\n            f\"set it to your predictor reference (e.g. {env_var}=predict.py:Predictor).\",\n            file=sys.stderr,\n        )\n        sys.exit(1)\n\n    coglet.server.serve(  # type: ignore[attr-defined]\n        predictor_ref=predictor_ref,\n        host=args.host,\n        port=port,\n        await_explicit_shutdown=args.await_explicit_shutdown,\n        is_train=is_train,\n        upload_url=args.upload_url,\n    )\n    sys.exit(0)\n"
  },
  {
    "path": "python/cog/suppress_output.py",
    "content": "import os\nimport sys\nfrom contextlib import contextmanager\nfrom typing import Iterator\n\n\n@contextmanager\ndef suppress_output() -> Iterator[None]:\n    out_fd = sys.stdout.fileno()\n    err_fd = sys.stderr.fileno()\n    out_dup_fd = os.dup(out_fd)\n    err_dup_fd = os.dup(err_fd)\n\n    try:\n        with (\n            open(os.devnull, \"w\", encoding=\"utf-8\") as null_out,\n            open(os.devnull, \"w\", encoding=\"utf-8\") as null_err,\n        ):\n            os.dup2(null_out.fileno(), out_fd)\n            os.dup2(null_err.fileno(), err_fd)\n            try:\n                yield\n            finally:\n                os.dup2(out_dup_fd, out_fd)\n                os.dup2(err_dup_fd, err_fd)\n    finally:\n        os.close(out_dup_fd)\n        os.close(err_dup_fd)\n"
  },
  {
    "path": "python/cog/types.py",
    "content": "\"\"\"\nCog SDK type definitions.\n\nThis module provides core types for defining predictor inputs and outputs:\n- Path: File path type that supports URL inputs\n- Secret: Secure string type that masks its value\n- File: Deprecated file type (use Path instead)\n- ConcatenateIterator: Streaming output iterator\n- AsyncConcatenateIterator: Async streaming output iterator\n\"\"\"\n\nimport io\nimport mimetypes\nimport os\nimport pathlib\nimport shutil\nimport tempfile\nimport urllib.parse\nimport urllib.request\nfrom abc import abstractmethod\nfrom dataclasses import dataclass\nfrom typing import (\n    Any,\n    AsyncIterator,\n    Dict,\n    Iterator,\n    Optional,\n    TypeVar,\n)\n\nimport requests\n\n# Constants for filename handling\nFILENAME_ILLEGAL_CHARS = set(\"\\u0000/\")\nFILENAME_MAX_LENGTH = 200\n\n\ndef _len_bytes(s: str) -> int:\n    \"\"\"Return the length of a string in bytes (UTF-8).\"\"\"\n    return len(s.encode(\"utf-8\"))\n\n\ndef _truncate_filename_bytes(filename: str, length: int) -> str:\n    \"\"\"Truncate a filename to a maximum byte length, preserving extension.\"\"\"\n    if _len_bytes(filename) <= length:\n        return filename\n\n    # Split filename and extension\n    name, ext = os.path.splitext(filename)\n\n    # Reserve space for tilde and extension\n    max_name_length = length - _len_bytes(ext) - 1\n\n    # Truncate name\n    encoded = name.encode(\"utf-8\")\n    truncated = encoded[:max_name_length].decode(\"utf-8\", errors=\"ignore\")\n\n    return f\"{truncated}~{ext}\"\n\n\ndef get_filename(url: str) -> str:\n    \"\"\"Extract a filename from a URL.\"\"\"\n    parsed_url = urllib.parse.urlparse(url)\n\n    if parsed_url.scheme == \"data\":\n        # Safe: scheme is validated to be 'data:' before urlopen\n        with urllib.request.urlopen(url) as resp:  # noqa: S310\n            mime_type = resp.headers.get_content_type()\n            extension = mimetypes.guess_extension(mime_type)\n            if extension is None:\n                return \"file\"\n            return \"file\" + extension\n\n    basename = os.path.basename(parsed_url.path)\n    basename = urllib.parse.unquote_plus(basename)\n\n    # Truncate if too long\n    if _len_bytes(basename) > FILENAME_MAX_LENGTH:\n        basename = _truncate_filename_bytes(basename, length=FILENAME_MAX_LENGTH)\n\n    # Replace illegal characters\n    for c in FILENAME_ILLEGAL_CHARS:\n        basename = basename.replace(c, \"_\")\n\n    return basename\n\n\n########################################\n# Secret\n########################################\n\n\n@dataclass(frozen=True)\nclass Secret:\n    \"\"\"\n    A secret string value that masks itself in string representations.\n\n    Use this type for sensitive data like API keys or passwords that should\n    not be logged or displayed.\n\n    Example:\n        def predict(self, api_key: Secret) -> str:\n            key = api_key.get_secret_value()\n            # Use key...\n    \"\"\"\n\n    secret_value: Optional[str] = None\n\n    def __repr__(self) -> str:\n        return f\"Secret({str(self)})\"\n\n    def __str__(self) -> str:\n        return \"**********\" if self.secret_value is not None else \"\"\n\n    def get_secret_value(self) -> Optional[str]:\n        \"\"\"Return the actual secret value.\"\"\"\n        return self.secret_value\n\n\n########################################\n# URLFile\n########################################\n\n\nclass URLFile(io.IOBase):\n    \"\"\"\n    URLFile is a proxy object for a :class:`urllib3.response.HTTPResponse`\n    object that is created lazily. It's a file-like object constructed from a\n    URL that can survive pickling/unpickling.\n    \"\"\"\n\n    __slots__ = (\"__target__\", \"__url__\", \"name\")\n\n    def __init__(self, url: str, filename: Optional[str] = None) -> None:\n        parsed = urllib.parse.urlparse(url)\n        if parsed.scheme not in {\"http\", \"https\"}:\n            raise ValueError(\n                \"URLFile requires URL to conform to HTTP or HTTPS protocol\"\n            )\n\n        if not filename:\n            filename = os.path.basename(parsed.path)\n\n        object.__setattr__(self, \"name\", filename)\n        object.__setattr__(self, \"__url__\", url)\n\n    def __del__(self) -> None:\n        try:\n            object.__getattribute__(self, \"__target__\")\n        except AttributeError:\n            # Do nothing when tearing down the object if the response object\n            # hasn't been created yet.\n            return\n\n        super().__del__()\n\n    # We provide __getstate__ and __setstate__ explicitly to ensure that the\n    # object is always picklable.\n    def __getstate__(self) -> Dict[str, Any]:\n        return {\n            \"name\": object.__getattribute__(self, \"name\"),\n            \"url\": object.__getattribute__(self, \"__url__\"),\n        }\n\n    def __setstate__(self, state: Dict[str, Any]) -> None:\n        object.__setattr__(self, \"name\", state[\"name\"])\n        object.__setattr__(self, \"__url__\", state[\"url\"])\n\n    # Proxy getattr/setattr/delattr through to the response object.\n    def __setattr__(self, name: str, value: Any) -> None:\n        if hasattr(type(self), name):\n            object.__setattr__(self, name, value)\n        else:\n            setattr(self.__wrapped__, name, value)\n\n    def __getattr__(self, name: str) -> Any:\n        if name in (\"__target__\", \"__wrapped__\", \"__url__\"):\n            raise AttributeError(name)\n        elif name == \"name\":\n            return object.__getattribute__(self, \"name\")\n        return getattr(self.__wrapped__, name)\n\n    def __delattr__(self, name: str) -> None:\n        if hasattr(type(self), name):\n            object.__delattr__(self, name)\n        else:\n            delattr(self.__wrapped__, name)\n\n    # Luckily the only dunder method on HTTPResponse is __iter__\n    def __iter__(self) -> Iterator[bytes]:\n        return iter(self.__wrapped__)\n\n    @property\n    def __wrapped__(self) -> Any:\n        try:\n            return object.__getattribute__(self, \"__target__\")\n        except AttributeError:\n            pass\n        url = object.__getattribute__(self, \"__url__\")\n        headers = {}\n        ua = os.getenv(\"COG_USER_AGENT\")\n        if ua:\n            headers[\"User-Agent\"] = ua\n\n        resp = requests.get(url, stream=True, timeout=10, headers=headers)\n        resp.raise_for_status()\n        resp.raw.decode_content = True\n        object.__setattr__(self, \"__target__\", resp.raw)\n        return resp.raw\n\n    def __repr__(self) -> str:\n        try:\n            target = object.__getattribute__(self, \"__target__\")\n        except AttributeError:\n            return f\"<{type(self).__name__} at 0x{id(self):x} for {object.__getattribute__(self, '__url__')!r}>\"\n\n        return f\"<{type(self).__name__} at 0x{id(self):x} wrapping {target!r}>\"\n\n\n########################################\n# File (Deprecated)\n########################################\n\n\nclass File(io.IOBase):\n    \"\"\"\n    Deprecated: use Path instead.\n\n    A file-like object that can be constructed from a URL or data URI.\n    \"\"\"\n\n    @classmethod\n    def validate(cls, value: Any) -> io.IOBase:\n        \"\"\"Validate and convert a value to a file-like object.\"\"\"\n        if isinstance(value, io.IOBase):\n            return value\n\n        parsed_url = urllib.parse.urlparse(value)\n        if parsed_url.scheme == \"data\":\n            # Safe: scheme is validated to be 'data:' before urlopen\n            with urllib.request.urlopen(value) as res:  # noqa: S310\n                return io.BytesIO(res.read())\n        if parsed_url.scheme in (\"http\", \"https\"):\n            return URLFile(value)\n        raise ValueError(\n            f\"'{parsed_url.scheme}' is not a valid URL scheme. \"\n            \"'data', 'http', or 'https' is supported.\"\n        )\n\n\n########################################\n# URLPath\n########################################\n\n\nclass URLPath(pathlib.PosixPath):\n    \"\"\"\n    URLPath is a nasty hack to ensure that we can defer the downloading of a\n    URL passed as a path until later in prediction dispatch.\n\n    It subclasses pathlib.PosixPath only so that it can pass isinstance(_,\n    pathlib.Path) checks.\n    \"\"\"\n\n    _path: Optional[\"Path\"]\n\n    # pylint: disable=super-init-not-called\n    def __init__(self, *, source: str, filename: str, fileobj: io.IOBase) -> None:\n        if len(filename) > FILENAME_MAX_LENGTH:\n            filename = _truncate_filename_bytes(filename, FILENAME_MAX_LENGTH)\n\n        self.source = source\n        self.filename = filename\n        self.fileobj = fileobj\n\n        self._path = None\n\n    def __new__(cls, *, source: str, filename: str, fileobj: io.IOBase) -> \"URLPath\":\n        # PosixPath.__new__ requires path segments, but we don't have a real path\n        # Use a placeholder that will be replaced\n        obj = super().__new__(cls, filename)\n        return obj\n\n    def convert(self) -> \"Path\":\n        \"\"\"Download the URL content to a temporary file and return its Path.\"\"\"\n        if self._path is None:\n            # pylint: disable=consider-using-with\n            dest = tempfile.NamedTemporaryFile(suffix=self.filename, delete=False)\n            shutil.copyfileobj(self.fileobj, dest)\n            dest.close()\n            self._path = Path(dest.name)\n        return self._path\n\n    def unlink(self, missing_ok: bool = False) -> None:\n        \"\"\"Remove the temporary file if it exists.\"\"\"\n        if self._path:\n            self._path.unlink(missing_ok=missing_ok)\n\n    def __str__(self) -> str:\n        # FastAPI's jsonable_encoder will encode subclasses of pathlib.Path by\n        # calling str() on them\n        return self.source\n\n\n########################################\n# Path\n########################################\n\n\nclass Path(pathlib.PosixPath):\n    \"\"\"\n    A path type that can be constructed from URLs.\n\n    When a URL is passed, it creates a URLPath that defers downloading\n    until the file is actually needed.\n\n    Example:\n        def predict(self, image: Path) -> Path:\n            # image could be a local path or downloaded from URL\n            return process(image)\n    \"\"\"\n\n    @classmethod\n    def validate(cls, value: Any) -> pathlib.Path:\n        \"\"\"Validate and convert a value to a Path.\"\"\"\n        if isinstance(value, pathlib.Path):\n            return value\n\n        parsed_url = urllib.parse.urlparse(value)\n        if parsed_url.scheme in (\"data\", \"http\", \"https\"):\n            return URLPath(\n                source=value,\n                filename=get_filename(value),\n                fileobj=File.validate(value),\n            )\n\n        return Path(value)\n\n\n########################################\n# Iterators\n########################################\n\nItem = TypeVar(\"Item\")\n\n\nclass ConcatenateIterator(Iterator[Item]):\n    \"\"\"\n    An iterator that yields items which should be concatenated for display.\n\n    Use this as a return type hint for streaming text output where the\n    individual chunks should be joined together.\n\n    Example:\n        def predict(self, prompt: str) -> ConcatenateIterator[str]:\n            for token in generate_tokens(prompt):\n                yield token\n    \"\"\"\n\n    @abstractmethod\n    def __next__(self) -> Item: ...\n\n\nclass AsyncConcatenateIterator(AsyncIterator[Item]):\n    \"\"\"\n    An async iterator that yields items which should be concatenated for display.\n\n    Use this as a return type hint for async streaming text output where the\n    individual chunks should be joined together.\n\n    Example:\n        async def predict(self, prompt: str) -> AsyncConcatenateIterator[str]:\n            async for token in generate_tokens_async(prompt):\n                yield token\n    \"\"\"\n\n    @abstractmethod\n    async def __anext__(self) -> Item: ...\n"
  },
  {
    "path": "python/tests/__init__.py",
    "content": "\"\"\"Tests for cog-dataclass SDK.\"\"\"\n"
  },
  {
    "path": "python/tests/test_emit_metric.py",
    "content": "\"\"\"Tests for the emit_metric backwards-compatibility shim.\"\"\"\n\nimport subprocess\nimport sys\n\n\nclass TestEmitMetric:\n    \"\"\"Tests for the deprecated emit_metric import.\"\"\"\n\n    def test_import_succeeds(self) -> None:\n        \"\"\"Importing emit_metric should not raise.\"\"\"\n        result = subprocess.run(\n            [sys.executable, \"-c\", \"from cog import emit_metric\"],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode == 0, f\"Import failed: {result.stderr}\"\n\n    def test_attribute_access_succeeds(self) -> None:\n        \"\"\"Accessing cog.emit_metric as a module attribute should not raise.\"\"\"\n        result = subprocess.run(\n            [sys.executable, \"-c\", \"import cog; cog.emit_metric\"],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode == 0, f\"Attribute access failed: {result.stderr}\"\n\n    def test_prints_deprecation_to_stderr(self) -> None:\n        \"\"\"First import should print a deprecation message to stderr.\"\"\"\n        result = subprocess.run(\n            [sys.executable, \"-c\", \"from cog import emit_metric\"],\n            capture_output=True,\n            text=True,\n        )\n        assert \"emit_metric() is deprecated\" in result.stderr\n\n    def test_message_prints_once(self) -> None:\n        \"\"\"The deprecation message should print only once per process, not on every call.\"\"\"\n        code = \"from cog import emit_metric\\nfrom cog import emit_metric\\n\"\n        result = subprocess.run(\n            [sys.executable, \"-c\", code],\n            capture_output=True,\n            text=True,\n        )\n        assert result.stderr.count(\"emit_metric() is deprecated\") == 1\n\n    def test_callable(self) -> None:\n        \"\"\"emit_metric should be callable and not raise outside a prediction context.\"\"\"\n        result = subprocess.run(\n            [\n                sys.executable,\n                \"-c\",\n                \"from cog import emit_metric; emit_metric('output_tokens', 42)\",\n            ],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode == 0, f\"Call failed: {result.stderr}\"\n\n    def test_module_attribute_callable(self) -> None:\n        \"\"\"cog.emit_metric(...) style (used in cog-triton, cog-arctic, etc.) should work.\"\"\"\n        result = subprocess.run(\n            [\n                sys.executable,\n                \"-c\",\n                \"import cog; cog.emit_metric('input_token_count', 100)\",\n            ],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode == 0, (\n            f\"Call via module attribute failed: {result.stderr}\"\n        )\n\n    def test_unknown_attr_still_raises(self) -> None:\n        \"\"\"Adding emit_metric shim should not break AttributeError for unknown attrs.\"\"\"\n        result = subprocess.run(\n            [sys.executable, \"-c\", \"import cog; cog.NoSuchAttribute\"],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode != 0\n        assert \"AttributeError\" in result.stderr\n"
  },
  {
    "path": "python/tests/test_experimental_feature_warning.py",
    "content": "\"\"\"Tests for the ExperimentalFeatureWarning backwards-compatibility shim.\"\"\"\n\nimport subprocess\nimport sys\n\n\nclass TestExperimentalFeatureWarning:\n    \"\"\"Tests for the deprecated ExperimentalFeatureWarning import.\"\"\"\n\n    def test_import_succeeds(self) -> None:\n        \"\"\"Importing ExperimentalFeatureWarning should not raise.\"\"\"\n        result = subprocess.run(\n            [sys.executable, \"-c\", \"from cog import ExperimentalFeatureWarning\"],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode == 0, f\"Import failed: {result.stderr}\"\n\n    def test_prints_deprecation_to_stderr(self) -> None:\n        \"\"\"First import should print a deprecation message to stderr.\"\"\"\n        result = subprocess.run(\n            [sys.executable, \"-c\", \"from cog import ExperimentalFeatureWarning\"],\n            capture_output=True,\n            text=True,\n        )\n        assert \"ExperimentalFeatureWarning is deprecated\" in result.stderr\n\n    def test_message_prints_once(self) -> None:\n        \"\"\"The deprecation message should print only once per process.\"\"\"\n        code = (\n            \"from cog import ExperimentalFeatureWarning\\n\"\n            \"from cog import ExperimentalFeatureWarning\\n\"\n        )\n        result = subprocess.run(\n            [sys.executable, \"-c\", code],\n            capture_output=True,\n            text=True,\n        )\n        assert result.stderr.count(\"ExperimentalFeatureWarning is deprecated\") == 1\n\n    def test_is_future_warning_subclass(self) -> None:\n        \"\"\"The shim class should be a subclass of FutureWarning.\"\"\"\n        code = (\n            \"from cog import ExperimentalFeatureWarning\\n\"\n            \"assert issubclass(ExperimentalFeatureWarning, FutureWarning)\\n\"\n        )\n        result = subprocess.run(\n            [sys.executable, \"-c\", code],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode == 0, f\"Assertion failed: {result.stderr}\"\n\n    def test_filterwarnings_compat(self) -> None:\n        \"\"\"The real use case: warnings.filterwarnings('ignore', ...) should work.\"\"\"\n        code = (\n            \"import warnings\\n\"\n            \"from cog import ExperimentalFeatureWarning\\n\"\n            \"warnings.filterwarnings('ignore', category=ExperimentalFeatureWarning)\\n\"\n        )\n        result = subprocess.run(\n            [sys.executable, \"-c\", code],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode == 0, f\"filterwarnings failed: {result.stderr}\"\n\n    def test_unknown_attr_raises(self) -> None:\n        \"\"\"Accessing a non-existent attribute should raise AttributeError.\"\"\"\n        code = \"import cog\\ncog.NoSuchAttribute\\n\"\n        result = subprocess.run(\n            [sys.executable, \"-c\", code],\n            capture_output=True,\n            text=True,\n        )\n        assert result.returncode != 0\n        assert \"AttributeError\" in result.stderr\n"
  },
  {
    "path": "python/tests/test_input.py",
    "content": "\"\"\"Tests for cog.input module (Input, FieldInfo).\"\"\"\n\nimport pytest\n\nfrom cog import Input\nfrom cog.input import FieldInfo\n\n\nclass TestInput:\n    \"\"\"Tests for Input() function.\"\"\"\n\n    def test_input_returns_fieldinfo(self) -> None:\n        result = Input(description=\"Test input\")\n        assert isinstance(result, FieldInfo)\n\n    def test_input_with_default(self) -> None:\n        result = Input(default=\"hello\", description=\"A string\")\n        assert result.default == \"hello\"\n        assert result.description == \"A string\"\n\n    def test_input_with_numeric_constraints(self) -> None:\n        result = Input(default=5, ge=0, le=10)\n        assert result.default == 5\n        assert result.ge == 0\n        assert result.le == 10\n\n    def test_input_with_string_constraints(self) -> None:\n        result = Input(min_length=1, max_length=100, regex=r\"^\\w+$\")\n        assert result.min_length == 1\n        assert result.max_length == 100\n        assert result.regex == r\"^\\w+$\"\n\n    def test_input_with_choices(self) -> None:\n        result = Input(default=\"a\", choices=[\"a\", \"b\", \"c\"])\n        assert result.default == \"a\"\n        assert result.choices == [\"a\", \"b\", \"c\"]\n\n    def test_input_with_deprecated(self) -> None:\n        result = Input(deprecated=True)\n        assert result.deprecated is True\n\n    def test_input_default_factory_raises_error(self) -> None:\n        with pytest.raises(TypeError, match=\"default_factory is not supported\"):\n            Input(default_factory=list)\n\n    def test_input_immutable_defaults_stored_directly(self) -> None:\n        for default in [\"string\", 42, 3.14, True, None, (1, 2), frozenset([1, 2])]:\n            result = Input(default=default)\n            assert result.default == default\n\n    def test_input_no_default(self) -> None:\n        # No default means the parameter is required\n        result = Input(description=\"Required input\")\n        assert result.default is None\n        assert result.description == \"Required input\"\n\n\nclass TestFieldInfo:\n    \"\"\"Tests for FieldInfo dataclass.\"\"\"\n\n    def test_fieldinfo_is_frozen(self) -> None:\n        info = FieldInfo(default=\"test\")\n        with pytest.raises(AttributeError):\n            info.default = \"new\"  # type: ignore[misc]\n\n    def test_fieldinfo_defaults(self) -> None:\n        info = FieldInfo(default=5, ge=0, le=10, description=\"A number\")\n        assert info.default == 5\n        assert info.ge == 0\n        assert info.le == 10\n        assert info.description == \"A number\"\n\n    def test_fieldinfo_none_defaults(self) -> None:\n        info = FieldInfo(description=\"Just a description\")\n        assert info.default is None\n        assert info.ge is None\n        assert info.le is None\n"
  },
  {
    "path": "python/tests/test_model.py",
    "content": "\"\"\"Tests for cog.model module (BaseModel).\"\"\"\n\nfrom dataclasses import is_dataclass\nfrom typing import Optional\n\nfrom cog import BaseModel, Path\n\n\nclass TestBaseModel:\n    \"\"\"Tests for BaseModel auto-dataclass behavior.\"\"\"\n\n    def test_subclass_becomes_dataclass(self) -> None:\n        class Output(BaseModel):\n            text: str\n            score: float\n\n        assert is_dataclass(Output)\n\n    def test_subclass_can_be_instantiated(self) -> None:\n        class Output(BaseModel):\n            text: str\n            score: float\n\n        output = Output(text=\"hello\", score=0.9)\n        assert output.text == \"hello\"\n        assert output.score == 0.9\n\n    def test_subclass_with_defaults(self) -> None:\n        class Output(BaseModel):\n            text: str\n            score: float = 0.5\n\n        output = Output(text=\"hello\")\n        assert output.text == \"hello\"\n        assert output.score == 0.5\n\n    def test_subclass_with_optional(self) -> None:\n        class Output(BaseModel):\n            text: str\n            metadata: Optional[str] = None\n\n        output = Output(text=\"hello\")\n        assert output.text == \"hello\"\n        assert output.metadata is None\n\n    def test_nested_models(self) -> None:\n        class Inner(BaseModel):\n            value: int\n\n        class Outer(BaseModel):\n            inner: Inner\n            name: str\n\n        inner = Inner(value=42)\n        outer = Outer(inner=inner, name=\"test\")\n        assert outer.inner.value == 42\n        assert outer.name == \"test\"\n\n    def test_inheritance(self) -> None:\n        class Base(BaseModel):\n            x: int\n\n        class Derived(Base):\n            y: str\n\n        derived = Derived(x=1, y=\"two\")\n        assert derived.x == 1\n        assert derived.y == \"two\"\n\n    def test_auto_dataclass_false(self) -> None:\n        class Manual(BaseModel, auto_dataclass=False):\n            x: int\n\n            def __init__(self, x: int) -> None:\n                self.x = x\n\n        # Should not be auto-dataclassed\n        assert not is_dataclass(Manual)\n\n        # But should still be usable\n        m = Manual(x=5)\n        assert m.x == 5\n\n    def test_primary_base_must_be_basemodel(self) -> None:\n        class NotBaseModel:\n            pass\n\n        try:\n\n            class Bad(NotBaseModel, BaseModel):  # type: ignore[misc]\n                x: int\n\n            assert False, \"Should have raised TypeError\"\n        except TypeError as e:\n            assert \"must inherit from BaseModel\" in str(e)\n\n    def test_cannot_mixin_dataclass(self) -> None:\n        from dataclasses import dataclass\n\n        @dataclass\n        class SomeDataclass:\n            y: int\n\n        try:\n\n            class Bad(BaseModel, SomeDataclass):  # type: ignore[misc]\n                x: int\n\n            assert False, \"Should have raised TypeError\"\n        except TypeError as e:\n            assert \"Cannot mixin dataclass\" in str(e)\n\n    def test_auto_dataclass_inheritance_mismatch(self) -> None:\n        class Parent(BaseModel):\n            x: int\n\n        try:\n\n            class Child(Parent, auto_dataclass=False):\n                y: str\n\n            assert False, \"Should have raised ValueError\"\n        except ValueError as e:\n            assert \"auto_dataclass=True\" in str(e)\n            assert \"auto_dataclass=False\" in str(e)\n\n    def test_basemodel_asdict(self) -> None:\n        from dataclasses import asdict\n\n        class Output(BaseModel):\n            weights: Path\n\n        output = Output(weights=Path(\"weights.bin\"))\n        assert asdict(output) == {\"weights\": Path(\"weights.bin\")}\n"
  },
  {
    "path": "python/tests/test_predictor.py",
    "content": "\"\"\"Tests for cog.predictor module (BasePredictor).\"\"\"\n\nfrom typing import Optional\n\nfrom cog import BasePredictor, Path\n\n\nclass TestBasePredictor:\n    \"\"\"Tests for BasePredictor class.\"\"\"\n\n    def test_subclass_can_override_predict(self) -> None:\n        class MyPredictor(BasePredictor):\n            def predict(self, text: str) -> str:\n                return text.upper()\n\n        predictor = MyPredictor()\n        result = predictor.predict(text=\"hello\")\n        assert result == \"HELLO\"\n\n    def test_default_predict_raises(self) -> None:\n        predictor = BasePredictor()\n        try:\n            predictor.predict()\n            assert False, \"Should have raised NotImplementedError\"\n        except NotImplementedError as e:\n            assert \"predict has not been implemented\" in str(e)\n\n    def test_setup_is_optional(self) -> None:\n        class MyPredictor(BasePredictor):\n            def predict(self, x: int) -> int:\n                return x * 2\n\n        predictor = MyPredictor()\n        # setup() should not raise\n        predictor.setup()\n        assert predictor.predict(x=5) == 10\n\n    def test_setup_with_weights(self) -> None:\n        class MyPredictor(BasePredictor):\n            weights_path: Optional[str] = None\n\n            def setup(self, weights: Optional[str] = None) -> None:\n                self.weights_path = weights\n\n            def predict(self, x: int) -> int:\n                return x\n\n        predictor = MyPredictor()\n        predictor.setup(weights=\"/path/to/weights\")\n        assert predictor.weights_path == \"/path/to/weights\"\n\n    def test_setup_with_path_weights(self) -> None:\n        class MyPredictor(BasePredictor):\n            weights_path: Optional[Path] = None\n\n            def setup(self, weights: Optional[Path] = None) -> None:\n                self.weights_path = weights\n\n            def predict(self, x: int) -> int:\n                return x\n\n        predictor = MyPredictor()\n        predictor.setup(weights=Path(\"/path/to/weights\"))\n        assert str(predictor.weights_path) == \"/path/to/weights\"\n\n    def test_predictor_with_multiple_inputs(self) -> None:\n        class MyPredictor(BasePredictor):\n            def predict(self, a: int, b: int, c: str = \"default\") -> str:\n                return f\"{a + b}: {c}\"\n\n        predictor = MyPredictor()\n        result = predictor.predict(a=1, b=2, c=\"test\")\n        assert result == \"3: test\"\n\n        result_default = predictor.predict(a=1, b=2)\n        assert result_default == \"3: default\"\n\n    def test_predictor_with_state(self) -> None:\n        class StatefulPredictor(BasePredictor):\n            count: int = 0\n\n            def setup(self, weights: Optional[str] = None) -> None:\n                self.count = 0\n\n            def predict(self, x: int) -> int:\n                self.count += 1\n                return x * self.count\n\n        predictor = StatefulPredictor()\n        predictor.setup()\n        assert predictor.predict(x=10) == 10\n        assert predictor.predict(x=10) == 20\n        assert predictor.predict(x=10) == 30\n"
  },
  {
    "path": "python/tests/test_types.py",
    "content": "\"\"\"Tests for cog.types module.\"\"\"\n\nimport io\nfrom dataclasses import is_dataclass\n\nfrom cog import (\n    AsyncConcatenateIterator,\n    ConcatenateIterator,\n    File,\n    Path,\n    Secret,\n    URLFile,\n)\n\n\nclass TestSecret:\n    \"\"\"Tests for Secret type.\"\"\"\n\n    def test_secret_creation(self) -> None:\n        secret = Secret(secret_value=\"my-api-key\")\n        assert secret.get_secret_value() == \"my-api-key\"\n\n    def test_secret_masks_in_str(self) -> None:\n        secret = Secret(secret_value=\"my-api-key\")\n        assert str(secret) == \"**********\"\n        assert \"my-api-key\" not in str(secret)\n\n    def test_secret_masks_in_repr(self) -> None:\n        secret = Secret(secret_value=\"my-api-key\")\n        assert \"my-api-key\" not in repr(secret)\n        assert \"**********\" in repr(secret)\n\n    def test_secret_none_value(self) -> None:\n        secret = Secret(secret_value=None)\n        assert secret.get_secret_value() is None\n        assert str(secret) == \"\"\n\n    def test_secret_default_none(self) -> None:\n        secret = Secret()\n        assert secret.get_secret_value() is None\n\n    def test_secret_is_dataclass(self) -> None:\n        assert is_dataclass(Secret)\n\n    def test_secret_is_frozen(self) -> None:\n        secret = Secret(secret_value=\"test\")\n        try:\n            secret.secret_value = \"new\"  # type: ignore[misc]\n            assert False, \"Should have raised FrozenInstanceError\"\n        except Exception:\n            pass  # Expected - frozen dataclass\n\n\nclass TestPath:\n    \"\"\"Tests for Path type.\"\"\"\n\n    def test_path_from_string(self) -> None:\n        p = Path(\"/tmp/test.txt\")\n        assert str(p) == \"/tmp/test.txt\"\n\n    def test_path_is_pathlib_subclass(self) -> None:\n        import pathlib\n\n        p = Path(\"/tmp/test.txt\")\n        assert isinstance(p, pathlib.PosixPath)\n\n\nclass TestFile:\n    \"\"\"Tests for File type (deprecated).\"\"\"\n\n    def test_file_validate_iobase(self) -> None:\n        buf = io.BytesIO(b\"test data\")\n        result = File.validate(buf)\n        assert result is buf\n\n    def test_file_validate_data_uri(self) -> None:\n        # data URI with plain text\n        data_uri = \"data:text/plain;base64,SGVsbG8gV29ybGQ=\"\n        result = File.validate(data_uri)\n        assert isinstance(result, io.BytesIO)\n        assert result.read() == b\"Hello World\"\n\n    def test_file_validate_invalid_scheme(self) -> None:\n        try:\n            File.validate(\"ftp://example.com/file.txt\")\n            assert False, \"Should have raised ValueError\"\n        except ValueError as e:\n            assert \"not a valid URL scheme\" in str(e)\n\n\nclass TestURLFile:\n    \"\"\"Tests for URLFile type.\"\"\"\n\n    def test_urlfile_creation(self) -> None:\n        url = \"https://example.com/image.jpg\"\n        uf = URLFile(url)\n        assert uf.name == \"image.jpg\"\n\n    def test_urlfile_invalid_scheme(self) -> None:\n        try:\n            URLFile(\"ftp://example.com/file.txt\")\n            assert False, \"Should have raised ValueError\"\n        except ValueError as e:\n            assert \"HTTP or HTTPS\" in str(e)\n\n    def test_urlfile_picklable(self) -> None:\n        import pickle\n\n        url = \"https://example.com/image.jpg\"\n        uf = URLFile(url)\n        pickled = pickle.dumps(uf)\n        restored = pickle.loads(pickled)\n        assert restored.name == \"image.jpg\"\n\n    def test_urlfile_custom_filename(self) -> None:\n        url = \"https://example.com/image.jpg\"\n        uf = URLFile(url, filename=\"custom.png\")\n        assert uf.name == \"custom.png\"\n\n\nclass TestIterators:\n    \"\"\"Tests for iterator types.\"\"\"\n\n    def test_concatenate_iterator_is_abstract(self) -> None:\n        # ConcatenateIterator should be usable as a type hint\n        from typing import Iterator\n\n        assert issubclass(ConcatenateIterator, Iterator)\n\n    def test_async_concatenate_iterator_is_abstract(self) -> None:\n        # AsyncConcatenateIterator should be usable as a type hint\n        from typing import AsyncIterator\n\n        assert issubclass(AsyncConcatenateIterator, AsyncIterator)\n"
  },
  {
    "path": "script/generate-compat",
    "content": "#!/usr/bin/env bash\n#\n# Regenerate CUDA/PyTorch/TensorFlow compatibility matrices.\n#\n# Usage:\n#   script/generate-compat          # regenerate all matrices\n#   script/generate-compat cuda     # regenerate CUDA base images only\n#   script/generate-compat torch    # regenerate PyTorch compatibility only\n#   script/generate-compat tensorflow  # regenerate TensorFlow compatibility only\n#\n# The generated JSON files are checked into source control and only need\n# to be regenerated when adding support for new framework versions.\n\nset -euo pipefail\n\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\nROOT_DIR=\"$(dirname \"$SCRIPT_DIR\")\"\nCONFIG_DIR=\"${ROOT_DIR}/pkg/config\"\n\ngenerate_cuda() {\n    echo \"Generating CUDA base images...\"\n    go run \"${ROOT_DIR}/tools/compatgen/main.go\" cuda -o \"${CONFIG_DIR}/cuda_base_images.json\"\n}\n\ngenerate_torch() {\n    echo \"Generating PyTorch compatibility matrix...\"\n    go run \"${ROOT_DIR}/tools/compatgen/main.go\" torch -o \"${CONFIG_DIR}/torch_compatibility_matrix.json\"\n}\n\ngenerate_tensorflow() {\n    echo \"Generating TensorFlow compatibility matrix...\"\n    go run \"${ROOT_DIR}/tools/compatgen/main.go\" tensorflow -o \"${CONFIG_DIR}/tf_compatibility_matrix.json\"\n}\n\ntarget=\"${1:-all}\"\n\ncase \"$target\" in\n    cuda)\n        generate_cuda\n        ;;\n    torch)\n        generate_torch\n        ;;\n    tensorflow|tf)\n        generate_tensorflow\n        ;;\n    all)\n        generate_cuda\n        generate_torch\n        generate_tensorflow\n        ;;\n    *)\n        echo \"Unknown target: $target\"\n        echo \"Usage: $0 [cuda|torch|tensorflow|all]\"\n        exit 1\n        ;;\nesac\n\necho \"Done.\"\n"
  },
  {
    "path": "test-helpers/https-server/go.mod",
    "content": "module github.com/replicate/cog/test-helpers/https-server\n\ngo 1.21\n"
  },
  {
    "path": "test-helpers/https-server/main.go",
    "content": "// Package main provides a simple HTTPS server for testing CA certificate injection.\n// Usage: go run ./test-helpers/https-server --cert=server.crt --key=server.key --addr=:8443\npackage main\n\nimport (\n\t\"flag\"\n\t\"fmt\"\n\t\"log\"\n\t\"net/http\"\n\t\"os\"\n\t\"os/signal\"\n\t\"syscall\"\n)\n\nfunc main() {\n\tcert := flag.String(\"cert\", \"\", \"Path to TLS certificate file\")\n\tkey := flag.String(\"key\", \"\", \"Path to TLS key file\")\n\taddr := flag.String(\"addr\", \":8443\", \"Address to listen on\")\n\tflag.Parse()\n\n\tif *cert == \"\" || *key == \"\" {\n\t\tfmt.Fprintln(os.Stderr, \"Usage: https-server --cert=server.crt --key=server.key [--addr=:8443]\")\n\t\tos.Exit(1)\n\t}\n\n\t// Simple handler that returns OK\n\thttp.HandleFunc(\"/\", func(w http.ResponseWriter, r *http.Request) {\n\t\tw.WriteHeader(http.StatusOK)\n\t\tfmt.Fprintln(w, \"OK\")\n\t})\n\n\t// Start server in a goroutine\n\tserver := &http.Server{Addr: *addr}\n\tgo func() {\n\t\tlog.Printf(\"Starting HTTPS server on %s\", *addr)\n\t\tif err := server.ListenAndServeTLS(*cert, *key); err != http.ErrServerClosed {\n\t\t\tlog.Fatalf(\"HTTPS server error: %v\", err)\n\t\t}\n\t}()\n\n\t// Print ready message for test synchronization\n\tfmt.Println(\"READY\")\n\n\t// Wait for interrupt signal\n\tsigChan := make(chan os.Signal, 1)\n\tsignal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)\n\t<-sigChan\n\n\tlog.Println(\"Shutting down HTTPS server\")\n}\n"
  },
  {
    "path": "test-integration/test_integration/fixtures/hello-image/cog.yaml",
    "content": "build:\n  python_version: \"3.11\"\npredict: \"predict.py:Predictor\"\nimage: \"r8.im/replicate/hello-image\"\n"
  },
  {
    "path": "test-integration/test_integration/fixtures/hello-image/predict.py",
    "content": "from cog import BasePredictor, Path\n\n\nclass Predictor(BasePredictor):\n    def predict(self, word: str) -> Path:\n        return Path(\"hello.webp\")\n"
  },
  {
    "path": "tools/compatgen/internal/cuda.go",
    "content": "package internal\n\nimport (\n\t\"cmp\"\n\t\"context\"\n\t\"encoding/json\"\n\t\"fmt\"\n\t\"slices\"\n\t\"sort\"\n\t\"strings\"\n\n\t\"github.com/anaskhan96/soup\"\n\t\"github.com/google/go-containerregistry/pkg/authn\"\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc FetchCUDABaseImages(ctx context.Context) ([]config.CUDABaseImage, error) {\n\turl := \"https://hub.docker.com/v2/repositories/nvidia/cuda/tags/?page_size=1000&name=devel-ubuntu&ordering=last_updated\"\n\ttags, err := fetchCUDABaseImageTags(url)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\tvar images []config.CUDABaseImage\n\tfor _, tag := range tags {\n\t\timage, err := parseCUDABaseImage(ctx, tag)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\timages = append(images, *image)\n\t}\n\n\t// stable sort for deterministic output\n\tslices.SortFunc(images, func(a, b config.CUDABaseImage) int {\n\t\treturn cmp.Or(\n\t\t\tcmp.Compare(a.CUDA, b.CUDA),\n\t\t\tcmp.Compare(a.CuDNN, b.CuDNN),\n\t\t\tcmp.Compare(a.Ubuntu, b.Ubuntu),\n\t\t\tcmp.Compare(a.Tag, b.Tag),\n\t\t)\n\t})\n\n\treturn images, nil\n}\n\nfunc fetchCUDABaseImageTags(url string) ([]string, error) {\n\ttags := []string{}\n\n\tresp, err := soup.Get(url)\n\tif err != nil {\n\t\treturn tags, fmt.Errorf(\"Failed to download %s: %w\", url, err)\n\t}\n\n\tvar results struct {\n\t\tNext    *string\n\t\tResults []struct {\n\t\t\tName string `json:\"name\"`\n\t\t} `json:\"results\"`\n\t}\n\tif err := json.Unmarshal([]byte(resp), &results); err != nil {\n\t\treturn tags, fmt.Errorf(\"Failed parse CUDA images json: %w\", err)\n\t}\n\n\tfor _, result := range results.Results {\n\t\ttag := result.Name\n\t\tif strings.Contains(tag, \"-cudnn\") && !strings.HasSuffix(tag, \"-rc\") {\n\t\t\ttags = append(tags, tag)\n\t\t}\n\t}\n\n\t// recursive case for pagination\n\tif results.Next != nil {\n\t\tnextURL := *results.Next\n\t\tnextTags, err := fetchCUDABaseImageTags(nextURL)\n\t\tif err != nil {\n\t\t\treturn tags, err\n\t\t}\n\t\ttags = append(tags, nextTags...)\n\t}\n\n\tsort.Sort(sort.Reverse(sort.StringSlice(tags)))\n\n\treturn tags, nil\n}\n\n// parseCUDABaseImage fetches the Docker image config for an nvidia/cuda tag\n// and extracts CUDA and CuDNN versions from environment variables. This is\n// necessary because newer nvidia/cuda tags no longer include the CuDNN version\n// in the tag itself (e.g. \"12.9.1-cudnn-devel-ubuntu24.04\" instead of\n// \"12.6.3-cudnn9-devel-ubuntu22.04\").\nfunc parseCUDABaseImage(ctx context.Context, tag string) (*config.CUDABaseImage, error) {\n\tfmt.Println(\"parsing\", tag)\n\n\tbaseImg := &config.CUDABaseImage{\n\t\tTag:     tag,\n\t\tIsDevel: strings.Contains(tag, \"-devel\"),\n\t}\n\n\tif parts := strings.Split(tag, \"ubuntu\"); len(parts) == 2 {\n\t\tbaseImg.Ubuntu = parts[1]\n\t} else {\n\t\treturn nil, fmt.Errorf(\"invalid tag, must end in ubuntu: %q\", tag)\n\t}\n\n\tref, err := name.ParseReference(fmt.Sprintf(\"nvidia/cuda:%s\", tag))\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to parse reference %s: %w\", tag, err)\n\t}\n\n\timg, err := remote.Image(ref, remote.WithContext(ctx), remote.WithAuthFromKeychain(authn.DefaultKeychain))\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to get image %s: %w\", tag, err)\n\t}\n\n\tcfg, err := img.ConfigFile()\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"failed to get config file %s: %w\", tag, err)\n\t}\n\n\tfor _, envVal := range cfg.Config.Env {\n\t\tparts := strings.SplitN(envVal, \"=\", 2)\n\t\tif len(parts) != 2 {\n\t\t\tcontinue\n\t\t}\n\t\tswitch parts[0] {\n\t\tcase \"CUDA_VERSION\":\n\t\t\tbaseImg.CUDA = parts[1]\n\t\tcase \"NV_CUDNN_VERSION\":\n\t\t\t// downstream code expects only the major version component\n\t\t\tbaseImg.CuDNN = strings.Split(parts[1], \".\")[0]\n\t\t}\n\t}\n\n\tif baseImg.CuDNN == \"\" {\n\t\treturn nil, fmt.Errorf(\"no CuDNN version found in image config for tag %s\", tag)\n\t}\n\n\treturn baseImg, nil\n}\n"
  },
  {
    "path": "tools/compatgen/internal/tensorflow.go",
    "content": "package internal\n\nimport (\n\t\"fmt\"\n\t\"strconv\"\n\t\"strings\"\n\n\t\"github.com/replicate/cog/pkg/util/version\"\n\n\t\"github.com/anaskhan96/soup\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n)\n\nfunc FetchTensorFlowCompatibilityMatrix() ([]config.TFCompatibility, error) {\n\turl := \"https://www.tensorflow.org/install/source\"\n\tminCudaVersion := strconv.Itoa(config.MinimumMajorCudaVersion)\n\n\tresp, err := soup.Get(url)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Failed to download %s: %w\", url, err)\n\t}\n\n\tdoc := soup.HTMLParse(resp)\n\tgpuHeading := doc.Find(\"h4\", \"id\", \"gpu\")\n\ttable := gpuHeading.FindNextElementSibling()\n\trows := table.FindAll(\"tr\")\n\n\tcompats := []config.TFCompatibility{}\n\tfor _, row := range rows[1:] {\n\t\tcells := row.FindAll(\"td\")\n\t\tgpuPackage, packageVersion := split2(cells[0].Text(), \"-\")\n\t\tpythonVersions, err := parsePythonVersionsCell(cells[1].Text())\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tcuDNN := cells[4].Text()\n\t\tcuda := cells[5].Text()\n\n\t\tif !version.Greater(cuda, minCudaVersion) && !version.Equal(cuda, minCudaVersion) {\n\t\t\tcontinue\n\t\t}\n\n\t\tcompat := config.TFCompatibility{\n\t\t\tTF:           packageVersion,\n\t\t\tTFCPUPackage: \"tensorflow==\" + packageVersion,\n\t\t\tTFGPUPackage: gpuPackage + \"==\" + packageVersion,\n\t\t\tCUDA:         cuda,\n\t\t\tCuDNN:        cuDNN,\n\t\t\tPythons:      pythonVersions,\n\t\t}\n\t\tcompats = append(compats, compat)\n\t}\n\n\t// sanity check\n\tif len(compats) < 12 {\n\t\treturn nil, fmt.Errorf(\"Tensorflow compatibility matrix only had %d rows, has the html changed?\", len(compats))\n\t}\n\n\treturn compats, nil\n}\n\nfunc parsePythonVersionsCell(val string) ([]string, error) {\n\tversions := []string{}\n\tparts := strings.SplitSeq(val, \",\")\n\tfor part := range parts {\n\t\tpart = strings.TrimSpace(part)\n\t\tif strings.Contains(part, \"-\") {\n\t\t\tstart, end := split2(part, \"-\")\n\t\t\tstartMajor, startMinor, err := splitPythonVersion(start)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t\tendMajor, endMinor, err := splitPythonVersion(end)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\n\t\t\tif startMajor != endMajor {\n\t\t\t\treturn nil, fmt.Errorf(\"Invalid start and end minor versions: %d, %d\", startMajor, endMajor)\n\t\t\t}\n\t\t\tfor minor := startMinor; minor <= endMinor; minor++ {\n\t\t\t\tversions = append(versions, newVersion(startMajor, minor))\n\t\t\t}\n\t\t} else {\n\t\t\tversions = append(versions, part)\n\t\t}\n\t}\n\treturn versions, nil\n}\n\nfunc newVersion(major int, minor int) string {\n\treturn fmt.Sprintf(\"%d.%d\", major, minor)\n}\n\nfunc splitPythonVersion(version string) (major int, minor int, err error) {\n\tversion = strings.TrimSpace(version)\n\tmajorStr, minorStr := split2(version, \".\")\n\tmajor, err = strconv.Atoi(majorStr)\n\tif err != nil {\n\t\treturn 0, 0, err\n\t}\n\tminor, err = strconv.Atoi(minorStr)\n\tif err != nil {\n\t\treturn 0, 0, err\n\t}\n\treturn major, minor, nil\n}\n"
  },
  {
    "path": "tools/compatgen/internal/torch.go",
    "content": "package internal\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"net/url\"\n\t\"regexp\"\n\t\"slices\"\n\t\"sort\"\n\t\"strconv\"\n\t\"strings\"\n\n\t\"github.com/anaskhan96/soup\"\n\n\t\"github.com/hashicorp/go-version\"\n\n\t\"github.com/replicate/cog/pkg/config\"\n\t\"github.com/replicate/cog/pkg/env\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nvar ErrorBadPytorchFormat = errors.New(\"The pytorch version format could not be parsed.\")\n\nfunc FetchTorchCompatibilityMatrix() ([]config.TorchCompatibility, error) {\n\tcompats := []config.TorchCompatibility{}\n\tvar err error\n\tcompats, err = fetchCurrentTorchVersions(compats)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tcompats, err = fetchPreviousTorchVersions(compats)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Remove entries with no supported Python versions\n\tfiltered := make([]config.TorchCompatibility, 0, len(compats))\n\tfor _, c := range compats {\n\t\tif len(c.Pythons) > 0 {\n\t\t\tfiltered = append(filtered, c)\n\t\t} else {\n\t\t\tconsole.Warnf(\"Dropping %s: no supported Python versions\", c.Torch)\n\t\t}\n\t}\n\tcompats = filtered\n\n\t// sanity check\n\tif len(compats) < 21 {\n\t\treturn nil, fmt.Errorf(\"PyTorch compatibility matrix only had %d rows, has the html changed?\", len(compats))\n\t}\n\n\treturn compats, nil\n}\n\nfunc FetchTorchPackages(name string) ([]TorchPackage, error) {\n\turl := pytorchURL(name)\n\treturn fetchTorchPackagesFromURL(url)\n}\n\nfunc getLatestVersion(packages []TorchPackage) string {\n\tlatestVersion, _ := version.NewVersion(\"0.0.0\")\n\tfor _, pkg := range packages {\n\t\tv, err := version.NewVersion(pkg.Version)\n\t\tif err != nil {\n\t\t\tfmt.Println(\"error parsing version:\", pkg.Version)\n\t\t\tcontinue\n\t\t}\n\t\tif v.GreaterThan(latestVersion) {\n\t\t\tlatestVersion = v\n\t\t}\n\t}\n\treturn latestVersion.String()\n}\n\nfunc fetchCurrentTorchVersions(compats []config.TorchCompatibility) ([]config.TorchCompatibility, error) {\n\t// For the latest PyTorch version, we can just grab the latest of each packages from the repository.\n\t// We then install the packages in the same way as we do for 1.12.x:\n\t// https://pytorch.org/get-started/previous-versions/#v1121\n\n\ttorchPackages, err := FetchTorchPackages(\"torch\")\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Error fetching PyTorch packages: %w\", err)\n\t}\n\ttorchVisionPackages, err := FetchTorchPackages(\"torchvision\")\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Error fetching PyTorch packages: %w\", err)\n\t}\n\ttorchAudioPackages, err := FetchTorchPackages(\"torchaudio\")\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Error fetching PyTorch packages: %w\", err)\n\t}\n\n\tlatestTorchVersion := getLatestVersion(torchPackages)\n\tlatestTorchvisionVersion := getLatestVersion(torchVisionPackages)\n\tlatestTorchaudioVersion := getLatestVersion(torchAudioPackages)\n\n\ttorchCompats := map[string]config.TorchCompatibility{}\n\n\tfor _, pkg := range torchPackages {\n\t\tif pkg.Version != latestTorchVersion {\n\t\t\tcontinue\n\t\t}\n\n\t\tif val, ok := torchCompats[pkg.Name]; ok {\n\t\t\tif !slices.Contains(val.Pythons, pkg.PythonVersion) {\n\t\t\t\tval.Pythons = append(val.Pythons, pkg.PythonVersion)\n\t\t\t}\n\t\t\ttorchCompats[pkg.Name] = val\n\t\t} else {\n\t\t\ttorchCompats[pkg.Name] = config.TorchCompatibility{\n\t\t\t\tTorch:         pkg.Name,\n\t\t\t\tTorchvision:   latestTorchvisionVersion,\n\t\t\t\tTorchaudio:    latestTorchaudioVersion,\n\t\t\t\tCUDA:          pkg.CUDA,\n\t\t\t\tExtraIndexURL: pytorchURL(pkg.Variant),\n\t\t\t\tPythons:       []string{pkg.PythonVersion},\n\t\t\t}\n\n\t\t}\n\t}\n\n\tfor _, compat := range torchCompats {\n\t\tcompats = append(compats, compat)\n\t}\n\n\treturn compats, nil\n}\n\nfunc parseTorchInstallString(s string, defaultVersions map[string]string, cuda *string) (*config.TorchCompatibility, error) {\n\t// for example:\n\t// pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113\n\t// pip install torch==1.8.0+cpu torchvision==0.9.0+cpu torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html\n\n\tlibVersions := map[string]string{}\n\n\tfindLinks := \"\"\n\textraIndexURL := \"\"\n\tskipNext := false\n\n\t// Simple parser for pip install strings\n\tfields := strings.Fields(s)\n\tfor i, item := range fields {\n\t\t// Ideally we want to be able to consume the next token, but golang has no simple way of doing that without constructing a channel\n\t\tif skipNext {\n\t\t\tskipNext = false\n\t\t\tcontinue\n\t\t}\n\t\tswitch item {\n\t\tcase \"pip\", \"pip3\", \"install\":\n\t\t\tcontinue\n\t\tcase \"-f\":\n\t\t\tfindLinks = fields[i+1]\n\t\t\tskipNext = true\n\t\t\tcontinue\n\t\tcase \"--index-url\", \"--extra-index-url\":\n\t\t\textraIndexURL = fields[i+1]\n\t\t\tskipNext = true\n\t\t\tcontinue\n\t\t}\n\n\t\tlibParts := strings.Split(item, \"==\")\n\t\tlibName := libParts[0]\n\t\tif _, ok := defaultVersions[libName]; !ok {\n\t\t\treturn nil, fmt.Errorf(\"Unknown token when parsing torch string: %s\", item)\n\t\t}\n\t\tif len(libParts) == 1 {\n\t\t\tlibVersions[libName] = defaultVersions[libName]\n\t\t} else {\n\t\t\tlibVersions[libName] = libParts[1]\n\t\t}\n\n\t}\n\n\ttorch, ok := libVersions[\"torch\"]\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"Missing torch version\")\n\t}\n\ttorchvision, ok := libVersions[\"torchvision\"]\n\tif !ok {\n\t\treturn nil, fmt.Errorf(\"Missing torchvision version\")\n\t}\n\ttorchaudio := libVersions[\"torchaudio\"]\n\n\tpythons, err := FindCompatiblePythonVersions(torch, torchvision, torchaudio, extraIndexURL, findLinks)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\treturn &config.TorchCompatibility{\n\t\tTorch:         torch,\n\t\tTorchvision:   torchvision,\n\t\tTorchaudio:    torchaudio,\n\t\tFindLinks:     findLinks,\n\t\tExtraIndexURL: extraIndexURL,\n\t\tCUDA:          cuda,\n\t\tPythons:       pythons,\n\t}, nil\n}\n\nfunc fetchPreviousTorchVersions(compats []config.TorchCompatibility) ([]config.TorchCompatibility, error) {\n\t// For previous versions, we need to scrape the PyTorch website.\n\t// The reason we can't fetch it from the PyPI repository like the latest version is\n\t// because we don't know what versions of torch, torchvision, and torchaudio are compatible with each other.\n\n\turl := \"https://pytorch.org/get-started/previous-versions/\"\n\tresp, err := soup.Get(url)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Failed to download %s: %w\", url, err)\n\t}\n\tdoc := soup.HTMLParse(resp)\n\n\tfor _, h5 := range doc.FindAll(\"h5\") {\n\t\tif strings.TrimSpace(h5.Text()) == \"Linux and Windows\" {\n\t\t\thighlight := h5.FindNextElementSibling()\n\t\t\tcode := highlight.Find(\"code\")\n\t\t\tcompats, err = parsePreviousTorchVersionsCode(code.Text(), compats)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, err\n\t\t\t}\n\t\t}\n\t}\n\treturn compats, nil\n}\n\nfunc parsePreviousTorchVersionsCode(code string, compats []config.TorchCompatibility) ([]config.TorchCompatibility, error) {\n\t// e.g.\n\t// # CUDA 10.1\n\t// pip install torch==1.5.0+cu101 torchvision==0.6.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html\n\n\tsupportedLibrarySet := map[string]string{\n\t\t\"torch\": \"\", \"torchvision\": \"\", \"torchaudio\": \"\",\n\t}\n\n\tvar cuda *string\n\tskipSection := false\n\n\tfor line := range strings.SplitSeq(code, \"\\n\") {\n\t\t// Set section\n\t\tif strings.HasPrefix(line, \"#\") {\n\t\t\tskipSection = false\n\t\t\trawArch := strings.ToLower(line[2:])\n\t\t\tswitch {\n\t\t\tcase strings.HasPrefix(rawArch, \"cuda\"):\n\t\t\t\t_, c := split2(rawArch, \" \")\n\t\t\t\tcuda = &c\n\t\t\tcase rawArch == \"cpu only\":\n\t\t\t\tcuda = nil\n\t\t\tcase strings.HasPrefix(rawArch, \"rocm\"):\n\t\t\t\tcuda = nil\n\t\t\t\tskipSection = true\n\t\t\tdefault:\n\t\t\t\t// Ignore additional heading lines (notes, etc)\n\t\t\t\tcontinue\n\t\t\t}\n\t\t}\n\n\t\t// In a ROCM section, so skip\n\t\tif skipSection {\n\t\t\tcontinue\n\t\t}\n\n\t\t// conda install etc\n\t\tif !strings.HasPrefix(line, \"pip install \") {\n\t\t\tcontinue\n\t\t}\n\t\tcompat, err := parseTorchInstallString(line, supportedLibrarySet, cuda)\n\t\tif err != nil {\n\t\t\treturn nil, err\n\t\t}\n\t\tfixTorchCompatibility(compat)\n\n\t\tcompats = append(compats, *compat)\n\t}\n\treturn compats, nil\n}\n\n// torchvision==0.8.0 should actually be 0.8.1, this is a bug on the website\nfunc fixTorchCompatibility(compat *config.TorchCompatibility) {\n\tif strings.HasPrefix(compat.Torchvision, \"0.8.0\") {\n\t\tcompat.Torchvision = strings.ReplaceAll(compat.Torchvision, \"0.8.0\", \"0.8.1\")\n\t}\n}\n\nfunc basePytorchURL() string {\n\treturn env.SchemeFromEnvironment() + \"://\" + env.PytorchHostFromEnvironment() + \"/whl\"\n}\n\nfunc pytorchURL(name string) string {\n\turl := fmt.Sprintf(basePytorchURL()+\"/%s/\", name)\n\treturn url\n}\n\nfunc ExtractSubFeaturesFromPytorchVersion(pytorchVersion string) (string, string, string, string, string, error) {\n\tdecoded, err := url.PathUnescape(pytorchVersion)\n\tif err != nil {\n\t\treturn \"\", \"\", \"\", \"\", \"\", fmt.Errorf(\"failed to decode filename: %w\", err)\n\t}\n\n\tpkgRegexp := regexp.MustCompile(\n\t\t`.+?-(?P\\d+(?:\\.\\d+)*)(?P(?:[._]?(?:post|dev|rc)\\d+)*)?(?:\\+(?P[a-z0-9_.]+))?-(?P[a-z0-9_.]+)-[a-z0-9_.]+-(?P.+?)\\.whl`,\n\t)\n\n\tmatches := pkgRegexp.FindStringSubmatch(decoded)\n\tif len(matches) == 0 {\n\t\treturn \"\", \"\", \"\", \"\", \"\", fmt.Errorf(\"invalid PyTorch wheel filename: %s\", decoded)\n\t}\n\n\tgroupMap := make(map[string]string)\n\tfor i, name := range pkgRegexp.SubexpNames() {\n\t\tif i != 0 && name != \"\" {\n\t\t\tgroupMap[name] = matches[i]\n\t\t}\n\t}\n\n\tbase := groupMap[\"basever\"]\n\tsuffix := groupMap[\"suffix\"]\n\tvariant := groupMap[\"variant\"]\n\tpyverRaw := groupMap[\"pyver\"]\n\tplatform := groupMap[\"platform\"]\n\n\tname := base + suffix\n\tif variant != \"\" {\n\t\tname += \"+\" + variant\n\t}\n\tversion := base\n\n\tpyver := pyverRaw\n\tif strings.HasPrefix(pyverRaw, \"cp\") {\n\t\tpyver = pyverRaw[len(\"cp\"):]\n\t}\n\n\treturn name, version, variant, pyver, platform, nil\n}\n\nfunc FindCompatiblePythonVersions(torchVersion string, torchVisionVersion string, torchAudioVersion string, extraIndexUrl string, findLinksUrl string) ([]string, error) {\n\tif extraIndexUrl == \"\" && findLinksUrl == \"\" {\n\t\textraIndexUrl = basePytorchURL()\n\t}\n\turl := extraIndexUrl\n\tif url == \"\" {\n\t\turl = findLinksUrl\n\t}\n\n\t// Correct 0.8.0 torchvision to 0.8.1, this is a bug on pytorch.org\n\tif strings.HasPrefix(torchVisionVersion, \"0.8.0\") {\n\t\ttorchVisionVersion = strings.ReplaceAll(torchVisionVersion, \"0.8.0\", \"0.8.1\")\n\t}\n\n\ttorchPkgs, err := findTorchPackagesWithVersion(\"torch\", url, torchVersion, url != findLinksUrl)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\ttorchVisionPkgs, err := findTorchPackagesWithVersion(\"torchvision\", url, torchVisionVersion, url != findLinksUrl)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\ttorchAudioPkgs, err := findTorchPackagesWithVersion(\"torchaudio\", url, torchAudioVersion, url != findLinksUrl)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\n\t// Get initial list of valid python versions from torch\n\tpythonVersions := map[string]bool{}\n\tfor _, pkg := range torchPkgs {\n\t\tpythonVersions[pkg.PythonVersion] = true\n\t}\n\n\t// Check that torchaudio/torchvision shares these python versions\n\textraPkgs := [][]TorchPackage{}\n\tif torchVisionVersion != \"\" {\n\t\textraPkgs = append(extraPkgs, torchVisionPkgs)\n\t}\n\tif torchAudioVersion != \"\" {\n\t\textraPkgs = append(extraPkgs, torchAudioPkgs)\n\t}\n\tfor _, pkgs := range extraPkgs {\n\t\tpkgPythonVersions := map[string]bool{}\n\t\tfor _, pkg := range pkgs {\n\t\t\tpkgPythonVersions[pkg.PythonVersion] = true\n\t\t}\n\t\tfor pythonVersion := range pythonVersions {\n\t\t\t_, ok := pkgPythonVersions[pythonVersion]\n\t\t\tif !ok {\n\t\t\t\tdelete(pythonVersions, pythonVersion)\n\t\t\t}\n\t\t}\n\t}\n\n\tvalidPythonVersions := make([]string, 0, len(pythonVersions))\n\tfor k := range pythonVersions {\n\t\tvalidPythonVersions = append(validPythonVersions, k)\n\t}\n\tsort.Strings(validPythonVersions)\n\n\treturn validPythonVersions, nil\n}\n\nfunc fetchTorchPackagesFromURL(url string) ([]TorchPackage, error) {\n\tresp, err := soup.Get(url)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"Failed to download %s: %w\", url, err)\n\t}\n\tdoc := soup.HTMLParse(resp)\n\tlinks := doc.FindAll(\"a\")\n\tpackages := []TorchPackage{}\n\tfor _, link := range links {\n\t\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(link.Text())\n\t\tif err != nil {\n\t\t\tconsole.Warnf(\"Failed to parse pytorch version: %v\", err)\n\t\t\tcontinue\n\t\t}\n\t\tif (platform != \"linux_x86_64\" && platform != \"manylinux_2_28_x86_64\" && platform != \"manylinux1_x86_64\") || strings.Contains(name, \".cxx\") {\n\t\t\tcontinue\n\t\t}\n\n\t\tvar cuda *string\n\t\tswitch {\n\t\tcase variant == \"cpu\":\n\t\t\tcuda = nil\n\t\tcase variant == \"\":\n\t\t\tcuda = nil\n\t\tcase strings.HasPrefix(variant, \"cu\"):\n\t\t\t// cu92 -> 9.2\n\t\t\tc := strings.TrimPrefix(variant, \"cu\")\n\t\t\tc = c[:len(c)-1] + \".\" + c[len(c)-1:]\n\t\t\tcuda = &c\n\t\tdefault:\n\t\t\t// rocm etc\n\t\t\tcontinue\n\t\t}\n\n\t\t// 310 -> 3.10\n\t\tpythonVersion = pythonVersion[:1] + \".\" + pythonVersion[1:]\n\t\tif minor, ok := strings.CutPrefix(pythonVersion, \"3.\"); ok {\n\t\t\tminorInt, err := strconv.Atoi(minor)\n\t\t\tif err != nil {\n\t\t\t\treturn nil, fmt.Errorf(\"invalid python version %q: %w\", pythonVersion, err)\n\t\t\t}\n\t\t\tif minorInt < config.MinimumMinorPythonVersion {\n\t\t\t\tcontinue\n\t\t\t}\n\t\t}\n\n\t\tpkg := TorchPackage{\n\t\t\tName:          name,\n\t\t\tVersion:       version,\n\t\t\tVariant:       variant,\n\t\t\tCUDA:          cuda,\n\t\t\tPythonVersion: pythonVersion,\n\t\t}\n\n\t\tfound := false\n\t\tfor _, currentPkg := range packages {\n\t\t\tif currentPkg.Equals(pkg) {\n\t\t\t\tfound = true\n\t\t\t\tbreak\n\t\t\t}\n\t\t}\n\t\tif found {\n\t\t\tcontinue\n\t\t}\n\n\t\tpackages = append(packages, pkg)\n\t}\n\treturn packages, nil\n}\n\nfunc findTorchPackagesWithVersion(pkgName string, url string, version string, appendPkg bool) ([]TorchPackage, error) {\n\tif appendPkg {\n\t\turl = url + \"/\" + pkgName\n\t}\n\tpkgs, err := fetchTorchPackagesFromURL(url)\n\tif err != nil {\n\t\treturn nil, err\n\t}\n\tvalidPkgs := []TorchPackage{}\n\tfor _, pkg := range pkgs {\n\t\tif pkg.Version != version && pkg.Name != version {\n\t\t\tcontinue\n\t\t}\n\t\tvalidPkgs = append(validPkgs, pkg)\n\t}\n\treturn validPkgs, nil\n}\n"
  },
  {
    "path": "tools/compatgen/internal/torch_package.go",
    "content": "package internal\n\ntype TorchPackage struct {\n\tName          string\n\tVersion       string\n\tVariant       string\n\tCUDA          *string\n\tPythonVersion string\n}\n\nfunc (c *TorchPackage) Equals(other TorchPackage) bool {\n\tif c.CUDA != other.CUDA {\n\t\tif c.CUDA != nil && other.CUDA != nil && *c.CUDA != *other.CUDA {\n\t\t\treturn false\n\t\t}\n\t}\n\treturn c.Name == other.Name && c.Version == other.Version && c.Variant == other.Variant && c.PythonVersion == other.PythonVersion\n}\n"
  },
  {
    "path": "tools/compatgen/internal/torch_test.go",
    "content": "package internal\n\nimport (\n\t\"log\"\n\t\"net/http\"\n\t\"net/http/httptest\"\n\t\"net/url\"\n\t\"os\"\n\t\"strings\"\n\t\"testing\"\n\n\t\"github.com/stretchr/testify/require\"\n\n\t\"github.com/replicate/cog/pkg/env\"\n)\n\nfunc TestFetchTorchPackages(t *testing.T) {\n\tserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {\n\t\tcontent, err := os.ReadFile(\"torch_test.html\")\n\t\tif err != nil {\n\t\t\tlog.Fatalf(\"Error reading file: %v\", err)\n\t\t}\n\t\tw.WriteHeader(http.StatusOK)\n\t\tw.Write(content)\n\t}))\n\tdefer server.Close()\n\turl, err := url.Parse(server.URL)\n\trequire.NoError(t, err)\n\tt.Setenv(env.SchemeEnvVarName, url.Scheme)\n\tt.Setenv(env.PytorchHostEnvVarName, url.Host)\n\n\ttorchPackages, err := FetchTorchPackages(\"torch\")\n\trequire.NoError(t, err)\n\ttorch271Packages := []TorchPackage{}\n\tfor _, pkg := range torchPackages {\n\t\tif strings.Contains(pkg.Name, \"2.7.1+cu128\") {\n\t\t\ttorch271Packages = append(torch271Packages, pkg)\n\t\t}\n\t}\n\tcuda128 := \"12.8\"\n\n\trequire.Equal(t, []TorchPackage{\n\t\t{\n\t\t\tName:          \"2.7.1+cu128\",\n\t\t\tVersion:       \"2.7.1\",\n\t\t\tVariant:       \"cu128\",\n\t\t\tCUDA:          &cuda128,\n\t\t\tPythonVersion: \"3.10\",\n\t\t},\n\t\t{\n\t\t\tName:          \"2.7.1+cu128\",\n\t\t\tVersion:       \"2.7.1\",\n\t\t\tVariant:       \"cu128\",\n\t\t\tCUDA:          &cuda128,\n\t\t\tPythonVersion: \"3.11\",\n\t\t},\n\t\t{\n\t\t\tName:          \"2.7.1+cu128\",\n\t\t\tVersion:       \"2.7.1\",\n\t\t\tVariant:       \"cu128\",\n\t\t\tCUDA:          &cuda128,\n\t\t\tPythonVersion: \"3.12\",\n\t\t},\n\t\t{\n\t\t\tName:          \"2.7.1+cu128\",\n\t\t\tVersion:       \"2.7.1\",\n\t\t\tVariant:       \"cu128\",\n\t\t\tCUDA:          &cuda128,\n\t\t\tPythonVersion: \"3.13\",\n\t\t},\n\t}, torch271Packages)\n}\n\nfunc TestIsValidPytorchVersionFormat(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torch-2.7.1+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"2.7.1+cpu.cxx11.abi\", name)\n\trequire.Equal(t, \"2.7.1\", version)\n\trequire.Equal(t, \"cpu.cxx11.abi\", variant)\n\trequire.Equal(t, \"312\", pythonVersion)\n\trequire.Equal(t, \"linux_x86_64\", platform)\n}\n\nfunc TestIsValidPytorchVersionFormatWithOldVersion(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torch-1.10.0+cpu-cp310-cp310-linux_x86_64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"1.10.0+cpu\", name)\n\trequire.Equal(t, \"1.10.0\", version)\n\trequire.Equal(t, \"cpu\", variant)\n\trequire.Equal(t, \"310\", pythonVersion)\n\trequire.Equal(t, \"linux_x86_64\", platform)\n}\n\nfunc TestIsValidPytorchAudioVersionFormat(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torchaudio-2.7.1+xpu-cp313-cp313t-win_amd64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"2.7.1+xpu\", name)\n\trequire.Equal(t, \"2.7.1\", version)\n\trequire.Equal(t, \"xpu\", variant)\n\trequire.Equal(t, \"313\", pythonVersion)\n\trequire.Equal(t, \"win_amd64\", platform)\n}\n\nfunc TestIsValidPytorchAudioVersionFormatBasic(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torchaudio-0.8.1-cp39-none-win_amd64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"0.8.1\", name)\n\trequire.Equal(t, \"0.8.1\", version)\n\trequire.Equal(t, \"\", variant)\n\trequire.Equal(t, \"39\", pythonVersion)\n\trequire.Equal(t, \"win_amd64\", platform)\n}\n\nfunc TestIsValidPytorchVisionVersionFormatPostRelease(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torchvision-0.4.1.post2-cp37-cp37m-macosx_10_9_x86_64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"0.4.1.post2\", name)\n\trequire.Equal(t, \"0.4.1\", version)\n\trequire.Equal(t, \"\", variant)\n\trequire.Equal(t, \"37\", pythonVersion)\n\trequire.Equal(t, \"macosx_10_9_x86_64\", platform)\n}\n\nfunc TestIsValidPytorchVisionEarlyVersion(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torchvision-0.14.1+cu116-cp310-cp310-linux_x86_64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"0.14.1+cu116\", name)\n\trequire.Equal(t, \"0.14.1\", version)\n\trequire.Equal(t, \"cu116\", variant)\n\trequire.Equal(t, \"310\", pythonVersion)\n\trequire.Equal(t, \"linux_x86_64\", platform)\n}\n\nfunc TestIsValidPytorchAudioEarlyVersion(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torchaudio-0.9.1-cp39-cp39-linux_x86_64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"0.9.1\", name)\n\trequire.Equal(t, \"0.9.1\", version)\n\trequire.Equal(t, \"\", variant)\n\trequire.Equal(t, \"39\", pythonVersion)\n\trequire.Equal(t, \"linux_x86_64\", platform)\n}\n\nfunc TestURLEncodedVersion(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torchtext-0.17.0%2Bcpu-cp39-cp39-win_amd64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"0.17.0+cpu\", name)\n\trequire.Equal(t, \"0.17.0\", version)\n\trequire.Equal(t, \"cpu\", variant)\n\trequire.Equal(t, \"39\", pythonVersion)\n\trequire.Equal(t, \"win_amd64\", platform)\n}\n\nfunc TestVersionUnderFolder(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"cu111/torch-1.8.0%2Bcu111-cp36-cp36m-linux_x86_64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"1.8.0+cu111\", name)\n\trequire.Equal(t, \"1.8.0\", version)\n\trequire.Equal(t, \"cu111\", variant)\n\trequire.Equal(t, \"36\", pythonVersion)\n\trequire.Equal(t, \"linux_x86_64\", platform)\n}\n\nfunc TestPythonMVersion(t *testing.T) {\n\tname, version, variant, pythonVersion, platform, err := ExtractSubFeaturesFromPytorchVersion(\"torchaudio-0.7.2-cp36-cp36m-linux_x86_64.whl\")\n\trequire.NoError(t, err)\n\trequire.Equal(t, \"0.7.2\", name)\n\trequire.Equal(t, \"0.7.2\", version)\n\trequire.Equal(t, \"\", variant)\n\trequire.Equal(t, \"36\", pythonVersion)\n\trequire.Equal(t, \"linux_x86_64\", platform)\n}\n"
  },
  {
    "path": "tools/compatgen/internal/torch_test.html",
    "content": "\n\n\n  \n    
Links for torch
\n    torch-2.0.0+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.0+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.0+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.0+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.1+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.1+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.1+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.1+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.0+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.0+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.0+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.0+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.1+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.1+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.1+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.1+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.2+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.2+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.2+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.2+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.0+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.0+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.0+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.0+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.0+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.1+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.1+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.1+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.1+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.1+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.2+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.2+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.2+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.2+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.2+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.0+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.0+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.0+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.0+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.0+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.1+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.1+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.1+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.1+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.1+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.0+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.0+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.0+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.0+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.0+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.1+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.1+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.1+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.1+cpu.cxx11.abi-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.1+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.0+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.0+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.0+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.0+cpu.cxx11.abi-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.0+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.1+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.1+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.1+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.1+cpu.cxx11.abi-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.1+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.6.0+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.6.0+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.6.0+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.6.0+cpu.cxx11.abi-cp313-cp313-linux_x86_64.whl
\n    torch-2.6.0+cpu.cxx11.abi-cp313-cp313t-linux_x86_64.whl
\n    torch-2.6.0+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.7.0+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.7.0+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.7.0+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.7.0+cpu.cxx11.abi-cp313-cp313-linux_x86_64.whl
\n    torch-2.7.0+cpu.cxx11.abi-cp313-cp313t-linux_x86_64.whl
\n    torch-2.7.0+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-2.7.1+cpu.cxx11.abi-cp310-cp310-linux_x86_64.whl
\n    torch-2.7.1+cpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
\n    torch-2.7.1+cpu.cxx11.abi-cp312-cp312-linux_x86_64.whl
\n    torch-2.7.1+cpu.cxx11.abi-cp313-cp313-linux_x86_64.whl
\n    torch-2.7.1+cpu.cxx11.abi-cp313-cp313t-linux_x86_64.whl
\n    torch-2.7.1+cpu.cxx11.abi-cp39-cp39-linux_x86_64.whl
\n    torch-0.3.0.post4-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post4-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.1-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.0-cp27-cp27m-linux_x86_64.whl
\n    torch-0.4.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-linux_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-win_amd64.whl
\n    torch-0.4.0-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.0-cp36-cp36m-win_amd64.whl
\n    torch-0.4.1-cp27-cp27m-linux_x86_64.whl
\n    torch-0.4.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-win_amd64.whl
\n    torch-0.4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.1-cp36-cp36m-win_amd64.whl
\n    torch-0.4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-0.4.1-cp37-cp37m-win_amd64.whl
\n    torch-0.4.1.post2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.0-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.0-cp27-none-macosx_10_6_x86_64.whl
\n    torch-1.0.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.0-cp35-cp35m-win_amd64.whl
\n    torch-1.0.0-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.0.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.0-cp36-cp36m-win_amd64.whl
\n    torch-1.0.0-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.0.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.0-cp37-cp37m-win_amd64.whl
\n    torch-1.0.0-cp37-none-macosx_10_7_x86_64.whl
\n    torch-1.0.1-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.1-cp27-none-macosx_10_6_x86_64.whl
\n    torch-1.0.1-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.1-cp35-cp35m-win_amd64.whl
\n    torch-1.0.1-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.1-cp36-cp36m-win_amd64.whl
\n    torch-1.0.1-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.1-cp37-cp37m-win_amd64.whl
\n    torch-1.0.1-cp37-none-macosx_10_7_x86_64.whl
\n    torch-1.0.1.post2-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.1.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.1.0-cp27-cp27m-linux_x86_64.whl
\n    torch-1.1.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.1.0-cp27-none-macosx_10_6_x86_64.whl
\n    torch-1.1.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.1.0-cp35-cp35m-win_amd64.whl
\n    torch-1.1.0-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.1.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.1.0-cp36-cp36m-win_amd64.whl
\n    torch-1.1.0-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.1.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.1.0-cp37-cp37m-win_amd64.whl
\n    torch-1.1.0-cp37-none-macosx_10_7_x86_64.whl
\n    torch-1.1.0.post2-cp27-none-macosx_10_6_x86_64.whl
\n    torch-1.1.0.post2-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.1.0.post2-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.1.0.post2-cp37-none-macosx_10_7_x86_64.whl
\n    torch-1.10.0+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.10.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.10.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.10.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.10.0-cp36-cp36m-manylinux2014_aarch64.whl
\n    torch-1.10.0-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.10.0-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.10.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.10.0-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.10.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.10.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.10.0-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.10.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.10.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.10.1+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.1+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.10.1+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.1+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.10.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.10.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.10.1-cp36-cp36m-manylinux2014_aarch64.whl
\n    torch-1.10.1-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.10.1-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.10.1-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.10.1-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.10.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.10.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.10.1-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.10.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.10.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.10.2+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.2+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.10.2+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.2+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.10.2+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.2+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.10.2+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.2+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.10.2-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-1.10.2-cp36-cp36m-manylinux2014_aarch64.whl
\n    torch-1.10.2-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.10.2-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.10.2-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.10.2-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.10.2-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.10.2-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.10.2-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.10.2-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.10.2-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.11.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-1.11.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-1.11.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.11.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.11.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.11.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.11.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.11.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.11.0-cp310-none-macosx_10_9_x86_64.whl
\n    torch-1.11.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-1.11.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.11.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.11.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.11.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.11.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.12.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-1.12.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.12.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.12.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.12.0-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-1.12.0-cp310-none-macosx_10_9_x86_64.whl
\n    torch-1.12.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-1.12.0-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.12.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.12.0-cp37-none-macosx_11_0_arm64.whl
\n    torch-1.12.0-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.12.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.12.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.12.0-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.12.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.12.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.12.1+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-1.12.1+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.1+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.12.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.12.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.12.1-cp310-none-macosx_10_9_x86_64.whl
\n    torch-1.12.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-1.12.1-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.12.1-cp37-none-macosx_11_0_arm64.whl
\n    torch-1.12.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.12.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.12.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.12.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.13.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-1.13.0+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-1.13.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.13.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.13.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.13.0-cp310-none-macosx_10_9_x86_64.whl
\n    torch-1.13.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-1.13.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.13.0-cp37-none-macosx_11_0_arm64.whl
\n    torch-1.13.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.13.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.13.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.13.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.13.1+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-1.13.1+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-1.13.1+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.1+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.13.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.13.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.13.1-cp310-none-macosx_10_9_x86_64.whl
\n    torch-1.13.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-1.13.1-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.13.1-cp37-none-macosx_11_0_arm64.whl
\n    torch-1.13.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.13.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.13.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.13.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.2.0+cpu-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-1.2.0+cpu-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-1.2.0+cpu-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-1.2.0+cpu-cp35-cp35m-win_amd64.whl
\n    torch-1.2.0+cpu-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-1.2.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.2.0+cpu-cp37-cp37m-manylinux1_x86_64.whl
\n    torch-1.2.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.2.0-cp27-none-macosx_10_6_x86_64.whl
\n    torch-1.2.0-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.2.0-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.2.0-cp37-none-macosx_10_7_x86_64.whl
\n    torch-1.3.0+cpu-cp27-cp27m-linux_x86_64.whl
\n    torch-1.3.0+cpu-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.3.0+cpu-cp35-cp35m-linux_x86_64.whl
\n    torch-1.3.0+cpu-cp35-cp35m-win_amd64.whl
\n    torch-1.3.0+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.3.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.3.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.3.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.3.0-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.3.0-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.3.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.3.0.post2-cp27-none-macosx_10_7_x86_64.whl
\n    torch-1.3.0.post2-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.3.0.post2-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.3.0.post2-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.3.1+cpu-cp27-cp27m-linux_x86_64.whl
\n    torch-1.3.1+cpu-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.3.1+cpu-cp35-cp35m-linux_x86_64.whl
\n    torch-1.3.1+cpu-cp35-cp35m-win_amd64.whl
\n    torch-1.3.1+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.3.1+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.3.1+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.3.1+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.3.1-cp27-none-macosx_10_7_x86_64.whl
\n    torch-1.3.1-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.3.1-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.3.1-cp37-none-macosx_10_7_x86_64.whl
\n    torch-1.4.0+cpu-cp27-cp27m-linux_x86_64.whl
\n    torch-1.4.0+cpu-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.4.0+cpu-cp35-cp35m-linux_x86_64.whl
\n    torch-1.4.0+cpu-cp35-cp35m-win_amd64.whl
\n    torch-1.4.0+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.4.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.4.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.4.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.4.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.4.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.4.0-cp27-none-macosx_10_7_x86_64.whl
\n    torch-1.4.0-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.4.0-cp36-none-macosx_10_7_x86_64.whl
\n    torch-1.4.0-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.4.0-cp37-none-macosx_10_7_x86_64.whl
\n    torch-1.4.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.4.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.5.0+cpu-cp27-cp27m-linux_x86_64.whl
\n    torch-1.5.0+cpu-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.5.0+cpu-cp35-cp35m-linux_x86_64.whl
\n    torch-1.5.0+cpu-cp35-cp35m-win_amd64.whl
\n    torch-1.5.0+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.5.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.5.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.5.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.5.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.5.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.5.0-cp27-none-macosx_10_7_x86_64.whl
\n    torch-1.5.0-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.5.0-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.5.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.5.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.5.1+cpu-cp35-cp35m-linux_x86_64.whl
\n    torch-1.5.1+cpu-cp35-cp35m-win_amd64.whl
\n    torch-1.5.1+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.5.1+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.5.1+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.5.1+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.5.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.5.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.5.1-cp35-none-macosx_10_6_x86_64.whl
\n    torch-1.5.1-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.5.1-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.5.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.6.0+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.6.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.6.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.6.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.6.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.6.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.6.0-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.6.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.6.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.7.0+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.7.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.7.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.7.0-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.7.0-cp36-none-macosx_11_0_x86_64.whl
\n    torch-1.7.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.7.0-cp37-none-macosx_11_0_x86_64.whl
\n    torch-1.7.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.7.0-cp38-none-macosx_11_0_x86_64.whl
\n    torch-1.7.1+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.1+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.7.1+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.1+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.7.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.7.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.7.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.7.1-cp36-cp36m-linux_aarch64.whl
\n    torch-1.7.1-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.7.1-cp37-cp37m-linux_aarch64.whl
\n    torch-1.7.1-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.7.1-cp38-cp38-linux_aarch64.whl
\n    torch-1.7.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.7.1-cp39-cp39-linux_aarch64.whl
\n    torch-1.7.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.8.0+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.8.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.8.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.8.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.8.0-cp36-cp36m-manylinux2014_aarch64.whl
\n    torch-1.8.0-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.8.0-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.8.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.8.0-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.8.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.8.0-cp38-none-macosx_11_1_arm64.whl
\n    torch-1.8.0-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.8.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.8.1+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.1+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.8.1+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.1+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.8.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.8.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.8.1-cp36-cp36m-manylinux2014_aarch64.whl
\n    torch-1.8.1-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.8.1-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.8.1-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.8.1-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.8.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.8.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.8.1-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.8.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.9.0+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.0+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.9.0+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.0+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.9.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.9.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.9.0-cp36-cp36m-linux_aarch64.whl
\n    torch-1.9.0-cp36-cp36m-manylinux2014_aarch64.whl
\n    torch-1.9.0-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.9.0-cp37-cp37m-linux_aarch64.whl
\n    torch-1.9.0-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.9.0-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.9.0-cp38-cp38-linux_aarch64.whl
\n    torch-1.9.0-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.9.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.9.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.9.0-cp39-cp39-linux_aarch64.whl
\n    torch-1.9.0-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.9.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.9.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-1.9.1+cpu-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.1+cpu-cp36-cp36m-win_amd64.whl
\n    torch-1.9.1+cpu-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.1+cpu-cp37-cp37m-win_amd64.whl
\n    torch-1.9.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-1.9.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-1.9.1-cp36-none-macosx_10_9_x86_64.whl
\n    torch-1.9.1-cp37-none-macosx_10_9_x86_64.whl
\n    torch-1.9.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-1.9.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-1.9.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-1.9.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.0.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.0.0+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.0+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.0.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.0.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.0.0-cp310-none-macosx_10_9_x86_64.whl
\n    torch-2.0.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.0.0-cp311-none-macosx_10_9_x86_64.whl
\n    torch-2.0.0-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.0.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-2.0.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.0.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-2.0.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.0.1+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.0.1+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.1+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.0.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.0.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.0.1-cp310-none-macosx_10_9_x86_64.whl
\n    torch-2.0.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.0.1-cp311-none-macosx_10_9_x86_64.whl
\n    torch-2.0.1-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.0.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-2.0.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.0.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-2.0.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.1.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.1.0+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.0+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.1.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.1.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.1.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.0-cp310-none-macosx_10_9_x86_64.whl
\n    torch-2.1.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.0-cp311-none-macosx_10_9_x86_64.whl
\n    torch-2.1.0-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.1.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-2.1.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.1.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-2.1.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.1.1+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.1.1+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.1+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.1.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.1.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.1.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.1-cp310-none-macosx_10_9_x86_64.whl
\n    torch-2.1.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.1.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.1-cp311-none-macosx_10_9_x86_64.whl
\n    torch-2.1.1-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.1.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-2.1.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.1.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-2.1.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.1.2+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.2+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.1.2+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.2+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.1.2+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.2+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.1.2+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.2+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.1.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.2-cp310-none-macosx_10_9_x86_64.whl
\n    torch-2.1.2-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.1.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.2-cp311-none-macosx_10_9_x86_64.whl
\n    torch-2.1.2-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.1.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.2-cp38-none-macosx_10_9_x86_64.whl
\n    torch-2.1.2-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.1.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.1.2-cp39-none-macosx_10_9_x86_64.whl
\n    torch-2.1.2-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.2.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.2.0+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.0+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.2.0+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.0+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.2.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.2.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.2.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.0-cp310-none-macosx_10_9_x86_64.whl
\n    torch-2.2.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.2.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.0-cp311-none-macosx_10_9_x86_64.whl
\n    torch-2.2.0-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.2.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.0-cp312-none-macosx_10_9_x86_64.whl
\n    torch-2.2.0-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.2.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.0-cp38-none-macosx_10_9_x86_64.whl
\n    torch-2.2.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.2.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.0-cp39-none-macosx_10_9_x86_64.whl
\n    torch-2.2.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.2.1+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.2.1+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.1+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.2.1+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.1+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.2.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.2.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.2.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.1-cp310-none-macosx_10_9_x86_64.whl
\n    torch-2.2.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.2.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.1-cp311-none-macosx_10_9_x86_64.whl
\n    torch-2.2.1-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.2.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.1-cp312-none-macosx_10_9_x86_64.whl
\n    torch-2.2.1-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.2.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.1-cp38-none-macosx_10_9_x86_64.whl
\n    torch-2.2.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.2.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.1-cp39-none-macosx_10_9_x86_64.whl
\n    torch-2.2.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.2.2+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.2+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.2.2+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.2+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.2.2+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.2+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.2.2+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.2+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.2.2+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.2+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.2.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.2-cp310-none-macosx_10_9_x86_64.whl
\n    torch-2.2.2-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.2.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.2-cp311-none-macosx_10_9_x86_64.whl
\n    torch-2.2.2-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.2.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.2-cp312-none-macosx_10_9_x86_64.whl
\n    torch-2.2.2-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.2.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.2-cp38-none-macosx_10_9_x86_64.whl
\n    torch-2.2.2-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.2.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.2.2-cp39-none-macosx_10_9_x86_64.whl
\n    torch-2.2.2-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.3.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.3.0+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.0+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.3.0+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.0+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.3.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.3.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.3.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.3.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.0-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.3.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.0-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.3.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.3.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.3.1+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.3.1+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.1+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.3.1+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.1+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.3.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.3.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.3.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.3.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.1-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.3.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.1-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.3.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.3.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.3.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.4.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.4.0+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.0+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.4.0+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.0+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.4.0+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.0+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.4.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.4.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.4.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.0-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.4.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.0-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.4.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.0-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.4.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.4.1+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.4.1+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.1+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.4.1+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.1+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.4.1+cpu-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.1+cpu-cp38-cp38-win_amd64.whl
\n    torch-2.4.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.4.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.4.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.1-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.1-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.4.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.1-cp38-none-macosx_11_0_arm64.whl
\n    torch-2.4.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.4.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.5.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.5.0+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.0+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.5.0+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.0+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.5.0+cpu-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.5.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.5.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.5.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.5.0-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.5.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.5.0-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.5.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.5.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.5.1+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.5.1+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.1+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.5.1+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.1+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.5.1+cpu-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.1+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.5.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.5.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.5.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.5.1-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.5.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.5.1-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.5.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
\n    torch-2.5.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.6.0+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.6.0+cpu-cp310-cp310-manylinux_2_28_aarch64.whl
\n    torch-2.6.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.6.0+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.6.0+cpu-cp311-cp311-manylinux_2_28_aarch64.whl
\n    torch-2.6.0+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.6.0+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.6.0+cpu-cp312-cp312-manylinux_2_28_aarch64.whl
\n    torch-2.6.0+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.6.0+cpu-cp313-cp313-linux_x86_64.whl
\n    torch-2.6.0+cpu-cp313-cp313-manylinux_2_28_aarch64.whl
\n    torch-2.6.0+cpu-cp313-cp313-win_amd64.whl
\n    torch-2.6.0+cpu-cp313-cp313t-linux_x86_64.whl
\n    torch-2.6.0+cpu-cp313-cp313t-manylinux_2_28_aarch64.whl
\n    torch-2.6.0+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.6.0+cpu-cp39-cp39-manylinux_2_28_aarch64.whl
\n    torch-2.6.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.6.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.6.0-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.6.0-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.6.0-cp313-none-macosx_11_0_arm64.whl
\n    torch-2.6.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.7.0+cpu-cp310-cp310-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cpu-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.7.0+cpu-cp311-cp311-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cpu-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.7.0+cpu-cp312-cp312-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cpu-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.7.0+cpu-cp313-cp313-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cpu-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cpu-cp313-cp313-win_amd64.whl
\n    torch-2.7.0+cpu-cp313-cp313t-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cpu-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cpu-cp313-cp313t-win_amd64.whl
\n    torch-2.7.0+cpu-cp39-cp39-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cpu-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.7.0-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.7.0-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.7.0-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.7.0-cp313-cp313t-macosx_14_0_arm64.whl
\n    torch-2.7.0-cp313-none-macosx_11_0_arm64.whl
\n    torch-2.7.0-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.7.1+cpu-cp310-cp310-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cpu-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cpu-cp310-cp310-win_amd64.whl
\n    torch-2.7.1+cpu-cp311-cp311-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cpu-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cpu-cp311-cp311-win_amd64.whl
\n    torch-2.7.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cpu-cp312-cp312-win_amd64.whl
\n    torch-2.7.1+cpu-cp313-cp313-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cpu-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cpu-cp313-cp313-win_amd64.whl
\n    torch-2.7.1+cpu-cp313-cp313t-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cpu-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cpu-cp313-cp313t-win_amd64.whl
\n    torch-2.7.1+cpu-cp39-cp39-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cpu-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cpu-cp39-cp39-win_amd64.whl
\n    torch-2.7.1-cp310-none-macosx_11_0_arm64.whl
\n    torch-2.7.1-cp311-none-macosx_11_0_arm64.whl
\n    torch-2.7.1-cp312-none-macosx_11_0_arm64.whl
\n    torch-2.7.1-cp313-cp313t-macosx_14_0_arm64.whl
\n    torch-2.7.1-cp313-none-macosx_11_0_arm64.whl
\n    torch-2.7.1-cp39-none-macosx_11_0_arm64.whl
\n    torch-2.6.0.dev20240914+cpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.6.0.dev20240914+cpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.6.0.dev20240914+cpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.6.0.dev20240914+cpu-cp39-cp39-linux_x86_64.whl
\n    torch-1.0.0-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.0-cp35-cp35m-win_amd64.whl
\n    torch-1.0.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.0-cp36-cp36m-win_amd64.whl
\n    torch-1.0.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.0-cp37-cp37m-win_amd64.whl
\n    torch-1.0.1-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.1-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.1-cp35-cp35m-win_amd64.whl
\n    torch-1.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.1-cp36-cp36m-win_amd64.whl
\n    torch-1.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.1-cp37-cp37m-win_amd64.whl
\n    torch-1.0.1.post2-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.1.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.1.0-cp27-cp27m-linux_x86_64.whl
\n    torch-1.1.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.1.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.1.0-cp35-cp35m-win_amd64.whl
\n    torch-1.1.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.1.0-cp36-cp36m-win_amd64.whl
\n    torch-1.1.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.1.0-cp37-cp37m-win_amd64.whl
\n    torch-1.2.0-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-1.2.0-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-1.2.0-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-1.2.0-cp35-cp35m-win_amd64.whl
\n    torch-1.2.0-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-1.2.0-cp36-cp36m-win_amd64.whl
\n    torch-1.2.0-cp37-cp37m-manylinux1_x86_64.whl
\n    torch-1.2.0-cp37-cp37m-win_amd64.whl
\n    torch-1.3.0+cu100-cp27-cp27m-linux_x86_64.whl
\n    torch-1.3.0+cu100-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.3.0+cu100-cp35-cp35m-linux_x86_64.whl
\n    torch-1.3.0+cu100-cp36-cp36m-linux_x86_64.whl
\n    torch-1.3.0+cu100-cp37-cp37m-linux_x86_64.whl
\n    torch-1.3.1+cu100-cp27-cp27m-linux_x86_64.whl
\n    torch-1.3.1+cu100-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.3.1+cu100-cp35-cp35m-linux_x86_64.whl
\n    torch-1.3.1+cu100-cp36-cp36m-linux_x86_64.whl
\n    torch-1.3.1+cu100-cp37-cp37m-linux_x86_64.whl
\n    torch-1.4.0+cu100-cp27-cp27m-linux_x86_64.whl
\n    torch-1.4.0+cu100-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.4.0+cu100-cp35-cp35m-linux_x86_64.whl
\n    torch-1.4.0+cu100-cp36-cp36m-linux_x86_64.whl
\n    torch-1.4.0+cu100-cp37-cp37m-linux_x86_64.whl
\n    torch-1.4.0+cu100-cp38-cp38-linux_x86_64.whl
\n    torch-1.3.0-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-1.3.0-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-1.3.0-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-1.3.0-cp35-cp35m-win_amd64.whl
\n    torch-1.3.0-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-1.3.0-cp36-cp36m-win_amd64.whl
\n    torch-1.3.0-cp37-cp37m-manylinux1_x86_64.whl
\n    torch-1.3.0-cp37-cp37m-win_amd64.whl
\n    torch-1.3.1-cp27-cp27m-linux_x86_64.whl
\n    torch-1.3.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.3.1-cp35-cp35m-linux_x86_64.whl
\n    torch-1.3.1-cp35-cp35m-win_amd64.whl
\n    torch-1.3.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.3.1-cp36-cp36m-win_amd64.whl
\n    torch-1.3.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.3.1-cp37-cp37m-win_amd64.whl
\n    torch-1.4.0-cp27-cp27m-linux_x86_64.whl
\n    torch-1.4.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.4.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.4.0-cp35-cp35m-win_amd64.whl
\n    torch-1.4.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.4.0-cp36-cp36m-win_amd64.whl
\n    torch-1.4.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.4.0-cp37-cp37m-win_amd64.whl
\n    torch-1.4.0-cp38-cp38-linux_x86_64.whl
\n    torch-1.4.0-cp38-cp38-win_amd64.whl
\n    torch-1.5.0+cu101-cp35-cp35m-linux_x86_64.whl
\n    torch-1.5.0+cu101-cp35-cp35m-win_amd64.whl
\n    torch-1.5.0+cu101-cp36-cp36m-linux_x86_64.whl
\n    torch-1.5.0+cu101-cp36-cp36m-win_amd64.whl
\n    torch-1.5.0+cu101-cp37-cp37m-linux_x86_64.whl
\n    torch-1.5.0+cu101-cp37-cp37m-win_amd64.whl
\n    torch-1.5.0+cu101-cp38-cp38-linux_x86_64.whl
\n    torch-1.5.0+cu101-cp38-cp38-win_amd64.whl
\n    torch-1.5.1+cu101-cp35-cp35m-linux_x86_64.whl
\n    torch-1.5.1+cu101-cp35-cp35m-win_amd64.whl
\n    torch-1.5.1+cu101-cp36-cp36m-linux_x86_64.whl
\n    torch-1.5.1+cu101-cp36-cp36m-win_amd64.whl
\n    torch-1.5.1+cu101-cp37-cp37m-linux_x86_64.whl
\n    torch-1.5.1+cu101-cp37-cp37m-win_amd64.whl
\n    torch-1.5.1+cu101-cp38-cp38-linux_x86_64.whl
\n    torch-1.5.1+cu101-cp38-cp38-win_amd64.whl
\n    torch-1.6.0+cu101-cp36-cp36m-linux_x86_64.whl
\n    torch-1.6.0+cu101-cp36-cp36m-win_amd64.whl
\n    torch-1.6.0+cu101-cp37-cp37m-linux_x86_64.whl
\n    torch-1.6.0+cu101-cp37-cp37m-win_amd64.whl
\n    torch-1.6.0+cu101-cp38-cp38-linux_x86_64.whl
\n    torch-1.6.0+cu101-cp38-cp38-win_amd64.whl
\n    torch-1.7.0+cu101-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.0+cu101-cp36-cp36m-win_amd64.whl
\n    torch-1.7.0+cu101-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.0+cu101-cp37-cp37m-win_amd64.whl
\n    torch-1.7.0+cu101-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.0+cu101-cp38-cp38-win_amd64.whl
\n    torch-1.7.1+cu101-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.1+cu101-cp36-cp36m-win_amd64.whl
\n    torch-1.7.1+cu101-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.1+cu101-cp37-cp37m-win_amd64.whl
\n    torch-1.7.1+cu101-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.1+cu101-cp38-cp38-win_amd64.whl
\n    torch-1.7.1+cu101-cp39-cp39-linux_x86_64.whl
\n    torch-1.7.1+cu101-cp39-cp39-win_amd64.whl
\n    torch-1.8.0+cu101-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.0+cu101-cp36-cp36m-win_amd64.whl
\n    torch-1.8.0+cu101-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.0+cu101-cp37-cp37m-win_amd64.whl
\n    torch-1.8.0+cu101-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.0+cu101-cp38-cp38-win_amd64.whl
\n    torch-1.8.0+cu101-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.0+cu101-cp39-cp39-win_amd64.whl
\n    torch-1.8.1+cu101-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.1+cu101-cp36-cp36m-win_amd64.whl
\n    torch-1.8.1+cu101-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.1+cu101-cp37-cp37m-win_amd64.whl
\n    torch-1.8.1+cu101-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.1+cu101-cp38-cp38-win_amd64.whl
\n    torch-1.8.1+cu101-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.1+cu101-cp39-cp39-win_amd64.whl
\n    torch-1.10.0+cu102-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.0+cu102-cp36-cp36m-win_amd64.whl
\n    torch-1.10.0+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.0+cu102-cp37-cp37m-win_amd64.whl
\n    torch-1.10.0+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.0+cu102-cp38-cp38-win_amd64.whl
\n    torch-1.10.0+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.0+cu102-cp39-cp39-win_amd64.whl
\n    torch-1.10.1+cu102-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.1+cu102-cp36-cp36m-win_amd64.whl
\n    torch-1.10.1+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.1+cu102-cp37-cp37m-win_amd64.whl
\n    torch-1.10.1+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.1+cu102-cp38-cp38-win_amd64.whl
\n    torch-1.10.1+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.1+cu102-cp39-cp39-win_amd64.whl
\n    torch-1.10.2+cu102-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.2+cu102-cp36-cp36m-win_amd64.whl
\n    torch-1.10.2+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.2+cu102-cp37-cp37m-win_amd64.whl
\n    torch-1.10.2+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.2+cu102-cp38-cp38-win_amd64.whl
\n    torch-1.10.2+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.2+cu102-cp39-cp39-win_amd64.whl
\n    torch-1.11.0+cu102-cp310-cp310-linux_x86_64.whl
\n    torch-1.11.0+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.11.0+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.11.0+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.0+cu102-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.0+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.0+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.0+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.1+cu102-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.1+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.1+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.1+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.5.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.5.0-cp35-cp35m-win_amd64.whl
\n    torch-1.5.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.5.0-cp36-cp36m-win_amd64.whl
\n    torch-1.5.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.5.0-cp37-cp37m-win_amd64.whl
\n    torch-1.5.0-cp38-cp38-linux_x86_64.whl
\n    torch-1.5.0-cp38-cp38-win_amd64.whl
\n    torch-1.5.1-cp35-cp35m-linux_x86_64.whl
\n    torch-1.5.1-cp35-cp35m-win_amd64.whl
\n    torch-1.5.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.5.1-cp36-cp36m-win_amd64.whl
\n    torch-1.5.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.5.1-cp37-cp37m-win_amd64.whl
\n    torch-1.5.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.5.1-cp38-cp38-win_amd64.whl
\n    torch-1.6.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.6.0-cp36-cp36m-win_amd64.whl
\n    torch-1.6.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.6.0-cp37-cp37m-win_amd64.whl
\n    torch-1.6.0-cp38-cp38-linux_x86_64.whl
\n    torch-1.6.0-cp38-cp38-win_amd64.whl
\n    torch-1.7.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.0-cp36-cp36m-win_amd64.whl
\n    torch-1.7.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.0-cp37-cp37m-win_amd64.whl
\n    torch-1.7.0-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.0-cp38-cp38-win_amd64.whl
\n    torch-1.7.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.1-cp36-cp36m-win_amd64.whl
\n    torch-1.7.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.1-cp37-cp37m-win_amd64.whl
\n    torch-1.7.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.1-cp38-cp38-win_amd64.whl
\n    torch-1.7.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.7.1-cp39-cp39-win_amd64.whl
\n    torch-1.8.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.0-cp36-cp36m-win_amd64.whl
\n    torch-1.8.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.0-cp37-cp37m-win_amd64.whl
\n    torch-1.8.0-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.0-cp38-cp38-win_amd64.whl
\n    torch-1.8.0-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.0-cp39-cp39-win_amd64.whl
\n    torch-1.8.1+cu102-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.1+cu102-cp36-cp36m-win_amd64.whl
\n    torch-1.8.1+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.1+cu102-cp37-cp37m-win_amd64.whl
\n    torch-1.8.1+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.1+cu102-cp38-cp38-win_amd64.whl
\n    torch-1.8.1+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.1+cu102-cp39-cp39-win_amd64.whl
\n    torch-1.9.0+cu102-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.0+cu102-cp36-cp36m-win_amd64.whl
\n    torch-1.9.0+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.0+cu102-cp37-cp37m-win_amd64.whl
\n    torch-1.9.0+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.0+cu102-cp38-cp38-win_amd64.whl
\n    torch-1.9.0+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.0+cu102-cp39-cp39-win_amd64.whl
\n    torch-1.9.1+cu102-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.1+cu102-cp36-cp36m-win_amd64.whl
\n    torch-1.9.1+cu102-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.1+cu102-cp37-cp37m-win_amd64.whl
\n    torch-1.9.1+cu102-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.1+cu102-cp38-cp38-win_amd64.whl
\n    torch-1.9.1+cu102-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.1+cu102-cp39-cp39-win_amd64.whl
\n    torch-1.7.0+cu110-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.0+cu110-cp36-cp36m-win_amd64.whl
\n    torch-1.7.0+cu110-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.0+cu110-cp37-cp37m-win_amd64.whl
\n    torch-1.7.0+cu110-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.0+cu110-cp38-cp38-win_amd64.whl
\n    torch-1.7.1+cu110-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.1+cu110-cp36-cp36m-win_amd64.whl
\n    torch-1.7.1+cu110-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.1+cu110-cp37-cp37m-win_amd64.whl
\n    torch-1.7.1+cu110-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.1+cu110-cp38-cp38-win_amd64.whl
\n    torch-1.7.1+cu110-cp39-cp39-linux_x86_64.whl
\n    torch-1.7.1+cu110-cp39-cp39-win_amd64.whl
\n    torch-1.10.0+cu111-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.0+cu111-cp36-cp36m-win_amd64.whl
\n    torch-1.10.0+cu111-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.0+cu111-cp37-cp37m-win_amd64.whl
\n    torch-1.10.0+cu111-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.0+cu111-cp38-cp38-win_amd64.whl
\n    torch-1.10.0+cu111-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.0+cu111-cp39-cp39-win_amd64.whl
\n    torch-1.10.1+cu111-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.1+cu111-cp36-cp36m-win_amd64.whl
\n    torch-1.10.1+cu111-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.1+cu111-cp37-cp37m-win_amd64.whl
\n    torch-1.10.1+cu111-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.1+cu111-cp38-cp38-win_amd64.whl
\n    torch-1.10.1+cu111-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.1+cu111-cp39-cp39-win_amd64.whl
\n    torch-1.10.2+cu111-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.2+cu111-cp36-cp36m-win_amd64.whl
\n    torch-1.10.2+cu111-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.2+cu111-cp37-cp37m-win_amd64.whl
\n    torch-1.10.2+cu111-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.2+cu111-cp38-cp38-win_amd64.whl
\n    torch-1.10.2+cu111-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.2+cu111-cp39-cp39-win_amd64.whl
\n    torch-1.8.0+cu111-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.0+cu111-cp36-cp36m-win_amd64.whl
\n    torch-1.8.0+cu111-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.0+cu111-cp37-cp37m-win_amd64.whl
\n    torch-1.8.0+cu111-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.0+cu111-cp38-cp38-win_amd64.whl
\n    torch-1.8.0+cu111-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.0+cu111-cp39-cp39-win_amd64.whl
\n    torch-1.8.1+cu111-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.1+cu111-cp36-cp36m-win_amd64.whl
\n    torch-1.8.1+cu111-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.1+cu111-cp37-cp37m-win_amd64.whl
\n    torch-1.8.1+cu111-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.1+cu111-cp38-cp38-win_amd64.whl
\n    torch-1.8.1+cu111-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.1+cu111-cp39-cp39-win_amd64.whl
\n    torch-1.9.0+cu111-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.0+cu111-cp36-cp36m-win_amd64.whl
\n    torch-1.9.0+cu111-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.0+cu111-cp37-cp37m-win_amd64.whl
\n    torch-1.9.0+cu111-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.0+cu111-cp38-cp38-win_amd64.whl
\n    torch-1.9.0+cu111-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.0+cu111-cp39-cp39-win_amd64.whl
\n    torch-1.9.1+cu111-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.1+cu111-cp36-cp36m-win_amd64.whl
\n    torch-1.9.1+cu111-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.1+cu111-cp37-cp37m-win_amd64.whl
\n    torch-1.9.1+cu111-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.1+cu111-cp38-cp38-win_amd64.whl
\n    torch-1.9.1+cu111-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.1+cu111-cp39-cp39-win_amd64.whl
\n    torch-1.10.0+cu113-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.0+cu113-cp36-cp36m-win_amd64.whl
\n    torch-1.10.0+cu113-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.0+cu113-cp37-cp37m-win_amd64.whl
\n    torch-1.10.0+cu113-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.0+cu113-cp38-cp38-win_amd64.whl
\n    torch-1.10.0+cu113-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.0+cu113-cp39-cp39-win_amd64.whl
\n    torch-1.10.1+cu113-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.1+cu113-cp36-cp36m-win_amd64.whl
\n    torch-1.10.1+cu113-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.1+cu113-cp37-cp37m-win_amd64.whl
\n    torch-1.10.1+cu113-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.1+cu113-cp38-cp38-win_amd64.whl
\n    torch-1.10.1+cu113-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.1+cu113-cp39-cp39-win_amd64.whl
\n    torch-1.10.2+cu113-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.2+cu113-cp36-cp36m-win_amd64.whl
\n    torch-1.10.2+cu113-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.2+cu113-cp37-cp37m-win_amd64.whl
\n    torch-1.10.2+cu113-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.2+cu113-cp38-cp38-win_amd64.whl
\n    torch-1.10.2+cu113-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.2+cu113-cp39-cp39-win_amd64.whl
\n    torch-1.11.0+cu113-cp310-cp310-linux_x86_64.whl
\n    torch-1.11.0+cu113-cp310-cp310-win_amd64.whl
\n    torch-1.11.0+cu113-cp37-cp37m-linux_x86_64.whl
\n    torch-1.11.0+cu113-cp37-cp37m-win_amd64.whl
\n    torch-1.11.0+cu113-cp38-cp38-linux_x86_64.whl
\n    torch-1.11.0+cu113-cp38-cp38-win_amd64.whl
\n    torch-1.11.0+cu113-cp39-cp39-linux_x86_64.whl
\n    torch-1.11.0+cu113-cp39-cp39-win_amd64.whl
\n    torch-1.12.0+cu113-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.0+cu113-cp310-cp310-win_amd64.whl
\n    torch-1.12.0+cu113-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.0+cu113-cp37-cp37m-win_amd64.whl
\n    torch-1.12.0+cu113-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.0+cu113-cp38-cp38-win_amd64.whl
\n    torch-1.12.0+cu113-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.0+cu113-cp39-cp39-win_amd64.whl
\n    torch-1.12.1+cu113-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.1+cu113-cp310-cp310-win_amd64.whl
\n    torch-1.12.1+cu113-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.1+cu113-cp37-cp37m-win_amd64.whl
\n    torch-1.12.1+cu113-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.1+cu113-cp38-cp38-win_amd64.whl
\n    torch-1.12.1+cu113-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.1+cu113-cp39-cp39-win_amd64.whl
\n    torch-1.11.0+cu115-cp310-cp310-linux_x86_64.whl
\n    torch-1.11.0+cu115-cp310-cp310-win_amd64.whl
\n    torch-1.11.0+cu115-cp37-cp37m-linux_x86_64.whl
\n    torch-1.11.0+cu115-cp37-cp37m-win_amd64.whl
\n    torch-1.11.0+cu115-cp38-cp38-linux_x86_64.whl
\n    torch-1.11.0+cu115-cp38-cp38-win_amd64.whl
\n    torch-1.11.0+cu115-cp39-cp39-linux_x86_64.whl
\n    torch-1.11.0+cu115-cp39-cp39-win_amd64.whl
\n    torch-1.12.0+cu116-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.0+cu116-cp310-cp310-win_amd64.whl
\n    torch-1.12.0+cu116-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.0+cu116-cp37-cp37m-win_amd64.whl
\n    torch-1.12.0+cu116-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.0+cu116-cp38-cp38-win_amd64.whl
\n    torch-1.12.0+cu116-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.0+cu116-cp39-cp39-win_amd64.whl
\n    torch-1.12.1+cu116-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.1+cu116-cp310-cp310-win_amd64.whl
\n    torch-1.12.1+cu116-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.1+cu116-cp37-cp37m-win_amd64.whl
\n    torch-1.12.1+cu116-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.1+cu116-cp38-cp38-win_amd64.whl
\n    torch-1.12.1+cu116-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.1+cu116-cp39-cp39-win_amd64.whl
\n    torch-1.13.0+cu116-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.0+cu116-cp310-cp310-win_amd64.whl
\n    torch-1.13.0+cu116-cp311-cp311-linux_x86_64.whl
\n    torch-1.13.0+cu116-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.0+cu116-cp37-cp37m-win_amd64.whl
\n    torch-1.13.0+cu116-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.0+cu116-cp38-cp38-win_amd64.whl
\n    torch-1.13.0+cu116-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.0+cu116-cp39-cp39-win_amd64.whl
\n    torch-1.13.1+cu116-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.1+cu116-cp310-cp310-win_amd64.whl
\n    torch-1.13.1+cu116-cp311-cp311-linux_x86_64.whl
\n    torch-1.13.1+cu116-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.1+cu116-cp37-cp37m-win_amd64.whl
\n    torch-1.13.1+cu116-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.1+cu116-cp38-cp38-win_amd64.whl
\n    torch-1.13.1+cu116-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.1+cu116-cp39-cp39-win_amd64.whl
\n    torch-1.13.0+cu117-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.0+cu117-cp310-cp310-win_amd64.whl
\n    torch-1.13.0+cu117-cp311-cp311-linux_x86_64.whl
\n    torch-1.13.0+cu117-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.0+cu117-cp37-cp37m-win_amd64.whl
\n    torch-1.13.0+cu117-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.0+cu117-cp38-cp38-win_amd64.whl
\n    torch-1.13.0+cu117-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.0+cu117-cp39-cp39-win_amd64.whl
\n    torch-1.13.1+cu117-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.1+cu117-cp310-cp310-win_amd64.whl
\n    torch-1.13.1+cu117-cp311-cp311-linux_x86_64.whl
\n    torch-1.13.1+cu117-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.1+cu117-cp37-cp37m-win_amd64.whl
\n    torch-1.13.1+cu117-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.1+cu117-cp38-cp38-win_amd64.whl
\n    torch-1.13.1+cu117-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.1+cu117-cp39-cp39-win_amd64.whl
\n    torch-2.0.0+cu117-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.0+cu117-cp310-cp310-win_amd64.whl
\n    torch-2.0.0+cu117-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.0+cu117-cp311-cp311-win_amd64.whl
\n    torch-2.0.0+cu117-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.0+cu117-cp38-cp38-win_amd64.whl
\n    torch-2.0.0+cu117-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.0+cu117-cp39-cp39-win_amd64.whl
\n    torch-2.0.1+cu117-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.1+cu117-cp310-cp310-win_amd64.whl
\n    torch-2.0.1+cu117-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.1+cu117-cp311-cp311-win_amd64.whl
\n    torch-2.0.1+cu117-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.1+cu117-cp38-cp38-win_amd64.whl
\n    torch-2.0.1+cu117-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.1+cu117-cp39-cp39-win_amd64.whl
\n    torch-1.13.0+cu117.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.0+cu117.with.pypi.cudnn-cp311-cp311-linux_x86_64.whl
\n    torch-1.13.0+cu117.with.pypi.cudnn-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.0+cu117.with.pypi.cudnn-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.0+cu117.with.pypi.cudnn-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.1+cu117.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.1+cu117.with.pypi.cudnn-cp311-cp311-linux_x86_64.whl
\n    torch-1.13.1+cu117.with.pypi.cudnn-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.1+cu117.with.pypi.cudnn-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.1+cu117.with.pypi.cudnn-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.0+cu117.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.0+cu117.with.pypi.cudnn-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.0+cu117.with.pypi.cudnn-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.0+cu117.with.pypi.cudnn-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.1+cu117.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.1+cu117.with.pypi.cudnn-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.1+cu117.with.pypi.cudnn-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.1+cu117.with.pypi.cudnn-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.0+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.0+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.0.0+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.0+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.0.0+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.0+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.0.0+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.0+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.0.1+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.1+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.0.1+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.1+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.0.1+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.1+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.0.1+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.1+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.1.0+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.0+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.1.0+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.0+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.1.0+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.0+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.1.0+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.0+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.1.1+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.1+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.1.1+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.1+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.1.1+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.1+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.1.1+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.1+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.1.2+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.2+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.1.2+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.2+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.1.2+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.2+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.1.2+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.2+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.2.0+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.0+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.2.0+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.0+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.2.0+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.0+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.2.0+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.0+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.2.0+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.0+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.2.1+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.1+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.2.1+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.1+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.2.1+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.1+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.2.1+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.1+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.2.1+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.1+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.2.2+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.2+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.2.2+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.2+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.2.2+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.2+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.2.2+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.2+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.2.2+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.2+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.3.0+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.0+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.3.0+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.0+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.3.0+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.0+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.3.0+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.0+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.3.0+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.0+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.3.1+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.1+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.3.1+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.1+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.3.1+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.1+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.3.1+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.1+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.3.1+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.1+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.4.0+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.0+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.4.0+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.0+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.4.0+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.0+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.4.0+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.0+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.4.0+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.0+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.4.1+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.1+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.4.1+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.1+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.4.1+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.1+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.4.1+cu118-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.1+cu118-cp38-cp38-win_amd64.whl
\n    torch-2.4.1+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.1+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.5.0+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.0+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.5.0+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.0+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.5.0+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.0+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.5.0+cu118-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.0+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.0+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.5.1+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.1+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.5.1+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.1+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.5.1+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.1+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.5.1+cu118-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.1+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.1+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.6.0+cu118-cp310-cp310-linux_x86_64.whl
\n    torch-2.6.0+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.6.0+cu118-cp311-cp311-linux_x86_64.whl
\n    torch-2.6.0+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.6.0+cu118-cp312-cp312-linux_x86_64.whl
\n    torch-2.6.0+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.6.0+cu118-cp313-cp313-linux_x86_64.whl
\n    torch-2.6.0+cu118-cp313-cp313-win_amd64.whl
\n    torch-2.6.0+cu118-cp313-cp313t-linux_x86_64.whl
\n    torch-2.6.0+cu118-cp39-cp39-linux_x86_64.whl
\n    torch-2.6.0+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.7.0+cu118-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.7.0+cu118-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.7.0+cu118-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.7.0+cu118-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu118-cp313-cp313-win_amd64.whl
\n    torch-2.7.0+cu118-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu118-cp313-cp313t-win_amd64.whl
\n    torch-2.7.0+cu118-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.7.1+cu118-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu118-cp310-cp310-win_amd64.whl
\n    torch-2.7.1+cu118-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu118-cp311-cp311-win_amd64.whl
\n    torch-2.7.1+cu118-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu118-cp312-cp312-win_amd64.whl
\n    torch-2.7.1+cu118-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu118-cp313-cp313-win_amd64.whl
\n    torch-2.7.1+cu118-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu118-cp313-cp313t-win_amd64.whl
\n    torch-2.7.1+cu118-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu118-cp39-cp39-win_amd64.whl
\n    torch-2.1.0+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.0+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.1.0+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.0+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.1.0+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.0+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.1.0+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.0+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.1.1+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.1+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.1.1+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.1+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.1.1+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.1+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.1.1+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.1+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.1.2+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.2+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.1.2+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.2+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.1.2+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.2+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.1.2+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.2+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.2.0+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.0+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.2.0+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.0+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.2.0+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.0+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.2.0+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.0+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.2.0+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.0+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.2.1+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.1+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.2.1+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.1+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.2.1+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.1+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.2.1+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.1+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.2.1+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.1+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.2.2+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.2+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.2.2+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.2+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.2.2+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.2+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.2.2+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.2+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.2.2+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.2+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.3.0+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.0+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.3.0+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.0+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.3.0+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.0+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.3.0+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.0+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.3.0+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.0+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.3.1+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.1+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.3.1+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.1+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.3.1+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.1+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.3.1+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.1+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.3.1+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.1+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.4.0+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.0+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.4.0+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.0+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.4.0+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.0+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.4.0+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.0+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.4.0+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.0+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.4.1+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.1+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.4.1+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.1+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.4.1+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.1+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.4.1+cu121-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.1+cu121-cp38-cp38-win_amd64.whl
\n    torch-2.4.1+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.1+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.5.0+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.0+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.5.0+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.0+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.5.0+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.0+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.5.0+cu121-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.0+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.0+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.5.1+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.1+cu121-cp310-cp310-win_amd64.whl
\n    torch-2.5.1+cu121-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.1+cu121-cp311-cp311-win_amd64.whl
\n    torch-2.5.1+cu121-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.1+cu121-cp312-cp312-win_amd64.whl
\n    torch-2.5.1+cu121-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.1+cu121-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.1+cu121-cp39-cp39-win_amd64.whl
\n    torch-2.4.0+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.1+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.0+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.1+cu121-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.0+cu121.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.0+cu121.with.pypi.cudnn-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.0+cu121.with.pypi.cudnn-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.0+cu121.with.pypi.cudnn-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.1+cu121.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.1+cu121.with.pypi.cudnn-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.1+cu121.with.pypi.cudnn-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.1+cu121.with.pypi.cudnn-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.2+cu121.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.2+cu121.with.pypi.cudnn-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.2+cu121.with.pypi.cudnn-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.2+cu121.with.pypi.cudnn-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.0+cu124-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.0+cu124-cp310-cp310-win_amd64.whl
\n    torch-2.4.0+cu124-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.0+cu124-cp311-cp311-win_amd64.whl
\n    torch-2.4.0+cu124-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.0+cu124-cp312-cp312-win_amd64.whl
\n    torch-2.4.0+cu124-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.0+cu124-cp38-cp38-win_amd64.whl
\n    torch-2.4.0+cu124-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.0+cu124-cp39-cp39-win_amd64.whl
\n    torch-2.4.0-cp310-cp310-linux_aarch64.whl
\n    torch-2.4.0-cp311-cp311-linux_aarch64.whl
\n    torch-2.4.0-cp312-cp312-linux_aarch64.whl
\n    torch-2.4.0-cp38-cp38-linux_aarch64.whl
\n    torch-2.4.0-cp39-cp39-linux_aarch64.whl
\n    torch-2.4.1+cu124-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.1+cu124-cp310-cp310-win_amd64.whl
\n    torch-2.4.1+cu124-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.1+cu124-cp311-cp311-win_amd64.whl
\n    torch-2.4.1+cu124-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.1+cu124-cp312-cp312-win_amd64.whl
\n    torch-2.4.1+cu124-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.1+cu124-cp38-cp38-win_amd64.whl
\n    torch-2.4.1+cu124-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.1+cu124-cp39-cp39-win_amd64.whl
\n    torch-2.4.1-cp310-cp310-linux_aarch64.whl
\n    torch-2.4.1-cp311-cp311-linux_aarch64.whl
\n    torch-2.4.1-cp312-cp312-linux_aarch64.whl
\n    torch-2.4.1-cp38-cp38-linux_aarch64.whl
\n    torch-2.4.1-cp39-cp39-linux_aarch64.whl
\n    torch-2.5.0+cu124-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.0+cu124-cp310-cp310-win_amd64.whl
\n    torch-2.5.0+cu124-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.0+cu124-cp311-cp311-win_amd64.whl
\n    torch-2.5.0+cu124-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.0+cu124-cp312-cp312-win_amd64.whl
\n    torch-2.5.0+cu124-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.0+cu124-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.0+cu124-cp39-cp39-win_amd64.whl
\n    torch-2.5.0-cp310-cp310-linux_aarch64.whl
\n    torch-2.5.0-cp311-cp311-linux_aarch64.whl
\n    torch-2.5.0-cp312-cp312-linux_aarch64.whl
\n    torch-2.5.0-cp39-cp39-linux_aarch64.whl
\n    torch-2.5.1+cu124-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.1+cu124-cp310-cp310-win_amd64.whl
\n    torch-2.5.1+cu124-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.1+cu124-cp311-cp311-win_amd64.whl
\n    torch-2.5.1+cu124-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.1+cu124-cp312-cp312-win_amd64.whl
\n    torch-2.5.1+cu124-cp313-cp313-linux_x86_64.whl
\n    torch-2.5.1+cu124-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.1+cu124-cp39-cp39-win_amd64.whl
\n    torch-2.5.1-cp310-cp310-linux_aarch64.whl
\n    torch-2.5.1-cp311-cp311-linux_aarch64.whl
\n    torch-2.5.1-cp312-cp312-linux_aarch64.whl
\n    torch-2.5.1-cp39-cp39-linux_aarch64.whl
\n    torch-2.6.0+cu124-cp310-cp310-linux_x86_64.whl
\n    torch-2.6.0+cu124-cp310-cp310-win_amd64.whl
\n    torch-2.6.0+cu124-cp311-cp311-linux_x86_64.whl
\n    torch-2.6.0+cu124-cp311-cp311-win_amd64.whl
\n    torch-2.6.0+cu124-cp312-cp312-linux_x86_64.whl
\n    torch-2.6.0+cu124-cp312-cp312-win_amd64.whl
\n    torch-2.6.0+cu124-cp313-cp313-linux_x86_64.whl
\n    torch-2.6.0+cu124-cp313-cp313-win_amd64.whl
\n    torch-2.6.0+cu124-cp313-cp313t-linux_x86_64.whl
\n    torch-2.6.0+cu124-cp39-cp39-linux_x86_64.whl
\n    torch-2.6.0+cu124-cp39-cp39-win_amd64.whl
\n    torch-2.6.0+cu124-cp311-cp311-linux_x86_64.whl
\n    torch-2.6.0+cu126-cp310-cp310-linux_aarch64.whl
\n    torch-2.6.0+cu126-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+cu126-cp310-cp310-win_amd64.whl
\n    torch-2.6.0+cu126-cp311-cp311-linux_aarch64.whl
\n    torch-2.6.0+cu126-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+cu126-cp311-cp311-win_amd64.whl
\n    torch-2.6.0+cu126-cp312-cp312-linux_aarch64.whl
\n    torch-2.6.0+cu126-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+cu126-cp312-cp312-win_amd64.whl
\n    torch-2.6.0+cu126-cp313-cp313-linux_aarch64.whl
\n    torch-2.6.0+cu126-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+cu126-cp313-cp313-win_amd64.whl
\n    torch-2.6.0+cu126-cp313-cp313t-linux_aarch64.whl
\n    torch-2.6.0+cu126-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+cu126-cp39-cp39-linux_aarch64.whl
\n    torch-2.6.0+cu126-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+cu126-cp39-cp39-win_amd64.whl
\n    torch-2.7.0+cu126-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu126-cp310-cp310-win_amd64.whl
\n    torch-2.7.0+cu126-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu126-cp311-cp311-win_amd64.whl
\n    torch-2.7.0+cu126-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu126-cp312-cp312-win_amd64.whl
\n    torch-2.7.0+cu126-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu126-cp313-cp313-win_amd64.whl
\n    torch-2.7.0+cu126-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu126-cp313-cp313t-win_amd64.whl
\n    torch-2.7.0+cu126-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu126-cp39-cp39-win_amd64.whl
\n    torch-2.7.1+cu126-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu126-cp310-cp310-win_amd64.whl
\n    torch-2.7.1+cu126-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu126-cp311-cp311-win_amd64.whl
\n    torch-2.7.1+cu126-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu126-cp312-cp312-win_amd64.whl
\n    torch-2.7.1+cu126-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu126-cp313-cp313-win_amd64.whl
\n    torch-2.7.1+cu126-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu126-cp313-cp313t-win_amd64.whl
\n    torch-2.7.1+cu126-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu126-cp39-cp39-win_amd64.whl
\n    torch-2.7.0+cu126-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu126-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu128-cp310-cp310-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cu128-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu128-cp310-cp310-win_amd64.whl
\n    torch-2.7.0+cu128-cp311-cp311-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cu128-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu128-cp311-cp311-win_amd64.whl
\n    torch-2.7.0+cu128-cp312-cp312-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cu128-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu128-cp312-cp312-win_amd64.whl
\n    torch-2.7.0+cu128-cp313-cp313-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cu128-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu128-cp313-cp313-win_amd64.whl
\n    torch-2.7.0+cu128-cp313-cp313t-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cu128-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu128-cp313-cp313t-win_amd64.whl
\n    torch-2.7.0+cu128-cp39-cp39-manylinux_2_28_aarch64.whl
\n    torch-2.7.0+cu128-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+cu128-cp39-cp39-win_amd64.whl
\n    torch-2.7.1+cu128-cp310-cp310-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cu128-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu128-cp310-cp310-win_amd64.whl
\n    torch-2.7.1+cu128-cp311-cp311-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cu128-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu128-cp311-cp311-win_amd64.whl
\n    torch-2.7.1+cu128-cp312-cp312-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cu128-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu128-cp312-cp312-win_amd64.whl
\n    torch-2.7.1+cu128-cp313-cp313-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cu128-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu128-cp313-cp313-win_amd64.whl
\n    torch-2.7.1+cu128-cp313-cp313t-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cu128-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu128-cp313-cp313t-win_amd64.whl
\n    torch-2.7.1+cu128-cp39-cp39-manylinux_2_28_aarch64.whl
\n    torch-2.7.1+cu128-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+cu128-cp39-cp39-win_amd64.whl
\n    torch-0.1.10.post1-cp27-none-linux_x86_64.whl
\n    torch-0.1.10.post1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.10.post1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.10.post2-cp27-none-linux_x86_64.whl
\n    torch-0.1.10.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.10.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.11.post4-cp27-none-linux_x86_64.whl
\n    torch-0.1.11.post4-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.11.post4-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.11.post5-cp27-none-linux_x86_64.whl
\n    torch-0.1.11.post5-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.11.post5-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.12.post1-cp27-none-linux_x86_64.whl
\n    torch-0.1.12.post1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.12.post1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.12.post2-cp27-none-linux_x86_64.whl
\n    torch-0.1.12.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.12.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.6.post20-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.1.6.post20-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.6.post22-cp27-none-linux_x86_64.whl
\n    torch-0.1.6.post22-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.7.post2-cp27-none-linux_x86_64.whl
\n    torch-0.1.7.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.7.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.8.post1-cp27-none-linux_x86_64.whl
\n    torch-0.1.8.post1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.8.post1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.9.post1-cp27-none-linux_x86_64.whl
\n    torch-0.1.9.post1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.9.post1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.9.post2-cp27-none-linux_x86_64.whl
\n    torch-0.1.9.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.9.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.2.0.post1-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-0.2.0.post1-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-0.2.0.post1-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-0.2.0.post1-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-0.2.0.post2-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-0.2.0.post2-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-0.2.0.post2-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-0.2.0.post2-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-0.2.0.post3-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-0.2.0.post3-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-0.2.0.post3-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-0.2.0.post3-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-0.3.0-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post3-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post4-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.10.post1-cp27-none-linux_x86_64.whl
\n    torch-0.1.10.post1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.10.post1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.10.post2-cp27-none-linux_x86_64.whl
\n    torch-0.1.10.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.10.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.11.post4-cp27-none-linux_x86_64.whl
\n    torch-0.1.11.post4-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.11.post4-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.11.post5-cp27-none-linux_x86_64.whl
\n    torch-0.1.11.post5-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.11.post5-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.12.post1-cp27-none-linux_x86_64.whl
\n    torch-0.1.12.post1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.12.post1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.12.post2-cp27-none-linux_x86_64.whl
\n    torch-0.1.12.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.12.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.6.post20-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.1.6.post20-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.6.post22-cp27-none-linux_x86_64.whl
\n    torch-0.1.6.post22-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.7.post2-cp27-none-linux_x86_64.whl
\n    torch-0.1.7.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.7.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.8.post1-cp27-none-linux_x86_64.whl
\n    torch-0.1.8.post1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.8.post1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.9.post1-cp27-none-linux_x86_64.whl
\n    torch-0.1.9.post1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.9.post1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.1.9.post2-cp27-none-linux_x86_64.whl
\n    torch-0.1.9.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.9.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.2.0.post2-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-0.2.0.post2-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-0.2.0.post2-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-0.2.0.post2-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-0.2.0.post3-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-0.2.0.post3-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-0.2.0.post3-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-0.2.0.post3-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-0.3.0-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post3-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post4-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.1-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.0-cp27-cp27m-linux_x86_64.whl
\n    torch-0.4.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-linux_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-win_amd64.whl
\n    torch-0.4.0-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.0-cp36-cp36m-win_amd64.whl
\n    torch-0.4.1-cp27-cp27m-linux_x86_64.whl
\n    torch-0.4.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-win_amd64.whl
\n    torch-0.4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.1-cp36-cp36m-win_amd64.whl
\n    torch-0.4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-0.4.1-cp37-cp37m-win_amd64.whl
\n    torch-0.4.1.post2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.0-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.0-cp35-cp35m-win_amd64.whl
\n    torch-1.0.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.0-cp36-cp36m-win_amd64.whl
\n    torch-1.0.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.0-cp37-cp37m-win_amd64.whl
\n    torch-1.0.1-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.1-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.1-cp35-cp35m-win_amd64.whl
\n    torch-1.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.1-cp36-cp36m-win_amd64.whl
\n    torch-1.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.1-cp37-cp37m-win_amd64.whl
\n    torch-1.0.1.post2-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.1.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp37-cp37m-linux_x86_64.whl
\n    torch-0.3.0-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post3-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post3-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.0.post4-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl
\n    torch-0.3.1-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.0-cp27-cp27m-linux_x86_64.whl
\n    torch-0.4.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-linux_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-win_amd64.whl
\n    torch-0.4.0-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.0-cp36-cp36m-win_amd64.whl
\n    torch-0.4.1-cp27-cp27m-linux_x86_64.whl
\n    torch-0.4.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-win_amd64.whl
\n    torch-0.4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.1-cp36-cp36m-win_amd64.whl
\n    torch-0.4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-0.4.1-cp37-cp37m-win_amd64.whl
\n    torch-0.4.1.post2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.0-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.0-cp35-cp35m-win_amd64.whl
\n    torch-1.0.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.0-cp36-cp36m-win_amd64.whl
\n    torch-1.0.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.0-cp37-cp37m-win_amd64.whl
\n    torch-1.0.1-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.1-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.1-cp35-cp35m-win_amd64.whl
\n    torch-1.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.1-cp36-cp36m-win_amd64.whl
\n    torch-1.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.0.1-cp37-cp37m-win_amd64.whl
\n    torch-1.0.1.post2-cp27-cp27m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.0.1.post2-cp35-cp35m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.0.1.post2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.1.0-cp27-cp27m-linux_x86_64.whl
\n    torch-1.1.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.1.0-cp35-cp35m-linux_x86_64.whl
\n    torch-1.1.0-cp35-cp35m-win_amd64.whl
\n    torch-1.1.0-cp36-cp36m-linux_x86_64.whl
\n    torch-1.1.0-cp36-cp36m-win_amd64.whl
\n    torch-1.1.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.1.0-cp37-cp37m-win_amd64.whl
\n    torch-0.3.1-cp27-cp27m-linux_x86_64.whl
\n    torch-0.3.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.3.1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.3.1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.0-cp27-cp27m-linux_x86_64.whl
\n    torch-0.4.0-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-linux_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-win_amd64.whl
\n    torch-0.4.0-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.0-cp36-cp36m-win_amd64.whl
\n    torch-0.4.1-cp27-cp27m-linux_x86_64.whl
\n    torch-0.4.1-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-linux_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-win_amd64.whl
\n    torch-0.4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-0.4.1-cp36-cp36m-win_amd64.whl
\n    torch-0.4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-0.4.1-cp37-cp37m-win_amd64.whl
\n    torch-0.4.1.post2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.2.0+cu92-cp27-cp27m-manylinux1_x86_64.whl
\n    torch-1.2.0+cu92-cp27-cp27mu-manylinux1_x86_64.whl
\n    torch-1.2.0+cu92-cp35-cp35m-manylinux1_x86_64.whl
\n    torch-1.2.0+cu92-cp35-cp35m-win_amd64.whl
\n    torch-1.2.0+cu92-cp36-cp36m-manylinux1_x86_64.whl
\n    torch-1.2.0+cu92-cp36-cp36m-win_amd64.whl
\n    torch-1.2.0+cu92-cp37-cp37m-manylinux1_x86_64.whl
\n    torch-1.2.0+cu92-cp37-cp37m-win_amd64.whl
\n    torch-1.3.0+cu92-cp27-cp27m-linux_x86_64.whl
\n    torch-1.3.0+cu92-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.3.0+cu92-cp35-cp35m-linux_x86_64.whl
\n    torch-1.3.0+cu92-cp35-cp35m-win_amd64.whl
\n    torch-1.3.0+cu92-cp36-cp36m-linux_x86_64.whl
\n    torch-1.3.0+cu92-cp36-cp36m-win_amd64.whl
\n    torch-1.3.0+cu92-cp37-cp37m-linux_x86_64.whl
\n    torch-1.3.0+cu92-cp37-cp37m-win_amd64.whl
\n    torch-1.3.1+cu92-cp27-cp27m-linux_x86_64.whl
\n    torch-1.3.1+cu92-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.3.1+cu92-cp35-cp35m-linux_x86_64.whl
\n    torch-1.3.1+cu92-cp35-cp35m-win_amd64.whl
\n    torch-1.3.1+cu92-cp36-cp36m-linux_x86_64.whl
\n    torch-1.3.1+cu92-cp36-cp36m-win_amd64.whl
\n    torch-1.3.1+cu92-cp37-cp37m-linux_x86_64.whl
\n    torch-1.3.1+cu92-cp37-cp37m-win_amd64.whl
\n    torch-1.4.0+cu92-cp27-cp27m-linux_x86_64.whl
\n    torch-1.4.0+cu92-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.4.0+cu92-cp35-cp35m-linux_x86_64.whl
\n    torch-1.4.0+cu92-cp35-cp35m-win_amd64.whl
\n    torch-1.4.0+cu92-cp36-cp36m-linux_x86_64.whl
\n    torch-1.4.0+cu92-cp36-cp36m-win_amd64.whl
\n    torch-1.4.0+cu92-cp37-cp37m-linux_x86_64.whl
\n    torch-1.4.0+cu92-cp37-cp37m-win_amd64.whl
\n    torch-1.4.0+cu92-cp38-cp38-linux_x86_64.whl
\n    torch-1.4.0+cu92-cp38-cp38-win_amd64.whl
\n    torch-1.5.0+cu92-cp27-cp27m-linux_x86_64.whl
\n    torch-1.5.0+cu92-cp27-cp27mu-linux_x86_64.whl
\n    torch-1.5.0+cu92-cp35-cp35m-linux_x86_64.whl
\n    torch-1.5.0+cu92-cp35-cp35m-win_amd64.whl
\n    torch-1.5.0+cu92-cp36-cp36m-linux_x86_64.whl
\n    torch-1.5.0+cu92-cp36-cp36m-win_amd64.whl
\n    torch-1.5.0+cu92-cp37-cp37m-linux_x86_64.whl
\n    torch-1.5.0+cu92-cp37-cp37m-win_amd64.whl
\n    torch-1.5.0+cu92-cp38-cp38-linux_x86_64.whl
\n    torch-1.5.0+cu92-cp38-cp38-win_amd64.whl
\n    torch-1.5.1+cu92-cp35-cp35m-linux_x86_64.whl
\n    torch-1.5.1+cu92-cp35-cp35m-win_amd64.whl
\n    torch-1.5.1+cu92-cp36-cp36m-linux_x86_64.whl
\n    torch-1.5.1+cu92-cp36-cp36m-win_amd64.whl
\n    torch-1.5.1+cu92-cp37-cp37m-linux_x86_64.whl
\n    torch-1.5.1+cu92-cp37-cp37m-win_amd64.whl
\n    torch-1.5.1+cu92-cp38-cp38-linux_x86_64.whl
\n    torch-1.5.1+cu92-cp38-cp38-win_amd64.whl
\n    torch-1.6.0+cu92-cp36-cp36m-linux_x86_64.whl
\n    torch-1.6.0+cu92-cp37-cp37m-linux_x86_64.whl
\n    torch-1.6.0+cu92-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.0+cu92-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.0+cu92-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.0+cu92-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.1+cu92-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.1+cu92-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.1+cu92-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.1+cu92-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.0+rocm3.10-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.0+rocm3.10-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.0+rocm3.10-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.0+rocm3.10-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.1+rocm3.10-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.1+rocm3.10-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.1+rocm3.10-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.1+rocm3.10-cp39-cp39-linux_x86_64.whl
\n    torch-1.7.1+rocm3.7-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.1+rocm3.7-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.1+rocm3.7-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.1+rocm3.7-cp39-cp39-linux_x86_64.whl
\n    torch-1.7.1+rocm3.8-cp36-cp36m-linux_x86_64.whl
\n    torch-1.7.1+rocm3.8-cp37-cp37m-linux_x86_64.whl
\n    torch-1.7.1+rocm3.8-cp38-cp38-linux_x86_64.whl
\n    torch-1.7.1+rocm3.8-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.0+rocm4.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.0+rocm4.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.0+rocm4.0.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.0+rocm4.0.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.1+rocm4.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.1+rocm4.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.1+rocm4.0.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.1+rocm4.0.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.2+rocm4.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.2+rocm4.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.2+rocm4.0.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.2+rocm4.0.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.0+rocm4.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.0+rocm4.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.0+rocm4.0.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.0+rocm4.0.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.8.1+rocm4.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.8.1+rocm4.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.8.1+rocm4.0.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.8.1+rocm4.0.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.0+rocm4.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.0+rocm4.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.0+rocm4.0.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.0+rocm4.0.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.1+rocm4.0.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.1+rocm4.0.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.1+rocm4.0.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.1+rocm4.0.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.0+rocm4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.0+rocm4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.0+rocm4.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.0+rocm4.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.1+rocm4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.1+rocm4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.1+rocm4.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.1+rocm4.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.2+rocm4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.2+rocm4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.2+rocm4.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.2+rocm4.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.0+rocm4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.0+rocm4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.0+rocm4.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.0+rocm4.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.1+rocm4.1-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.1+rocm4.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.1+rocm4.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.1+rocm4.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.0+rocm4.2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.0+rocm4.2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.0+rocm4.2-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.0+rocm4.2-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.1+rocm4.2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.1+rocm4.2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.1+rocm4.2-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.1+rocm4.2-cp39-cp39-linux_x86_64.whl
\n    torch-1.10.2+rocm4.2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.10.2+rocm4.2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.10.2+rocm4.2-cp38-cp38-linux_x86_64.whl
\n    torch-1.10.2+rocm4.2-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.0+rocm4.2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.0+rocm4.2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.0+rocm4.2-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.0+rocm4.2-cp39-cp39-linux_x86_64.whl
\n    torch-1.9.1+rocm4.2-cp36-cp36m-linux_x86_64.whl
\n    torch-1.9.1+rocm4.2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.9.1+rocm4.2-cp38-cp38-linux_x86_64.whl
\n    torch-1.9.1+rocm4.2-cp39-cp39-linux_x86_64.whl
\n    torch-1.11.0+rocm4.3.1-cp310-cp310-linux_x86_64.whl
\n    torch-1.11.0+rocm4.3.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.11.0+rocm4.3.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.11.0+rocm4.3.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.11.0+rocm4.5.2-cp310-cp310-linux_x86_64.whl
\n    torch-1.11.0+rocm4.5.2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.11.0+rocm4.5.2-cp38-cp38-linux_x86_64.whl
\n    torch-1.11.0+rocm4.5.2-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.0+rocm5.0-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.0+rocm5.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.0+rocm5.0-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.0+rocm5.0-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.1+rocm5.0-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.1+rocm5.0-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.1+rocm5.0-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.1+rocm5.0-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.0+rocm5.1.1-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.0+rocm5.1.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.0+rocm5.1.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.0+rocm5.1.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.12.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl
\n    torch-1.12.1+rocm5.1.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.12.1+rocm5.1.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.12.1+rocm5.1.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.0+rocm5.1.1-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.0+rocm5.1.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.0+rocm5.1.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.0+rocm5.1.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.1+rocm5.1.1-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.1+rocm5.1.1-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.1+rocm5.1.1-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.0+rocm5.2-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.0+rocm5.2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.0+rocm5.2-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.0+rocm5.2-cp39-cp39-linux_x86_64.whl
\n    torch-1.13.1+rocm5.2-cp310-cp310-linux_x86_64.whl
\n    torch-1.13.1+rocm5.2-cp37-cp37m-linux_x86_64.whl
\n    torch-1.13.1+rocm5.2-cp38-cp38-linux_x86_64.whl
\n    torch-1.13.1+rocm5.2-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.0+rocm5.3-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.0+rocm5.3-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.0+rocm5.3-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.1+rocm5.3-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.1+rocm5.3-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.1+rocm5.3-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.1+rocm5.3-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.0+rocm5.4.2-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.0+rocm5.4.2-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.0+rocm5.4.2-cp39-cp39-linux_x86_64.whl
\n    torch-2.0.1+rocm5.4.2-cp310-cp310-linux_x86_64.whl
\n    torch-2.0.1+rocm5.4.2-cp311-cp311-linux_x86_64.whl
\n    torch-2.0.1+rocm5.4.2-cp38-cp38-linux_x86_64.whl
\n    torch-2.0.1+rocm5.4.2-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.0+rocm5.5-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.0+rocm5.5-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.0+rocm5.5-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.0+rocm5.5-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.1+rocm5.5-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.1+rocm5.5-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.1+rocm5.5-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.1+rocm5.5-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.2+rocm5.5-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.2+rocm5.5-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.2+rocm5.5-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.2+rocm5.5-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.0+rocm5.6-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.0+rocm5.6-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.0+rocm5.6-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.0+rocm5.6-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.1+rocm5.6-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.1+rocm5.6-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.1+rocm5.6-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.1+rocm5.6-cp39-cp39-linux_x86_64.whl
\n    torch-2.1.2+rocm5.6-cp310-cp310-linux_x86_64.whl
\n    torch-2.1.2+rocm5.6-cp311-cp311-linux_x86_64.whl
\n    torch-2.1.2+rocm5.6-cp38-cp38-linux_x86_64.whl
\n    torch-2.1.2+rocm5.6-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.0+rocm5.6-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.0+rocm5.6-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.0+rocm5.6-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.0+rocm5.6-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.0+rocm5.6-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.1+rocm5.6-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.1+rocm5.6-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.1+rocm5.6-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.1+rocm5.6-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.1+rocm5.6-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.2+rocm5.6-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.2+rocm5.6-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.2+rocm5.6-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.2+rocm5.6-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.2+rocm5.6-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.0+rocm5.7-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.0+rocm5.7-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.0+rocm5.7-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.0+rocm5.7-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.0+rocm5.7-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.1+rocm5.7-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.1+rocm5.7-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.1+rocm5.7-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.1+rocm5.7-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.1+rocm5.7-cp39-cp39-linux_x86_64.whl
\n    torch-2.2.2+rocm5.7-cp310-cp310-linux_x86_64.whl
\n    torch-2.2.2+rocm5.7-cp311-cp311-linux_x86_64.whl
\n    torch-2.2.2+rocm5.7-cp312-cp312-linux_x86_64.whl
\n    torch-2.2.2+rocm5.7-cp38-cp38-linux_x86_64.whl
\n    torch-2.2.2+rocm5.7-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.0+rocm5.7-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.0+rocm5.7-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.0+rocm5.7-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.0+rocm5.7-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.0+rocm5.7-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.1+rocm5.7-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.1+rocm5.7-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.1+rocm5.7-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.1+rocm5.7-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.1+rocm5.7-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.0+rocm6.0-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.0+rocm6.0-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.0+rocm6.0-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.0+rocm6.0-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.0+rocm6.0-cp39-cp39-linux_x86_64.whl
\n    torch-2.3.1+rocm6.0-cp310-cp310-linux_x86_64.whl
\n    torch-2.3.1+rocm6.0-cp311-cp311-linux_x86_64.whl
\n    torch-2.3.1+rocm6.0-cp312-cp312-linux_x86_64.whl
\n    torch-2.3.1+rocm6.0-cp38-cp38-linux_x86_64.whl
\n    torch-2.3.1+rocm6.0-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.0+rocm6.0-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.0+rocm6.0-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.0+rocm6.0-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.0+rocm6.0-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.0+rocm6.0-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.1+rocm6.0-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.1+rocm6.0-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.1+rocm6.0-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.1+rocm6.0-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.1+rocm6.0-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.0+rocm6.1-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.0+rocm6.1-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.0+rocm6.1-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.0+rocm6.1-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.0+rocm6.1-cp39-cp39-linux_x86_64.whl
\n    torch-2.4.1+rocm6.1-cp310-cp310-linux_x86_64.whl
\n    torch-2.4.1+rocm6.1-cp311-cp311-linux_x86_64.whl
\n    torch-2.4.1+rocm6.1-cp312-cp312-linux_x86_64.whl
\n    torch-2.4.1+rocm6.1-cp38-cp38-linux_x86_64.whl
\n    torch-2.4.1+rocm6.1-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.0+rocm6.1-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.0+rocm6.1-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.0+rocm6.1-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.0+rocm6.1-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.1+rocm6.1-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.1+rocm6.1-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.1+rocm6.1-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.1+rocm6.1-cp39-cp39-linux_x86_64.whl
\n    torch-2.6.0+rocm6.1-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.1-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.1-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.1-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.1-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.1-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.2.4-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.2.4-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.2.4-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.2.4-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.2.4-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.6.0+rocm6.2.4-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.2.4-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.2.4-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.2.4-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.2.4-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.2.4-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.2.4-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.2.4-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.2.4-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.2.4-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.2.4-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.2.4-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.2.4-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.5.0+rocm6.2-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.0+rocm6.2-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.0+rocm6.2-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.0+rocm6.2-cp39-cp39-linux_x86_64.whl
\n    torch-2.5.1+rocm6.2-cp310-cp310-linux_x86_64.whl
\n    torch-2.5.1+rocm6.2-cp311-cp311-linux_x86_64.whl
\n    torch-2.5.1+rocm6.2-cp312-cp312-linux_x86_64.whl
\n    torch-2.5.1+rocm6.2-cp39-cp39-linux_x86_64.whl
\n    torch-2.7.0+rocm6.3-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.3-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.3-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.3-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.3-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.0+rocm6.3-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.3-cp310-cp310-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.3-cp311-cp311-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.3-cp312-cp312-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.3-cp313-cp313-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.3-cp313-cp313t-manylinux_2_28_x86_64.whl
\n    torch-2.7.1+rocm6.3-cp39-cp39-manylinux_2_28_x86_64.whl
\n    torch-0.1-cp27-cp27m-macosx_10_6_x86_64.whl
\n    torch-0.1-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.1.10.post1-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.10.post1-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.1.10.post1-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.1.11.post4-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.11.post4-cp35-cp35m-macosx_10_7_x86_64.whl
\n    torch-0.1.11.post4-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.1.11.post5-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.11.post5-cp35-cp35m-macosx_10_7_x86_64.whl
\n    torch-0.1.11.post5-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.1.12.post1-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.12.post1-cp35-cp35m-macosx_10_7_x86_64.whl
\n    torch-0.1.12.post1-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.1.12.post2-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.12.post2-cp35-cp35m-macosx_10_7_x86_64.whl
\n    torch-0.1.12.post2-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.1.6.post17-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.1.6.post17-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.6.post20-cp27-cp27mu-linux_x86_64.whl
\n    torch-0.1.6.post20-cp35-cp35m-linux_x86_64.whl
\n    torch-0.1.6.post22-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.6.post22-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.1.7.post2-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.7.post2-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.1.7.post2-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.1.8.post1-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.8.post1-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.1.8.post1-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.1.9.post1-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.9.post1-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.1.9.post1-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.1.9.post2-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.1.9.post2-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.1.9.post2-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post1-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post1-cp35-cp35m-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post1-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post2-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post2-cp35-cp35m-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post2-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post3-cp27-none-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post3-cp35-cp35m-macosx_10_7_x86_64.whl
\n    torch-0.2.0.post3-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.3.0-cp27-none-macosx_10_6_x86_64.whl
\n    torch-0.3.0-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.3.0-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.3.0.post2-cp27-none-macosx_10_6_x86_64.whl
\n    torch-0.3.0.post2-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.3.0.post2-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.3.0.post3-cp27-none-macosx_10_6_x86_64.whl
\n    torch-0.3.0.post3-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.3.0.post3-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.3.0.post4-cp27-none-macosx_10_6_x86_64.whl
\n    torch-0.3.0.post4-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.3.0.post4-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.3.1-cp27-none-macosx_10_6_x86_64.whl
\n    torch-0.3.1-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.3.1-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.4.0-cp27-none-macosx_10_6_x86_64.whl
\n    torch-0.4.0-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.4.0-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.4.1-cp27-none-macosx_10_6_x86_64.whl
\n    torch-0.4.1-cp35-cp35m-macosx_10_6_x86_64.whl
\n    torch-0.4.1-cp36-cp36m-macosx_10_7_x86_64.whl
\n    torch-0.4.1-cp37-cp37m-macosx_10_7_x86_64.whl
\n    torch-1.11.0-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-1.11.0-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.11.0-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.11.0-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.12.0-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-1.12.0-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.12.0-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.12.0-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.12.1-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-1.12.1-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.12.1-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.12.1-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.13.0-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-1.13.0-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.13.0-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.13.0-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-1.13.1-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-1.13.1-cp37-cp37m-manylinux2014_aarch64.whl
\n    torch-1.13.1-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-1.13.1-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-2.0.0-1-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-2.0.0-1-cp311-cp311-manylinux2014_aarch64.whl
\n    torch-2.0.0-1-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-2.0.0-1-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-2.0.0-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-2.0.0-cp311-cp311-manylinux2014_aarch64.whl
\n    torch-2.0.0-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-2.0.0-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-2.0.1-cp310-cp310-manylinux2014_aarch64.whl
\n    torch-2.0.1-cp311-cp311-manylinux2014_aarch64.whl
\n    torch-2.0.1-cp38-cp38-manylinux2014_aarch64.whl
\n    torch-2.0.1-cp39-cp39-manylinux2014_aarch64.whl
\n    torch-2.6.0+xpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.6.0+xpu-cp310-cp310-win_amd64.whl
\n    torch-2.6.0+xpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.6.0+xpu-cp311-cp311-win_amd64.whl
\n    torch-2.6.0+xpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.6.0+xpu-cp312-cp312-win_amd64.whl
\n    torch-2.6.0+xpu-cp313-cp313-linux_x86_64.whl
\n    torch-2.6.0+xpu-cp313-cp313-win_amd64.whl
\n    torch-2.6.0+xpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.6.0+xpu-cp39-cp39-win_amd64.whl
\n    torch-2.7.0+xpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.7.0+xpu-cp310-cp310-win_amd64.whl
\n    torch-2.7.0+xpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.7.0+xpu-cp311-cp311-win_amd64.whl
\n    torch-2.7.0+xpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.7.0+xpu-cp312-cp312-win_amd64.whl
\n    torch-2.7.0+xpu-cp313-cp313-linux_x86_64.whl
\n    torch-2.7.0+xpu-cp313-cp313-win_amd64.whl
\n    torch-2.7.0+xpu-cp313-cp313t-linux_x86_64.whl
\n    torch-2.7.0+xpu-cp313-cp313t-win_amd64.whl
\n    torch-2.7.0+xpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.7.0+xpu-cp39-cp39-win_amd64.whl
\n    torch-2.7.1+xpu-cp310-cp310-linux_x86_64.whl
\n    torch-2.7.1+xpu-cp310-cp310-win_amd64.whl
\n    torch-2.7.1+xpu-cp311-cp311-linux_x86_64.whl
\n    torch-2.7.1+xpu-cp311-cp311-win_amd64.whl
\n    torch-2.7.1+xpu-cp312-cp312-linux_x86_64.whl
\n    torch-2.7.1+xpu-cp312-cp312-win_amd64.whl
\n    torch-2.7.1+xpu-cp313-cp313-linux_x86_64.whl
\n    torch-2.7.1+xpu-cp313-cp313-win_amd64.whl
\n    torch-2.7.1+xpu-cp313-cp313t-linux_x86_64.whl
\n    torch-2.7.1+xpu-cp313-cp313t-win_amd64.whl
\n    torch-2.7.1+xpu-cp39-cp39-linux_x86_64.whl
\n    torch-2.7.1+xpu-cp39-cp39-win_amd64.whl
\n  \n\n\n"
  },
  {
    "path": "tools/compatgen/internal/util.go",
    "content": "package internal\n\nimport (\n\t\"strings\"\n)\n\nfunc split2(s string, sep string) (string, string) {\n\tparts := strings.SplitN(s, sep, 2)\n\treturn parts[0], parts[1]\n}\n"
  },
  {
    "path": "tools/compatgen/main.go",
    "content": "package main\n\nimport (\n\t\"context\"\n\t\"encoding/json\"\n\t\"os\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/util/console\"\n\t\"github.com/replicate/cog/tools/compatgen/internal\"\n)\n\nfunc main() {\n\tvar output string\n\n\tvar rootCmd = &cobra.Command{\n\t\tUse:   \"compatgen {cuda|torch|tensorflow}\",\n\t\tShort: \"Generate compatibility matrix for Cog base images\",\n\t\tArgs:  cobra.ExactArgs(1),\n\t\tRun: func(cmd *cobra.Command, args []string) {\n\t\t\tctx := context.Background()\n\t\t\ttarget := args[0]\n\n\t\t\tvar v any\n\t\t\tvar err error\n\n\t\t\tswitch target {\n\t\t\tcase \"cuda\":\n\t\t\t\tv, err = internal.FetchCUDABaseImages(ctx)\n\t\t\t\tif err != nil {\n\t\t\t\t\tconsole.Fatalf(\"Failed to fetch CUDA base image tags: %s\", err)\n\t\t\t\t}\n\t\t\tcase \"tensorflow\":\n\t\t\t\tv, err = internal.FetchTensorFlowCompatibilityMatrix()\n\t\t\t\tif err != nil {\n\t\t\t\t\tconsole.Fatalf(\"Failed to fetch TensorFlow compatibility matrix: %s\", err)\n\t\t\t\t}\n\t\t\tcase \"torch\":\n\t\t\t\tv, err = internal.FetchTorchCompatibilityMatrix()\n\t\t\t\tif err != nil {\n\t\t\t\t\tconsole.Fatalf(\"Failed to fetch PyTorch compatibility matrix: %s\", err)\n\t\t\t\t}\n\t\t\tdefault:\n\t\t\t\tconsole.Fatalf(\"Unknown target: %s\", target)\n\t\t\t}\n\n\t\t\tdata, err := json.MarshalIndent(v, \"\", \"  \")\n\t\t\tif err != nil {\n\t\t\t\tconsole.Fatalf(\"Failed to marshal value: %s\", err)\n\t\t\t}\n\n\t\t\tif output != \"\" {\n\t\t\t\tif err := os.WriteFile(output, data, 0o644); err != nil {\n\t\t\t\t\tconsole.Fatalf(\"Failed to write to %s: %s\", output, err)\n\t\t\t\t}\n\t\t\t\tconsole.Infof(\"Wrote to %s\", output)\n\t\t\t} else {\n\t\t\t\tconsole.Output(string(data))\n\t\t\t}\n\t\t},\n\t}\n\n\trootCmd.Flags().StringVarP(&output, \"output\", \"o\", \"\", \"Output flag (optional)\")\n\tif err := rootCmd.Execute(); err != nil {\n\t\tconsole.Fatal(err.Error())\n\t}\n}\n"
  },
  {
    "path": "tools/gendocs/main.go",
    "content": "package main\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"slices\"\n\t\"sort\"\n\t\"strings\"\n\n\t\"github.com/spf13/cobra\"\n\t\"github.com/spf13/cobra/doc\"\n\n\t\"github.com/replicate/cog/pkg/cli\"\n\t\"github.com/replicate/cog/pkg/util/console\"\n)\n\nfunc main() {\n\tvar output string\n\n\trootCmd := &cobra.Command{\n\t\tUse:   \"gendocs\",\n\t\tShort: \"Generate CLI reference documentation for Cog\",\n\t\tRun: func(cmd *cobra.Command, args []string) {\n\t\t\tif err := generateDocs(output); err != nil {\n\t\t\t\tconsole.Fatalf(\"Failed to generate docs: %s\", err)\n\t\t\t}\n\t\t\tconsole.Infof(\"Generated CLI docs at %s\", output)\n\t\t},\n\t}\n\n\trootCmd.Flags().StringVarP(&output, \"output\", \"o\", \"docs/cli.md\", \"Output file path\")\n\tif err := rootCmd.Execute(); err != nil {\n\t\tconsole.Fatal(err.Error())\n\t}\n}\n\nfunc generateDocs(outputPath string) error {\n\t// Create temporary directory for cobra doc generation\n\ttmpDir, err := os.MkdirTemp(\"\", \"cog-cli-docs-*\")\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to create temp dir: %w\", err)\n\t}\n\tdefer os.RemoveAll(tmpDir)\n\n\t// Get the cog command\n\tcmd, err := cli.NewRootCommand()\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to create root command: %w\", err)\n\t}\n\n\t// Generate markdown files using cobra/doc\n\tif err := doc.GenMarkdownTree(cmd, tmpDir); err != nil {\n\t\treturn fmt.Errorf(\"failed to generate markdown: %w\", err)\n\t}\n\n\t// Read all generated files\n\tfiles, err := os.ReadDir(tmpDir)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to read temp dir: %w\", err)\n\t}\n\n\t// Sort files to ensure consistent ordering\n\t// Order: cog (root), then alphabetically by command name\n\tvar fileNames []string\n\tfor _, file := range files {\n\t\tif !file.IsDir() && strings.HasSuffix(file.Name(), \".md\") {\n\t\t\tfileNames = append(fileNames, file.Name())\n\t\t}\n\t}\n\tsort.Strings(fileNames)\n\n\t// Build the combined markdown content\n\tvar content strings.Builder\n\n\t// Write header\n\tcontent.WriteString(\"# CLI reference\\n\\n\")\n\tcontent.WriteString(\"\\n\\n\")\n\n\t// Process each command file\n\tfor _, fileName := range fileNames {\n\t\tfilePath := filepath.Join(tmpDir, fileName)\n\t\tdata, err := os.ReadFile(filePath)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"failed to read %s: %w\", fileName, err)\n\t\t}\n\n\t\t// Process the content\n\t\tprocessed := processCommandDoc(string(data), fileName)\n\t\tcontent.WriteString(processed)\n\t\tcontent.WriteString(\"\\n\")\n\t}\n\n\t// Ensure output directory exists\n\toutputDir := filepath.Dir(outputPath)\n\tif err := os.MkdirAll(outputDir, 0o755); err != nil {\n\t\treturn fmt.Errorf(\"failed to create output directory: %w\", err)\n\t}\n\n\t// Write the combined file\n\tif err := os.WriteFile(outputPath, []byte(content.String()), 0o644); err != nil {\n\t\treturn fmt.Errorf(\"failed to write output file: %w\", err)\n\t}\n\n\treturn nil\n}\n\nfunc processCommandDoc(content string, fileName string) string {\n\t// Remove the \"SEE ALSO\" section and everything after it\n\tif idx := strings.Index(content, \"### SEE ALSO\"); idx != -1 {\n\t\tcontent = content[:idx]\n\t}\n\n\t// Remove the \"Options inherited from parent commands\" section\n\tif idx := strings.Index(content, \"### Options inherited from parent commands\"); idx != -1 {\n\t\tcontent = content[:idx]\n\t}\n\n\t// Remove trailing whitespace\n\tcontent = strings.TrimRight(content, \"\\n\")\n\n\t// Fix command headers to use backticks\n\t// Change \"## cog init\" to \"## `cog init`\"\n\t// Change \"### Options\" to \"**Options**\" (not a heading, won't appear in TOC)\n\t// Change \"### Examples\" to \"**Examples**\" (not a heading, won't appear in TOC)\n\t// Remove \"### Synopsis\" heading but keep its content\n\t// Skip the short description if there's a Synopsis section (to avoid duplication)\n\tlines := strings.Split(content, \"\\n\")\n\tvar result []string\n\tskipSynopsis := false\n\tskipShortDesc := false\n\tfor _, line := range lines {\n\t\tswitch {\n\t\tcase strings.HasPrefix(line, \"## cog\"):\n\t\t\t// Extract the command name\n\t\t\tcommand := strings.TrimPrefix(line, \"## \")\n\t\t\tresult = append(result, \"## `\"+command+\"`\")\n\t\t\t// Check if next non-empty line is \"### Synopsis\" - if so, skip the short desc\n\t\t\tskipShortDesc = hasSynopsisSection(lines)\n\t\tcase skipShortDesc:\n\t\t\t// Skip the short description line (first non-empty line after header)\n\t\t\t// Also skip any blank lines that follow the header\n\t\t\tif strings.TrimSpace(line) != \"\" && !strings.HasPrefix(line, \"###\") {\n\t\t\t\t// This is the short description line, skip it\n\t\t\t\tskipShortDesc = false\n\t\t\t}\n\t\t\t// If line is blank, we continue skipping until we hit the short desc\n\t\tcase line == \"### Synopsis\":\n\t\t\t// Skip the \"### Synopsis\" heading line, but keep content after it\n\t\t\tskipSynopsis = true\n\t\tcase skipSynopsis:\n\t\t\t// Keep synopsis content until we hit the usage block (```) or another heading\n\t\t\tswitch {\n\t\t\tcase line == \"### Examples\":\n\t\t\t\tskipSynopsis = false\n\t\t\t\t// Add blank line before if needed\n\t\t\t\tif len(result) > 0 && strings.TrimSpace(result[len(result)-1]) != \"\" {\n\t\t\t\t\tresult = append(result, \"\")\n\t\t\t\t}\n\t\t\t\tresult = append(result, \"**Examples**\")\n\t\t\tcase strings.HasPrefix(line, \"###\"), strings.HasPrefix(line, \"```\"):\n\t\t\t\tskipSynopsis = false\n\t\t\t\t// Add blank line before if needed\n\t\t\t\tif len(result) > 0 && strings.TrimSpace(result[len(result)-1]) != \"\" {\n\t\t\t\t\tresult = append(result, \"\")\n\t\t\t\t}\n\t\t\t\tresult = append(result, line)\n\t\t\tdefault:\n\t\t\t\t// Keep all lines from synopsis (including blank lines for paragraph breaks)\n\t\t\t\tresult = append(result, line)\n\t\t\t}\n\t\tcase line == \"### Options\":\n\t\t\t// Add blank line before if needed\n\t\t\tif len(result) > 0 && strings.TrimSpace(result[len(result)-1]) != \"\" {\n\t\t\t\tresult = append(result, \"\")\n\t\t\t}\n\t\t\tresult = append(result, \"**Options**\")\n\t\tcase line == \"### Examples\":\n\t\t\t// Add blank line before if needed\n\t\t\tif len(result) > 0 && strings.TrimSpace(result[len(result)-1]) != \"\" {\n\t\t\t\tresult = append(result, \"\")\n\t\t\t}\n\t\t\tresult = append(result, \"**Examples**\")\n\t\tdefault:\n\t\t\tresult = append(result, line)\n\t\t}\n\t}\n\n\t// Remove consecutive blank lines\n\tresult = removeConsecutiveBlankLines(result)\n\n\treturn strings.Join(result, \"\\n\")\n}\n\n// removeConsecutiveBlankLines removes consecutive blank lines, keeping only one\nfunc removeConsecutiveBlankLines(lines []string) []string {\n\tvar result []string\n\tprevBlank := false\n\tfor _, line := range lines {\n\t\tisBlank := strings.TrimSpace(line) == \"\"\n\t\tif isBlank && prevBlank {\n\t\t\t// Skip consecutive blank lines\n\t\t\tcontinue\n\t\t}\n\t\tresult = append(result, line)\n\t\tprevBlank = isBlank\n\t}\n\treturn result\n}\n\n// hasSynopsisSection checks if the content has a \"### Synopsis\" section\nfunc hasSynopsisSection(lines []string) bool {\n\treturn slices.Contains(lines, \"### Synopsis\")\n}\n"
  },
  {
    "path": "tools/install.sh",
    "content": "#!/bin/sh\n#\n# This script should be run via curl:\n#   sh -c \"$(curl -fsSL https://raw.githubusercontent.com/replicate/cog/main/tools/install.sh)\"\n# or via wget:\n#   sh -c \"$(wget -qO- https://raw.githubusercontent.com/replicate/cog/main/tools/install.sh)\"\n# or via fetch:\n#   sh -c \"$(fetch -o - https://raw.githubusercontent.com/replicate/cog/main/tools/install.sh)\"\n#\n# As an alternative, you can first download the install script and run it afterwards:\n#   wget https://raw.githubusercontent.com/replicate/cog/main/tools/install.sh\n#   sh install.sh\n#\n# You can tweak the install location by setting the INSTALL_DIR env var when running the script.\n#   INSTALL_DIR=~/my/custom/install/location sh install.sh\n#\n# By default, cog will be installed at /usr/local/bin/cog\n\n\n# This install script is based on that of ohmyzsh[1], which is licensed under the MIT License\n# [1] https://github.com/ohmyzsh/ohmyzsh/blob/master/tools/install.sh\n# MIT License\n\n# Copyright (c) 2009-2022 Robby Russell and contributors (https://github.com/ohmyzsh/ohmyzsh/contributors)\n\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE.\nset -e\n\nset_install_dir() {\n  # Set install directory\n  DEFAULT_INSTALL_DIR=\"/usr/local/bin\"\n  if [ -z \"${INSTALL_DIR}\" ]; then\n    read -p \"Install location? [$DEFAULT_INSTALL_DIR]: \" INSTALL_DIR\n    INSTALL_DIR=${INSTALL_DIR:-$DEFAULT_INSTALL_DIR}\n  fi\n  if [ ! -d \"$INSTALL_DIR\" ]; then\n    echo \"The directory $INSTALL_DIR does not exist. Please create it and re-run this script.\"\n    # Ask user to manually create directory rather than making it for them,\n    # so they don't just type in \"y\" again and accidentally install at ./y\n    exit 1\n  fi\n  # Expand abbreviations in INSTALL_DIR\n  INSTALL_DIR=$(cd \"$INSTALL_DIR\"; pwd)\n}\n\ncommand_exists() {\n  command -v \"$@\" >/dev/null 2>&1\n}\n\nuser_can_sudo() {\n  # Check if sudo is installed\n  command_exists $SUDO || return 1\n  # Termux can't run sudo, so we can detect it and exit the function early.\n  case \"$PREFIX\" in\n  *com.termux*) return 1 ;;\n  esac\n  # The following command has 3 parts:\n  #\n  # 1. Run `sudo` with `-v`. Does the following:\n  #    • with privilege: asks for a password immediately.\n  #    • without privilege: exits with error code 1 and prints the message:\n  #      Sorry, user  may not run sudo on \n  #\n  # 2. Pass `-n` to `sudo` to tell it to not ask for a password. If the\n  #    password is not required, the command will finish with exit code 0.\n  #    If one is required, sudo will exit with error code 1 and print the\n  #    message:\n  #    sudo: a password is required\n  #\n  # 3. Check for the words \"may not run sudo\" in the output to really tell\n  #    whether the user has privileges or not. For that we have to make sure\n  #    to run `sudo` in the default locale (with `LANG=`) so that the message\n  #    stays consistent regardless of the user's locale.\n  #\n  ! LANG= $SUDO -n -v 2>&1 | grep -q \"may not run $SUDO\"\n}\n\ncheck_docker() {\n  if ! command_exists docker; then\n  echo \"Docker is not installed on your system. Please install Docker before proceeding.\"\n    exit 1\n  fi\n\n  if ! docker run hello-world >/dev/null 2>&1; then\n    echo \"WARNING: Docker engine is not running, or docker cannot be run without sudo. Please setup Docker so that your user has permission to run it: https://docs.docker.com/engine/install/linux-postinstall/\"\n  fi\n}\n\nsetup_cog() {\n  COG_LOCATION=\"${INSTALL_DIR}/cog\"\n  BINARY_URI=\"https://github.com/replicate/cog/releases/latest/download/cog_$(uname -s)_$(uname -m)\"\n  if [ -f \"$COG_LOCATION\" ]; then\n    echo \"A file already exists at $COG_LOCATION\"\n    echo \"Do you want to delete this file and continue with this installation anyway?\"\n    read -p \"Delete file? (y/N): \" choice\n    case \"$choice\" in \n      y|Y ) echo \"Deleting existing file and continuing with installation...\"; $SUDO rm $COG_LOCATION;;\n      * ) echo \"Exiting installation.\"; exit 1;;\n    esac\n  fi\n  if command_exists curl; then\n    $SUDO curl -o $COG_LOCATION -L $BINARY_URI\n  elif command_exists wget; then\n    $SUDO wget $BINARY_URI -O $COG_LOCATION\n  elif command_exists fetch; then\n    $SUDO fetch -o $COG_LOCATION $BINARY_URI\n  else\n    echo \"One of curl, wget, or fetch must be present for this installer to work.\"\n    exit 1\n  fi\n  if [ \"$(cat $COG_LOCATION)\" = \"Not Found\" ]; then\n    echo \"Error: Cog binary not found at ${BINARY_URI}. Check releases to see if a binary is available for your system.\"\n    rm $COG_LOCATION\n    exit 1\n  fi\n\n  $SUDO chmod +x $COG_LOCATION\n\n  # On macOS, remove the quarantine attribute that triggers Gatekeeper's\n  # \"cannot be opened because the developer cannot be verified\" warning.\n  if [ \"$(uname -s)\" = \"Darwin\" ]; then\n    $SUDO xattr -d com.apple.quarantine \"$COG_LOCATION\" 2>/dev/null || true\n  fi\n\n  SHELL_NAME=$(basename \"$SHELL\")\n  if [[ \":$PATH:\" != *\":$INSTALL_DIR:\"* ]]; then\n    echo \"Adding $INSTALL_DIR to PATH in .$SHELL_NAME\"rc\n    echo \"\" >> ~/.$SHELL_NAME\"rc\"\n    echo \"# Created by \\`cog\\` install script on $(date)\" >> ~/.$SHELL_NAME\"rc\"\n    echo \"export PATH=\\$PATH:$INSTALL_DIR\" >> ~/.$SHELL_NAME\"rc\"\n    source ~/.$SHELL_NAME\"rc\"\n\n    echo \"You may need to open a new terminal window to run cog for the first time.\"\n  fi\n    \n  echo\n}\n\n\nprint_success() {\n  echo \"Successfully installed cog. Run \\`cog login\\` to configure Replicate access\"\n}\n\nmain() {\n\n  # Check if macOS\n  if [ \"$(uname -s)\" = \"Darwin\" ]; then\n    echo \"On macOS, it is recommended to install cog using Homebrew instead:\"\n    echo \\`brew install replicate/tap/cog\\`\n    echo \"Do you want to continue with this installation anyway?\"\n    \n    read -p \"Continue? (y/N): \" choice\n    case \"$choice\" in \n      y|Y ) echo \"Continuing with installation...\";;\n      * ) echo \"Exiting installation.\"; exit 1;;\n    esac\n  fi\n\n  set_install_dir\n\n  # Check if `cog` command already exists\n  if command_exists cog; then\n    echo \"A cog command already exists on your system at the following location: $(which cog)\".\n    echo \"The installations may interfere with one another.\"\n    echo \"Do you want to continue with this installation anyway?\"\n    read -p \"Continue? (y/N): \" choice\n    case \"$choice\" in \n      y|Y ) echo \"Continuing with installation...\";;\n      * ) echo \"Exiting installation.\"; exit 1;;\n    esac\n  fi\n\n  # Check the users sudo privileges\n  if [ -z \"${SUDO+set}\" ]; then\n    SUDO=\"sudo\"\n  fi\n  if [ ! user_can_sudo ] && [ \"${SUDO}\" != \"\" ]; then\n    echo \"You need sudo permissions to run this install script. Please try again as a sudoer.\"\n    exit 1\n  fi\n\n  check_docker\n  setup_cog\n\n  if command_exists cog; then\n    print_success\n  else\n    echo 'Error: cog not installed.'\n    exit 1\n  fi\n}\n\nmain \"$@\"\n"
  },
  {
    "path": "tools/test-harness/.gitignore",
    "content": "results/*.json\n__pycache__/\n*.pyc\n.venv/\n"
  },
  {
    "path": "tools/test-harness/README.md",
    "content": "# Cog Model Test Harness\n\nAutomated test harness for validating cog models against new SDK versions.\nDesigned to test any cog model from any repo.\n\n## Quick Start\n\n```bash\ncd tools/test-harness\n\n# Create a venv and install dependencies\npython3 -m venv .venv\nsource .venv/bin/activate\npip install pyyaml\n\n# List all models in the manifest\npython -m harness list\n\n# Run all non-GPU models\npython -m harness run --no-gpu\n\n# Run a specific model\npython -m harness run --model hello-world\n\n# Run GPU models only (requires NVIDIA GPU + nvidia-docker)\npython -m harness run --gpu-only\n\n# Output JSON report\npython -m harness run --no-gpu --output json --output-file results/report.json\n\n# Build images only (no predictions)\npython -m harness build --no-gpu\n```\n\n## Prerequisites\n\n- Python 3.10+\n- Docker\n- For GPU models: NVIDIA GPU + nvidia-docker runtime\n\n### Version Resolution\n\nBy default the harness automatically resolves the **latest stable** versions\nof both the cog CLI (from GitHub releases) and the Python SDK (from PyPI),\nskipping any alpha/beta/rc tags. You can override either via the CLI or in\n`manifest.yaml`:\n\n```bash\n# Use the latest stable CLI + SDK (default)\npython -m harness run --no-gpu\n\n# Pin a specific CLI version\npython -m harness run --cog-version v0.16.12 --no-gpu\n\n# Pin a specific SDK version\npython -m harness run --sdk-version 0.16.12 --no-gpu\n\n# Use a pre-release CLI\npython -m harness run --cog-version v0.17.0-rc.2 --no-gpu\n\n# Use a locally-built binary (overrides --cog-version)\npython -m harness run --cog-binary ./dist/go/darwin-arm64/cog --no-gpu\n```\n\nYou can also pin versions in `manifest.yaml` under `defaults`:\n\n```yaml\ndefaults:\n  sdk_version: \"latest\"    # or pin e.g. \"0.16.12\"\n  cog_version: \"latest\"    # or pin e.g. \"v0.16.12\"\n```\n\n**Resolution priority** (for both CLI and SDK): CLI flag > manifest default > latest stable.\n\n## Manifest Format\n\nModels are defined in `manifest.yaml`. Each entry specifies a GitHub repo,\nsubdirectory, test inputs, and expected outputs:\n\n```yaml\nmodels:\n  - name: hello-world\n    repo: replicate/cog-examples\n    path: hello-world\n    gpu: false\n    tests:\n      - description: \"basic predict\"\n        inputs:\n          text: \"world\"\n        expect:\n          type: exact\n          value: \"hello world\"\n```\n\n### Model Fields\n\n| Field | Required | Description |\n|-------|----------|-------------|\n| `name` | yes | Unique identifier for the model |\n| `repo` | yes | GitHub `owner/repo` to clone |\n| `path` | no | Subdirectory within the repo (default: `.`) |\n| `gpu` | no | Whether the model requires a GPU (default: `false`) |\n| `sdk_version` | no | Override the SDK version (default: from `defaults.sdk_version`) |\n| `timeout` | no | Per-prediction timeout in seconds (default: 300) |\n| `requires_env` | no | List of env vars that must be set; model is skipped if missing |\n| `env` | no | Extra env vars to pass; supports `${VAR}` expansion from host |\n| `cog_yaml_overrides` | no | Dict deep-merged into the model's cog.yaml |\n| `tests` | no | List of predict test cases |\n| `train_tests` | no | List of train test cases |\n\n### Input References\n\nPrefix a value with `@` to reference a file in `fixtures/`:\n\n```yaml\ninputs:\n  image: \"@test_image.png\"    # resolves to fixtures/test_image.png\n```\n\n### Validation Types\n\n| Type | Fields | Description |\n|------|--------|-------------|\n| `exact` | `value` | Output must equal value exactly |\n| `contains` | `value` | Output must contain the substring |\n| `regex` | `pattern` | Output must match the regex |\n| `file_exists` | `mime` (optional) | Output is a file path that must exist |\n| `json_match` | `match` | Output parsed as JSON must contain the given subset |\n| `json_keys` | `keys` (optional) | Output parsed as JSON dict must have entries |\n| `not_empty` | — | Output must be non-empty |\n\n## Adding a New Model\n\nAdd an entry to `manifest.yaml`:\n\n```yaml\n  - name: my-model\n    repo: myorg/my-model-repo\n    path: \".\"\n    gpu: true\n    # sdk_version: \"0.16.12\"  # optional per-model override\n    env:\n      HF_TOKEN: \"${HF_TOKEN}\"\n    timeout: 600\n    tests:\n      - description: \"smoke test\"\n        inputs:\n          prompt: \"hello\"\n        expect:\n          type: contains\n          value: \"result\"\n```\n\nNo code changes required.\n\n## CLI Reference\n\n```\nusage: cog-test {run,build,list} [options]\n\nCommands:\n  run     Build and test models (full pipeline)\n  build   Build Docker images only (no predictions)\n  list    List models defined in the manifest\n\nCommon options:\n  --manifest PATH       Path to manifest.yaml\n  --model NAME          Run only this model (repeatable)\n  --no-gpu              Skip GPU models\n  --gpu-only            Only run GPU models\n  --sdk-version VER     SDK version (default: latest stable from PyPI)\n  --cog-version TAG     CLI version to download (default: latest stable)\n  --cog-binary PATH     Path to local cog binary (overrides --cog-version)\n  --keep-images         Don't clean up Docker images after run\n\nRun-specific options:\n  --output {console,json}   Output format (default: console)\n  --output-file PATH        Write report to file\n```\n\n## Architecture\n\n```\ntools/test-harness/\n├── manifest.yaml           # Declarative test definitions\n├── fixtures/               # Test input files (images, etc.)\n├── harness/\n│   ├── cli.py              # CLI entry point\n│   ├── cog_resolver.py     # Resolves + downloads cog CLI and SDK versions\n│   ├── runner.py           # Clone -> patch -> build -> predict -> validate\n│   ├── patcher.py          # Patches cog.yaml with sdk_version + overrides\n│   ├── validators.py       # Output validation strategies\n│   └── report.py           # Console + JSON report generation\n├── results/                # Output reports (gitignored)\n└── pyproject.toml\n```\n"
  },
  {
    "path": "tools/test-harness/harness/__init__.py",
    "content": ""
  },
  {
    "path": "tools/test-harness/harness/__main__.py",
    "content": "\"\"\"Allow running as ``python -m harness``.\"\"\"\n\nfrom .cli import main\n\nmain()\n"
  },
  {
    "path": "tools/test-harness/harness/cli.py",
    "content": "\"\"\"CLI entry point for the cog test harness.\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport logging\nimport sys\nfrom pathlib import Path\n\nimport yaml\n\nfrom .cog_resolver import resolve_cog_binary, resolve_sdk_version\nfrom .report import console_report, write_json_report\nfrom .runner import ModelResult, Runner\n\n\ndef main(argv: list[str] | None = None) -> None:\n    parser = argparse.ArgumentParser(\n        prog=\"cog-test\",\n        description=\"Test harness for validating cog models against new SDK versions\",\n    )\n    subparsers = parser.add_subparsers(dest=\"command\")\n\n    # ── run ─────────────────────────────────────────────────────────\n    run_parser = subparsers.add_parser(\"run\", help=\"Build and test models\")\n    _add_common_args(run_parser)\n    run_parser.add_argument(\n        \"--output\",\n        choices=[\"console\", \"json\"],\n        default=\"console\",\n        help=\"Output format (default: console)\",\n    )\n    run_parser.add_argument(\n        \"--output-file\",\n        type=str,\n        default=None,\n        help=\"Write report to file instead of stdout\",\n    )\n\n    # ── build ───────────────────────────────────────────────────────\n    build_parser = subparsers.add_parser(\n        \"build\", help=\"Build model images only (no predict)\"\n    )\n    _add_common_args(build_parser)\n\n    # ── list ────────────────────────────────────────────────────────\n    list_parser = subparsers.add_parser(\"list\", help=\"List models in manifest\")\n    list_parser.add_argument(\n        \"--manifest\",\n        type=str,\n        default=None,\n        help=\"Path to manifest.yaml\",\n    )\n\n    args = parser.parse_args(argv)\n\n    if args.command is None:\n        parser.print_help()\n        sys.exit(1)\n\n    logging.basicConfig(\n        format=\"%(asctime)s %(levelname)-8s %(message)s\",\n        level=logging.INFO,\n        datefmt=\"%H:%M:%S\",\n    )\n\n    if args.command == \"list\":\n        _cmd_list(args)\n    elif args.command == \"build\":\n        _cmd_build(args)\n    elif args.command == \"run\":\n        _cmd_run(args)\n\n\ndef _cmd_list(args: argparse.Namespace) -> None:\n    manifest = _load_manifest(args.manifest)\n    models = manifest.get(\"models\", [])\n\n    for m in models:\n        gpu_tag = \" [GPU]\" if m.get(\"gpu\") else \"\"\n        req_env = m.get(\"requires_env\", [])\n        env_tag = f\" (requires: {', '.join(req_env)})\" if req_env else \"\"\n        print(f\"  {m['name']:<25} {m['repo']}/{m.get('path', '.')}{gpu_tag}{env_tag}\")\n\n    print(f\"\\n{len(models)} models total\")\n\n\ndef _cmd_build(args: argparse.Namespace) -> None:\n    manifest = _load_manifest(args.manifest)\n    models = _filter_models(manifest, args)\n    defaults = manifest.get(\"defaults\", {})\n\n    sdk_version, _ = resolve_sdk_version(\n        cli_sdk_version=args.sdk_version,\n        manifest_defaults=defaults,\n    )\n    cog_binary, cog_version_label = resolve_cog_binary(\n        cog_version=args.cog_version,\n        cog_binary=args.cog_binary,\n        manifest_defaults=defaults,\n    )\n    log = logging.getLogger(__name__)\n    log.info(\"Using cog CLI: %s (%s)\", cog_binary, cog_version_label)\n    log.info(\"Using SDK version: %s\", sdk_version)\n\n    runner = Runner(\n        cog_binary=cog_binary,\n        sdk_version=sdk_version,\n        keep_images=True,\n    )\n\n    results: list[ModelResult] = []\n    for model in models:\n        result = ModelResult(\n            name=model[\"name\"], passed=True, gpu=model.get(\"gpu\", False)\n        )\n        try:\n            model_dir = runner.prepare_model(model)\n            import time\n\n            start = time.monotonic()\n            runner.build_model(model_dir, model)\n            result.build_duration_s = time.monotonic() - start\n            logging.getLogger(__name__).info(\n                \"BUILD OK %s (%.1fs)\", model[\"name\"], result.build_duration_s\n            )\n        except Exception as exc:\n            result.passed = False\n            result.error = str(exc)\n\n        results.append(result)\n\n    console_report(\n        results, sdk_version=sdk_version or \"\", cog_version=cog_version_label\n    )\n\n    failed = any(not r.passed for r in results)\n    sys.exit(1 if failed else 0)\n\n\ndef _cmd_run(args: argparse.Namespace) -> None:\n    manifest = _load_manifest(args.manifest)\n    models = _filter_models(manifest, args)\n    defaults = manifest.get(\"defaults\", {})\n\n    sdk_version, _ = resolve_sdk_version(\n        cli_sdk_version=args.sdk_version,\n        manifest_defaults=defaults,\n    )\n    cog_binary, cog_version_label = resolve_cog_binary(\n        cog_version=args.cog_version,\n        cog_binary=args.cog_binary,\n        manifest_defaults=defaults,\n    )\n    log = logging.getLogger(__name__)\n    log.info(\"Using cog CLI: %s (%s)\", cog_binary, cog_version_label)\n    log.info(\"Using SDK version: %s\", sdk_version)\n\n    runner = Runner(\n        cog_binary=cog_binary,\n        sdk_version=sdk_version,\n        keep_images=args.keep_images,\n    )\n\n    results: list[ModelResult] = []\n    try:\n        for model in models:\n            result = runner.run_model(model)\n            results.append(result)\n    finally:\n        if not args.keep_images:\n            runner.cleanup()\n\n    # Output\n    if args.output == \"json\":\n        if args.output_file:\n            with open(args.output_file, \"w\") as f:\n                write_json_report(\n                    results,\n                    sdk_version=sdk_version or \"\",\n                    cog_version=cog_version_label,\n                    stream=f,\n                )\n        else:\n            write_json_report(\n                results,\n                sdk_version=sdk_version or \"\",\n                cog_version=cog_version_label,\n            )\n    else:\n        console_report(\n            results,\n            sdk_version=sdk_version or \"\",\n            cog_version=cog_version_label,\n        )\n        if args.output_file:\n            with open(args.output_file, \"w\") as f:\n                write_json_report(\n                    results,\n                    sdk_version=sdk_version or \"\",\n                    cog_version=cog_version_label,\n                    stream=f,\n                )\n\n    failed = any(not r.passed for r in results)\n    sys.exit(1 if failed else 0)\n\n\n# ── Helpers ────────────────────────────────────────────────────────────\n\n\ndef _add_common_args(parser: argparse.ArgumentParser) -> None:\n    parser.add_argument(\n        \"--manifest\",\n        type=str,\n        default=None,\n        help=\"Path to manifest.yaml (default: auto-detect)\",\n    )\n    parser.add_argument(\n        \"--model\",\n        type=str,\n        action=\"append\",\n        default=None,\n        help=\"Run only specific model(s) by name (repeatable)\",\n    )\n    parser.add_argument(\n        \"--no-gpu\",\n        action=\"store_true\",\n        help=\"Skip models that require a GPU\",\n    )\n    parser.add_argument(\n        \"--gpu-only\",\n        action=\"store_true\",\n        help=\"Only run models that require a GPU\",\n    )\n    parser.add_argument(\n        \"--sdk-version\",\n        type=str,\n        default=None,\n        help=(\n            \"SDK version to inject into cog.yaml (e.g. 0.16.12). \"\n            \"Default: latest stable release from PyPI.\"\n        ),\n    )\n    parser.add_argument(\n        \"--cog-version\",\n        type=str,\n        default=None,\n        help=(\n            \"Cog CLI version to download and use (e.g. v0.16.12). \"\n            \"Default: latest stable release. Ignored if --cog-binary is set.\"\n        ),\n    )\n    parser.add_argument(\n        \"--cog-binary\",\n        type=str,\n        default=\"cog\",\n        help=\"Path to a local cog binary (overrides --cog-version)\",\n    )\n    parser.add_argument(\n        \"--keep-images\",\n        action=\"store_true\",\n        help=\"Don't clean up Docker images after run\",\n    )\n\n\ndef _load_manifest(manifest_path: str | None) -> dict:\n    if manifest_path:\n        path = Path(manifest_path)\n    else:\n        # Search up from CWD, then fall back to the default location\n        path = Path(__file__).parent.parent / \"manifest.yaml\"\n\n    if not path.exists():\n        print(f\"Error: manifest not found at {path}\", file=sys.stderr)\n        sys.exit(1)\n\n    with open(path) as f:\n        return yaml.safe_load(f)\n\n\ndef _filter_models(manifest: dict, args: argparse.Namespace) -> list[dict]:\n    models = manifest.get(\"models\", [])\n\n    if args.model:\n        names = set(args.model)\n        models = [m for m in models if m[\"name\"] in names]\n        found = {m[\"name\"] for m in models}\n        missing = names - found\n        if missing:\n            print(f\"Warning: models not found in manifest: {missing}\", file=sys.stderr)\n\n    if getattr(args, \"no_gpu\", False):\n        models = [m for m in models if not m.get(\"gpu\")]\n\n    if getattr(args, \"gpu_only\", False):\n        models = [m for m in models if m.get(\"gpu\")]\n\n    return models\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tools/test-harness/harness/cog_resolver.py",
    "content": "\"\"\"Resolve and download specific cog CLI and SDK versions.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport logging\nimport os\nimport platform\nimport re\nimport stat\nimport tempfile\nimport urllib.request\nfrom pathlib import Path\n\nlogger = logging.getLogger(__name__)\n\nGITHUB_API = \"https://api.github.com/repos/replicate/cog/releases\"\nDOWNLOAD_BASE = (\n    \"https://github.com/replicate/cog/releases/download/{tag}/cog_{os}_{arch}\"\n)\nPYPI_API = \"https://pypi.org/pypi/cog/json\"\n\n# Pre-release patterns to skip when resolving \"latest\"\n_PRERELEASE_RE = re.compile(r\"-(alpha|beta|rc|dev)\", re.IGNORECASE)\n\n\ndef resolve_latest_stable_version() -> str:\n    \"\"\"Query GitHub releases and return the tag of the latest stable release.\n\n    Skips any release marked as a prerelease or whose tag contains\n    alpha/beta/rc/dev suffixes.\n    \"\"\"\n    url = f\"{GITHUB_API}?per_page=50\"\n    headers = {\"Accept\": \"application/vnd.github+json\"}\n\n    # Use a token if available to avoid rate limits\n    token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\n    if token:\n        headers[\"Authorization\"] = f\"Bearer {token}\"\n\n    req = urllib.request.Request(url, headers=headers)\n    with urllib.request.urlopen(req, timeout=30) as resp:\n        releases = json.loads(resp.read().decode())\n\n    for release in releases:\n        tag = release.get(\"tag_name\", \"\")\n        if release.get(\"prerelease\") or release.get(\"draft\"):\n            continue\n        if _PRERELEASE_RE.search(tag):\n            continue\n        return tag\n\n    raise RuntimeError(\n        \"Could not find a stable cog release. \"\n        \"Check https://github.com/replicate/cog/releases\"\n    )\n\n\ndef _platform_asset_name() -> str:\n    \"\"\"Return the cog binary asset name for the current platform.\"\"\"\n    system = platform.system()  # Darwin, Linux\n    machine = platform.machine()  # arm64, x86_64, aarch64\n\n    if system not in (\"Darwin\", \"Linux\"):\n        raise RuntimeError(f\"Unsupported OS: {system}\")\n\n    # Normalise architecture names\n    arch_map = {\n        \"arm64\": \"arm64\",\n        \"aarch64\": \"arm64\",\n        \"x86_64\": \"x86_64\",\n        \"amd64\": \"x86_64\",\n    }\n    arch = arch_map.get(machine)\n    if not arch:\n        raise RuntimeError(f\"Unsupported architecture: {machine}\")\n\n    return f\"cog_{system}_{arch}\"\n\n\ndef download_cog_binary(tag: str, dest_dir: Path | None = None) -> Path:\n    \"\"\"Download the cog binary for *tag* and return the path to it.\n\n    The binary is placed in *dest_dir* (default: a new temp directory) and\n    made executable.\n    \"\"\"\n    asset = _platform_asset_name()\n    url = DOWNLOAD_BASE.format(tag=tag, os=platform.system(), arch=asset.split(\"_\")[-1])\n\n    if dest_dir is None:\n        dest_dir = Path(tempfile.mkdtemp(prefix=\"cog-bin-\"))\n    dest_dir.mkdir(parents=True, exist_ok=True)\n\n    dest = dest_dir / \"cog\"\n\n    logger.info(\"Downloading cog %s from %s ...\", tag, url)\n\n    req = urllib.request.Request(url)\n    token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\n    if token:\n        req.add_header(\"Authorization\", f\"Bearer {token}\")\n\n    with urllib.request.urlopen(req, timeout=120) as resp, open(dest, \"wb\") as f:\n        while True:\n            chunk = resp.read(1 << 16)\n            if not chunk:\n                break\n            f.write(chunk)\n\n    # Make executable\n    dest.chmod(dest.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)\n\n    # Verify it works\n    logger.info(\"Downloaded cog %s -> %s\", tag, dest)\n    return dest\n\n\ndef resolve_cog_binary(\n    cog_version: str | None,\n    cog_binary: str | None,\n    manifest_defaults: dict | None = None,\n) -> tuple[str, str]:\n    \"\"\"Resolve which cog binary to use. Returns ``(binary_path, version_label)``.\n\n    Priority:\n    1. ``--cog-binary`` (explicit path) — use as-is, version label = \"custom\"\n    2. ``--cog-version`` — download that specific tag\n    3. ``defaults.cog_version`` from manifest — download that tag\n    4. No version specified — resolve latest stable, download it\n\n    If *cog_binary* is provided and is not the default ``\"cog\"``, it takes\n    top priority (the user wants their own binary).\n    \"\"\"\n    defaults = manifest_defaults or {}\n\n    # 1. Explicit --cog-binary (non-default)\n    if cog_binary and cog_binary != \"cog\":\n        return cog_binary, \"custom\"\n\n    # 2. Explicit --cog-version\n    if cog_version:\n        tag = cog_version if cog_version.startswith(\"v\") else f\"v{cog_version}\"\n        path = download_cog_binary(tag)\n        return str(path), tag\n\n    # 3. Manifest default\n    manifest_version = defaults.get(\"cog_version\")\n    if manifest_version and manifest_version != \"latest\":\n        tag = (\n            manifest_version\n            if manifest_version.startswith(\"v\")\n            else f\"v{manifest_version}\"\n        )\n        path = download_cog_binary(tag)\n        return str(path), tag\n\n    # 4. Resolve latest stable\n    tag = resolve_latest_stable_version()\n    logger.info(\"Resolved latest stable cog version: %s\", tag)\n    path = download_cog_binary(tag)\n    return str(path), tag\n\n\n# ── SDK version resolution ─────────────────────────────────────────────\n\n\ndef resolve_latest_sdk_version() -> str:\n    \"\"\"Query PyPI and return the latest stable version of the ``cog`` package.\n\n    PyPI's ``info.version`` field always returns the latest non-prerelease\n    version, so no extra filtering is needed.\n    \"\"\"\n    req = urllib.request.Request(PYPI_API, headers={\"Accept\": \"application/json\"})\n    with urllib.request.urlopen(req, timeout=30) as resp:\n        data = json.loads(resp.read().decode())\n    version = data[\"info\"][\"version\"]\n    logger.info(\"Resolved latest stable SDK version from PyPI: %s\", version)\n    return version\n\n\ndef resolve_sdk_version(\n    cli_sdk_version: str | None,\n    manifest_defaults: dict | None = None,\n) -> tuple[str, bool]:\n    \"\"\"Resolve which SDK version to use. Returns ``(version, was_resolved)``.\n\n    Priority:\n    1. ``--sdk-version`` CLI flag — use as-is\n    2. ``defaults.sdk_version`` from manifest (if not ``\"latest\"``)\n    3. Resolve latest stable from PyPI\n\n    *was_resolved* is ``True`` when the version was auto-resolved from PyPI.\n    \"\"\"\n    defaults = manifest_defaults or {}\n\n    # 1. Explicit --sdk-version\n    if cli_sdk_version:\n        return cli_sdk_version, False\n\n    # 2. Manifest default\n    manifest_version = defaults.get(\"sdk_version\")\n    if manifest_version and manifest_version != \"latest\":\n        return manifest_version, False\n\n    # 3. Resolve latest stable from PyPI\n    version = resolve_latest_sdk_version()\n    return version, True\n"
  },
  {
    "path": "tools/test-harness/harness/patcher.py",
    "content": "\"\"\"Patch cog.yaml files with sdk_version and arbitrary overrides.\"\"\"\n\nfrom __future__ import annotations\n\nimport copy\nfrom pathlib import Path\nfrom typing import Any\n\nimport yaml\n\n\ndef deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:\n    \"\"\"Recursively merge *override* into *base*, returning a new dict.\"\"\"\n    result = copy.deepcopy(base)\n    for key, value in override.items():\n        if key in result and isinstance(result[key], dict) and isinstance(value, dict):\n            result[key] = deep_merge(result[key], value)\n        else:\n            result[key] = copy.deepcopy(value)\n    return result\n\n\ndef patch_cog_yaml(\n    cog_yaml_path: Path,\n    sdk_version: str | None = None,\n    overrides: dict[str, Any] | None = None,\n) -> dict[str, Any]:\n    \"\"\"Read a cog.yaml, apply patches, write it back, and return the final config.\n\n    Parameters\n    ----------\n    cog_yaml_path:\n        Path to the cog.yaml file to patch (modified in-place).\n    sdk_version:\n        If set, inject ``build.sdk_version`` into the config.\n    overrides:\n        Arbitrary dict that is deep-merged into the config.  Useful for\n        changing python_version, adding system_packages, etc.\n\n    Returns\n    -------\n    The patched config dict.\n    \"\"\"\n    with open(cog_yaml_path) as f:\n        config = yaml.safe_load(f) or {}\n\n    if sdk_version:\n        config.setdefault(\"build\", {})\n        config[\"build\"][\"sdk_version\"] = sdk_version\n\n    if overrides:\n        config = deep_merge(config, overrides)\n\n    with open(cog_yaml_path, \"w\") as f:\n        yaml.dump(config, f, default_flow_style=False, sort_keys=False)\n\n    return config\n"
  },
  {
    "path": "tools/test-harness/harness/report.py",
    "content": "\"\"\"Generate human-readable and machine-readable test reports.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport sys\nfrom datetime import datetime, timezone\nfrom typing import Any, TextIO\n\nfrom .runner import ModelResult\n\n\ndef console_report(\n    results: list[ModelResult],\n    *,\n    sdk_version: str = \"\",\n    cog_version: str = \"\",\n    stream: TextIO = sys.stdout,\n) -> None:\n    \"\"\"Print a coloured summary table to the terminal.\"\"\"\n    parts = []\n    if cog_version:\n        parts.append(f\"CLI {cog_version}\")\n    if sdk_version:\n        parts.append(f\"SDK {sdk_version}\")\n    version_str = \" / \".join(parts)\n    header = (\n        f\"Cog Compatibility Report ({version_str})\"\n        if version_str\n        else \"Cog Compatibility Report\"\n    )\n    stream.write(f\"\\n{'=' * len(header)}\\n\")\n    stream.write(f\"{header}\\n\")\n    stream.write(f\"{'=' * len(header)}\\n\\n\")\n\n    passed = 0\n    failed = 0\n    skipped = 0\n\n    for r in results:\n        if r.skipped:\n            _write(stream, \"SKIP\", r.name, r.skip_reason or \"\", gpu=r.gpu)\n            skipped += 1\n            continue\n\n        if r.error:\n            _write(stream, \"FAIL\", r.name, r.error.splitlines()[0], gpu=r.gpu)\n            failed += 1\n            continue\n\n        all_tests = r.test_results + r.train_results\n        if r.passed:\n            timing = _timing_str(r.build_duration_s, all_tests)\n            _write(stream, \"PASS\", r.name, timing, gpu=r.gpu)\n            passed += 1\n        else:\n            failures = [t for t in all_tests if not t.passed]\n            msg = f\"{len(failures)} test(s) failed\"\n            if failures:\n                msg += f\": {failures[0].message[:60]}\"\n            _write(stream, \"FAIL\", r.name, msg, gpu=r.gpu)\n            failed += 1\n\n            # Print individual test failures indented\n            for t in failures:\n                stream.write(f\"    FAIL {t.description}: {t.message[:100]}\\n\")\n\n    stream.write(f\"\\n{'-' * 40}\\n\")\n    total = passed + failed + skipped\n    stream.write(f\"{passed}/{total} passed\")\n    if skipped:\n        stream.write(f\", {skipped} skipped\")\n    if failed:\n        stream.write(f\", {failed} FAILED\")\n    stream.write(\"\\n\\n\")\n\n\ndef json_report(\n    results: list[ModelResult],\n    *,\n    sdk_version: str = \"\",\n    cog_version: str = \"\",\n) -> dict[str, Any]:\n    \"\"\"Return a JSON-serializable report dict.\"\"\"\n    models = []\n    for r in results:\n        entry: dict[str, Any] = {\n            \"name\": r.name,\n            \"passed\": r.passed,\n            \"skipped\": r.skipped,\n            \"gpu\": r.gpu,\n            \"build_duration_s\": round(r.build_duration_s, 2),\n        }\n        if r.skipped:\n            entry[\"skip_reason\"] = r.skip_reason\n        if r.error:\n            entry[\"error\"] = r.error\n        if r.test_results:\n            entry[\"tests\"] = [\n                {\n                    \"description\": t.description,\n                    \"passed\": t.passed,\n                    \"message\": t.message,\n                    \"duration_s\": round(t.duration_s, 2),\n                }\n                for t in r.test_results\n            ]\n        if r.train_results:\n            entry[\"train_tests\"] = [\n                {\n                    \"description\": t.description,\n                    \"passed\": t.passed,\n                    \"message\": t.message,\n                    \"duration_s\": round(t.duration_s, 2),\n                }\n                for t in r.train_results\n            ]\n        models.append(entry)\n\n    total = len(results)\n    passed = sum(1 for r in results if r.passed and not r.skipped)\n    failed = sum(1 for r in results if not r.passed)\n    skipped_count = sum(1 for r in results if r.skipped)\n\n    return {\n        \"timestamp\": datetime.now(timezone.utc).isoformat(),\n        \"cog_version\": cog_version,\n        \"sdk_version\": sdk_version,\n        \"summary\": {\n            \"total\": total,\n            \"passed\": passed,\n            \"failed\": failed,\n            \"skipped\": skipped_count,\n        },\n        \"models\": models,\n    }\n\n\ndef write_json_report(\n    results: list[ModelResult],\n    *,\n    sdk_version: str = \"\",\n    cog_version: str = \"\",\n    stream: TextIO = sys.stdout,\n) -> None:\n    \"\"\"Write JSON report to a stream.\"\"\"\n    report = json_report(results, sdk_version=sdk_version, cog_version=cog_version)\n    json.dump(report, stream, indent=2)\n    stream.write(\"\\n\")\n\n\n# ── Helpers ────────────────────────────────────────────────────────────\n\n\ndef _write(\n    stream: TextIO, status: str, name: str, detail: str, *, gpu: bool = False\n) -> None:\n    icon = {\"PASS\": \"+\", \"FAIL\": \"x\", \"SKIP\": \"-\"}[status]\n    gpu_tag = \" [GPU]\" if gpu else \"\"\n    stream.write(f\"  {icon} {name:<25} {detail}{gpu_tag}\\n\")\n\n\ndef _timing_str(build_s: float, tests: list[Any]) -> str:\n    parts = [f\"{build_s:.1f}s build\"]\n    if tests:\n        total_predict = sum(t.duration_s for t in tests)\n        parts.append(f\"{total_predict:.1f}s predict\")\n    return f\"({', '.join(parts)})\"\n"
  },
  {
    "path": "tools/test-harness/harness/runner.py",
    "content": "\"\"\"Core test runner: clone, patch, build, predict, validate.\"\"\"\n\nfrom __future__ import annotations\n\nimport logging\nimport os\nimport re\nimport shutil\nimport subprocess\nimport tempfile\nimport time\nfrom dataclasses import dataclass, field\nfrom pathlib import Path\nfrom typing import Any\n\nfrom .patcher import patch_cog_yaml\nfrom .validators import ValidationResult, validate\n\nlogger = logging.getLogger(__name__)\n\n# ── Data types ─────────────────────────────────────────────────────────\n\n\n@dataclass\nclass TestCaseResult:\n    description: str\n    passed: bool\n    message: str\n    duration_s: float = 0.0\n\n\n@dataclass\nclass ModelResult:\n    name: str\n    passed: bool\n    build_duration_s: float = 0.0\n    test_results: list[TestCaseResult] = field(default_factory=list)\n    train_results: list[TestCaseResult] = field(default_factory=list)\n    error: str | None = None\n    skipped: bool = False\n    skip_reason: str | None = None\n    gpu: bool = False\n\n\n# ── Runner ─────────────────────────────────────────────────────────────\n\n\nclass Runner:\n    \"\"\"Orchestrates the clone -> patch -> build -> predict -> validate cycle.\"\"\"\n\n    def __init__(\n        self,\n        *,\n        cog_binary: str = \"cog\",\n        sdk_version: str | None = None,\n        fixtures_dir: Path | None = None,\n        work_dir: Path | None = None,\n        keep_images: bool = False,\n        default_timeout: int = 300,\n    ) -> None:\n        self.cog_binary = cog_binary\n        self.sdk_version = sdk_version\n        self.fixtures_dir = fixtures_dir or Path(__file__).parent.parent / \"fixtures\"\n        self.work_dir = work_dir or Path(tempfile.mkdtemp(prefix=\"cog-harness-\"))\n        self.keep_images = keep_images\n        self.default_timeout = default_timeout\n        self._cloned_repos: dict[str, Path] = {}\n\n    def prepare_model(self, model: dict[str, Any]) -> Path:\n        \"\"\"Public wrapper around model preparation (clone + patch).\"\"\"\n        return self._prepare_model(model)\n\n    def build_model(self, model_dir: Path, model: dict[str, Any]) -> None:\n        \"\"\"Public wrapper around ``cog build``.\"\"\"\n        self._cog_build(model_dir, model)\n\n    def run_model(self, model: dict[str, Any]) -> ModelResult:\n        \"\"\"Run all tests for a single model definition from the manifest.\"\"\"\n        name = model[\"name\"]\n        gpu = model.get(\"gpu\", False)\n        result = ModelResult(name=name, passed=True, gpu=gpu)\n\n        # Check required env vars\n        required_env = model.get(\"requires_env\", [])\n        missing = [v for v in required_env if not os.environ.get(v)]\n        if missing:\n            result.passed = True  # not a failure, just skipped\n            result.skipped = True\n            result.skip_reason = f\"Missing env vars: {', '.join(missing)}\"\n            logger.info(\"SKIP %s: %s\", name, result.skip_reason)\n            return result\n\n        try:\n            model_dir = self._prepare_model(model)\n        except Exception as exc:\n            result.passed = False\n            result.error = f\"Preparation failed: {exc}\"\n            logger.error(\"FAIL %s: %s\", name, result.error)\n            return result\n\n        # Build\n        build_start = time.monotonic()\n        try:\n            self._cog_build(model_dir, model)\n            result.build_duration_s = time.monotonic() - build_start\n            logger.info(\"BUILD OK %s (%.1fs)\", name, result.build_duration_s)\n        except subprocess.CalledProcessError as exc:\n            result.passed = False\n            result.build_duration_s = time.monotonic() - build_start\n            stderr = exc.stderr or \"\"\n            result.error = f\"Build failed:\\n{stderr[-2000:]}\"\n            logger.error(\"BUILD FAIL %s:\\n%s\", name, stderr[-500:])\n            return result\n\n        # Train tests\n        for tc in model.get(\"train_tests\", []):\n            tc_result = self._run_train_test(model_dir, model, tc)\n            result.train_results.append(tc_result)\n            if not tc_result.passed:\n                result.passed = False\n\n        # Predict tests\n        for tc in model.get(\"tests\", []):\n            tc_result = self._run_predict_test(model_dir, model, tc)\n            result.test_results.append(tc_result)\n            if not tc_result.passed:\n                result.passed = False\n\n        return result\n\n    # ── Internal helpers ───────────────────────────────────────────────\n\n    def _prepare_model(self, model: dict[str, Any]) -> Path:\n        \"\"\"Clone the repo (if needed) and patch cog.yaml. Returns model dir.\"\"\"\n        repo = model[\"repo\"]\n        subpath = model.get(\"path\", \".\")\n\n        repo_dir = self._clone_repo(repo)\n        model_dir = repo_dir / subpath\n\n        if not (model_dir / \"cog.yaml\").exists():\n            raise FileNotFoundError(f\"No cog.yaml in {model_dir}\")\n\n        sdk_version = model.get(\"sdk_version\", self.sdk_version)\n        overrides = model.get(\"cog_yaml_overrides\")\n\n        patch_cog_yaml(\n            model_dir / \"cog.yaml\",\n            sdk_version=sdk_version,\n            overrides=overrides,\n        )\n\n        return model_dir\n\n    def _clone_repo(self, repo: str) -> Path:\n        \"\"\"Shallow-clone a GitHub repo into the work dir, caching by repo name.\"\"\"\n        if repo in self._cloned_repos:\n            return self._cloned_repos[repo]\n\n        dest = self.work_dir / repo.replace(\"/\", \"--\")\n        if dest.exists():\n            shutil.rmtree(dest)\n\n        url = f\"https://github.com/{repo}.git\"\n        logger.info(\"Cloning %s ...\", url)\n        subprocess.run(\n            [\"git\", \"clone\", \"--depth=1\", url, str(dest)],\n            check=True,\n            capture_output=True,\n            text=True,\n        )\n        self._cloned_repos[repo] = dest\n        return dest\n\n    def _cog_build(self, model_dir: Path, model: dict[str, Any]) -> None:\n        \"\"\"Run ``cog build`` in the model directory.\"\"\"\n        image_tag = f\"cog-harness-{model['name']}:test\"\n        cmd = [self.cog_binary, \"build\", \"-t\", image_tag]\n\n        env = self._build_env(model)\n        timeout = model.get(\"timeout\", self.default_timeout)\n\n        subprocess.run(\n            cmd,\n            cwd=model_dir,\n            check=True,\n            capture_output=True,\n            text=True,\n            env=env,\n            timeout=timeout,\n        )\n\n    def _run_predict_test(\n        self, model_dir: Path, model: dict[str, Any], tc: dict[str, Any]\n    ) -> TestCaseResult:\n        description = tc.get(\"description\", \"predict\")\n        start = time.monotonic()\n\n        cmd = [self.cog_binary, \"predict\"]\n        for key, value in tc.get(\"inputs\", {}).items():\n            resolved = self._resolve_input(value)\n            cmd.extend([\"-i\", f\"{key}={resolved}\"])\n\n        env = self._build_env(model)\n        timeout = model.get(\"timeout\", self.default_timeout)\n\n        try:\n            proc = subprocess.run(\n                cmd,\n                cwd=model_dir,\n                capture_output=True,\n                text=True,\n                env=env,\n                timeout=timeout,\n            )\n            duration = time.monotonic() - start\n\n            if proc.returncode != 0:\n                return TestCaseResult(\n                    description=description,\n                    passed=False,\n                    message=f\"cog predict exited {proc.returncode}:\\n{proc.stderr[-1000:]}\",\n                    duration_s=duration,\n                )\n\n            output = self._extract_output(proc, model_dir)\n            vr: ValidationResult = validate(output, tc.get(\"expect\", {}))\n            logger.info(\n                \"  %s %s: %s (%.1fs)\",\n                \"PASS\" if vr.passed else \"FAIL\",\n                description,\n                vr.message[:80],\n                duration,\n            )\n            return TestCaseResult(\n                description=description,\n                passed=vr.passed,\n                message=vr.message,\n                duration_s=duration,\n            )\n\n        except subprocess.TimeoutExpired:\n            duration = time.monotonic() - start\n            return TestCaseResult(\n                description=description,\n                passed=False,\n                message=f\"Timed out after {timeout}s\",\n                duration_s=duration,\n            )\n        except Exception as exc:\n            duration = time.monotonic() - start\n            return TestCaseResult(\n                description=description,\n                passed=False,\n                message=f\"Unexpected error: {exc}\",\n                duration_s=duration,\n            )\n\n    def _run_train_test(\n        self, model_dir: Path, model: dict[str, Any], tc: dict[str, Any]\n    ) -> TestCaseResult:\n        description = tc.get(\"description\", \"train\")\n        start = time.monotonic()\n\n        cmd = [self.cog_binary, \"train\"]\n        for key, value in tc.get(\"inputs\", {}).items():\n            resolved = self._resolve_input(value)\n            cmd.extend([\"-i\", f\"{key}={resolved}\"])\n\n        env = self._build_env(model)\n        timeout = model.get(\"timeout\", self.default_timeout)\n\n        try:\n            proc = subprocess.run(\n                cmd,\n                cwd=model_dir,\n                capture_output=True,\n                text=True,\n                env=env,\n                timeout=timeout,\n            )\n            duration = time.monotonic() - start\n\n            if proc.returncode != 0:\n                return TestCaseResult(\n                    description=description,\n                    passed=False,\n                    message=f\"cog train exited {proc.returncode}:\\n{proc.stderr[-1000:]}\",\n                    duration_s=duration,\n                )\n\n            output = self._extract_output(proc, model_dir)\n            vr: ValidationResult = validate(output, tc.get(\"expect\", {}))\n            logger.info(\n                \"  %s %s: %s (%.1fs)\",\n                \"PASS\" if vr.passed else \"FAIL\",\n                description,\n                vr.message[:80],\n                duration,\n            )\n            return TestCaseResult(\n                description=description,\n                passed=vr.passed,\n                message=vr.message,\n                duration_s=duration,\n            )\n\n        except subprocess.TimeoutExpired:\n            duration = time.monotonic() - start\n            return TestCaseResult(\n                description=description,\n                passed=False,\n                message=f\"Timed out after {timeout}s\",\n                duration_s=duration,\n            )\n        except Exception as exc:\n            duration = time.monotonic() - start\n            return TestCaseResult(\n                description=description,\n                passed=False,\n                message=f\"Unexpected error: {exc}\",\n                duration_s=duration,\n            )\n\n    @staticmethod\n    def _extract_output(proc: subprocess.CompletedProcess[str], model_dir: Path) -> str:\n        \"\"\"Extract the prediction output from cog's stdout/stderr.\n\n        ``cog predict`` prints text/JSON output to **stdout**.  For file\n        outputs (e.g. images) it writes the file to the CWD and prints\n        ``Written output to: `` on **stderr**.  We detect the\n        latter pattern and return the absolute path to the file so that\n        the ``file_exists`` validator can verify it.\n        \"\"\"\n\n        # If there's meaningful stdout, prefer that\n        stdout = proc.stdout.strip()\n        if stdout:\n            return proc.stdout\n\n        # Check stderr for \"Written output to: \"\n        m = re.search(r\"Written output to:\\s*(.+)\", proc.stderr)\n        if m:\n            rel_path = m.group(1).strip()\n            abs_path = model_dir / rel_path\n            return str(abs_path)\n\n        # Fallback: return whatever stdout had (possibly empty)\n        return proc.stdout\n\n    def _resolve_input(self, value: Any) -> str:\n        \"\"\"Resolve input values — ``@filename`` becomes an absolute fixture path.\n\n        The path is resolved to an absolute, canonical path (no symlinks or\n        ``..`` components) so that ``cog predict -i image=@/abs/path`` works\n        correctly when cog mounts the file into the container.\n        \"\"\"\n        s = str(value)\n        if s.startswith(\"@\"):\n            fixture_path = (self.fixtures_dir / s[1:]).resolve()\n            if not fixture_path.exists():\n                raise FileNotFoundError(\n                    f\"Fixture not found: {fixture_path} (referenced as {s!r})\"\n                )\n            return f\"@{fixture_path}\"\n        return s\n\n    def _build_env(self, model: dict[str, Any]) -> dict[str, str]:\n        \"\"\"Build environment dict, expanding ${VAR} references from host env.\"\"\"\n        env = os.environ.copy()\n        for key, value in model.get(\"env\", {}).items():\n            resolved = os.path.expandvars(value)\n            env[key] = resolved\n        return env\n\n    def cleanup(self) -> None:\n        \"\"\"Remove work directory and optionally docker images.\"\"\"\n        if not self.keep_images:\n            # Clean up docker images we created\n            try:\n                proc = subprocess.run(\n                    [\n                        \"docker\",\n                        \"images\",\n                        \"--filter\",\n                        \"reference=cog-harness-*\",\n                        \"--format\",\n                        \"{{.Repository}}:{{.Tag}}\",\n                    ],\n                    capture_output=True,\n                    text=True,\n                )\n                images = [\n                    line.strip() for line in proc.stdout.splitlines() if line.strip()\n                ]\n                if images:\n                    subprocess.run(\n                        [\"docker\", \"rmi\", \"--force\"] + images,\n                        capture_output=True,\n                        text=True,\n                    )\n            except Exception as exc:\n                logger.warning(\"Failed to clean up Docker images in cleanup(): %s\", exc)\n\n        if self.work_dir.exists():\n            shutil.rmtree(self.work_dir, ignore_errors=True)\n"
  },
  {
    "path": "tools/test-harness/harness/validators.py",
    "content": "\"\"\"Output validation strategies for cog model predictions.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport mimetypes\nimport re\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Any\n\n\n@dataclass\nclass ValidationResult:\n    passed: bool\n    message: str\n\n\ndef validate(output: str, expect: dict[str, Any]) -> ValidationResult:\n    \"\"\"Dispatch to the appropriate validator based on ``expect[\"type\"]``.\"\"\"\n    vtype = expect.get(\"type\", \"not_empty\")\n    validator = _VALIDATORS.get(vtype)\n    if validator is None:\n        return ValidationResult(\n            passed=False,\n            message=f\"Unknown validation type: {vtype!r}\",\n        )\n    return validator(output, expect)\n\n\n# ── Individual validators ──────────────────────────────────────────────\n\n\ndef _validate_exact(output: str, expect: dict[str, Any]) -> ValidationResult:\n    expected = str(expect[\"value\"])\n    clean = output.strip()\n    if clean == expected:\n        return ValidationResult(passed=True, message=\"Exact match\")\n    return ValidationResult(\n        passed=False,\n        message=f\"Expected exact match:\\n  expected: {expected!r}\\n  got:      {clean!r}\",\n    )\n\n\ndef _validate_contains(output: str, expect: dict[str, Any]) -> ValidationResult:\n    substring = str(expect[\"value\"])\n    if substring in output:\n        return ValidationResult(passed=True, message=f\"Contains {substring!r}\")\n    return ValidationResult(\n        passed=False,\n        message=f\"Expected output to contain {substring!r}, got:\\n  {output[:200]!r}\",\n    )\n\n\ndef _validate_regex(output: str, expect: dict[str, Any]) -> ValidationResult:\n    pattern = expect[\"pattern\"]\n    if re.search(pattern, output):\n        return ValidationResult(passed=True, message=f\"Matches pattern {pattern!r}\")\n    return ValidationResult(\n        passed=False,\n        message=f\"Output does not match regex {pattern!r}:\\n  {output[:200]!r}\",\n    )\n\n\ndef _validate_file_exists(output: str, expect: dict[str, Any]) -> ValidationResult:\n    \"\"\"Validate that the output references an existing file.\n\n    ``cog predict`` prints the output file path to stdout.  It may be an\n    absolute path or a relative path.  We also handle the common case\n    where cog wraps the path in quotes or prints extra whitespace.\n    \"\"\"\n    path_str = output.strip().strip(\"'\\\"\")\n\n    # cog predict may output a URL or a path -- for local testing it's a path\n    if path_str.startswith(\"http://\") or path_str.startswith(\"https://\"):\n        # Can't verify remote files; treat as pass\n        return ValidationResult(passed=True, message=f\"Output is a URL: {path_str}\")\n\n    path = Path(path_str)\n    if not path.exists():\n        return ValidationResult(\n            passed=False,\n            message=f\"Output file does not exist: {path}\",\n        )\n\n    expected_mime = expect.get(\"mime\")\n    if expected_mime:\n        guessed, _ = mimetypes.guess_type(str(path))\n        if guessed != expected_mime:\n            return ValidationResult(\n                passed=False,\n                message=f\"Expected MIME {expected_mime}, got {guessed} for {path}\",\n            )\n\n    return ValidationResult(passed=True, message=f\"File exists: {path}\")\n\n\ndef _validate_json_match(output: str, expect: dict[str, Any]) -> ValidationResult:\n    \"\"\"Parse output as JSON and verify that ``expect[\"match\"]`` is a subset.\"\"\"\n    try:\n        parsed = json.loads(output.strip())\n    except json.JSONDecodeError as exc:\n        return ValidationResult(\n            passed=False,\n            message=f\"Output is not valid JSON: {exc}\\n  {output[:200]!r}\",\n        )\n\n    match = expect[\"match\"]\n    if not _is_subset(match, parsed):\n        return ValidationResult(\n            passed=False,\n            message=f\"JSON subset mismatch:\\n  expected subset: {match}\\n  got: {parsed}\",\n        )\n    return ValidationResult(passed=True, message=\"JSON subset match\")\n\n\ndef _validate_json_keys(output: str, expect: dict[str, Any]) -> ValidationResult:\n    \"\"\"Parse output as JSON dict and verify it has entries (non-empty).\"\"\"\n    try:\n        parsed = json.loads(output.strip())\n    except json.JSONDecodeError as exc:\n        return ValidationResult(\n            passed=False,\n            message=f\"Output is not valid JSON: {exc}\\n  {output[:200]!r}\",\n        )\n\n    if not isinstance(parsed, dict):\n        return ValidationResult(\n            passed=False,\n            message=f\"Expected JSON object, got {type(parsed).__name__}\",\n        )\n\n    required_keys = expect.get(\"keys\", [])\n    if required_keys:\n        missing = [k for k in required_keys if k not in parsed]\n        if missing:\n            return ValidationResult(\n                passed=False,\n                message=f\"Missing keys: {missing}. Got: {list(parsed.keys())}\",\n            )\n    elif not parsed:\n        return ValidationResult(\n            passed=False,\n            message=\"Expected non-empty JSON object, got empty dict\",\n        )\n\n    return ValidationResult(\n        passed=True,\n        message=f\"JSON dict with {len(parsed)} keys: {list(parsed.keys())[:5]}\",\n    )\n\n\ndef _validate_not_empty(output: str, _expect: dict[str, Any]) -> ValidationResult:\n    if output.strip():\n        return ValidationResult(passed=True, message=\"Output is non-empty\")\n    return ValidationResult(passed=False, message=\"Output is empty\")\n\n\n# ── Helpers ────────────────────────────────────────────────────────────\n\n\ndef _is_subset(subset: Any, superset: Any) -> bool:\n    \"\"\"Check that *subset* is recursively contained in *superset*.\"\"\"\n    if isinstance(subset, dict) and isinstance(superset, dict):\n        return all(\n            k in superset and _is_subset(v, superset[k]) for k, v in subset.items()\n        )\n    if isinstance(subset, list) and isinstance(superset, list):\n        return all(\n            any(_is_subset(s_item, p_item) for p_item in superset) for s_item in subset\n        )\n    return subset == superset\n\n\n# ── Registry ───────────────────────────────────────────────────────────\n\n_VALIDATORS = {\n    \"exact\": _validate_exact,\n    \"contains\": _validate_contains,\n    \"regex\": _validate_regex,\n    \"file_exists\": _validate_file_exists,\n    \"json_match\": _validate_json_match,\n    \"json_keys\": _validate_json_keys,\n    \"not_empty\": _validate_not_empty,\n}\n"
  },
  {
    "path": "tools/test-harness/manifest.yaml",
    "content": "# Cog Model Test Manifest\n# =======================\n# Each entry defines a model to test, its inputs, and expected outputs.\n#\n# Input values prefixed with \"@\" are resolved as fixture file paths relative\n# to the fixtures/ directory (e.g. \"@test_image.png\" -> fixtures/test_image.png).\n#\n# Validation types:\n#   exact       - output string must equal `value` exactly\n#   contains    - output string must contain `value` as a substring\n#   regex       - output string must match `pattern`\n#   file_exists - output is a file path; optionally check `mime` type\n#   json_match  - parse output as JSON, assert `match` is a subset\n#   json_keys   - parse output as JSON dict, assert it has entries\n#   not_empty   - output is non-empty (loose smoke test)\n\ndefaults:\n  sdk_version: \"latest\"          # \"latest\" = newest stable from PyPI; or pin e.g. \"0.16.12\"\n  cog_version: \"latest\"          # \"latest\" = newest stable release; or pin e.g. \"v0.16.12\"\n\nmodels:\n  # ── cog-examples (CPU) ──────────────────────────────────────────────\n\n  - name: hello-world\n    repo: replicate/cog-examples\n    path: hello-world\n    gpu: false\n    tests:\n      - description: \"basic predict\"\n        inputs:\n          text: \"world\"\n        expect:\n          type: exact\n          value: \"hello world\"\n\n  - name: canary\n    repo: replicate/cog-examples\n    path: canary\n    gpu: false\n    tests:\n      - description: \"streaming concatenate iterator\"\n        inputs:\n          text: \"friend\"\n        expect:\n          type: contains\n          value: \"friend\"\n\n  - name: blur\n    repo: replicate/cog-examples\n    path: blur\n    gpu: false\n    tests:\n      - description: \"blur an image\"\n        inputs:\n          image: \"@test_image.png\"\n          blur: 5\n        expect:\n          type: file_exists\n          mime: \"image/png\"\n\n  - name: hello-image\n    repo: replicate/cog-examples\n    path: hello-image\n    gpu: false\n    tests:\n      - description: \"return a static image\"\n        inputs: {}\n        expect:\n          type: file_exists\n\n  - name: hello-concurrency\n    repo: replicate/cog-examples\n    path: hello-concurrency\n    gpu: false\n    tests:\n      - description: \"async streaming output\"\n        inputs:\n          total: 3\n          interval: 0\n        expect:\n          type: contains\n          value: \"Apple\"\n\n  - name: hello-context\n    repo: replicate/cog-examples\n    path: hello-context\n    gpu: false\n    # NOTE: This model uses current_scope().context which may not be\n    # available in coglet yet. A failure here is a real compatibility signal.\n    tests:\n      - description: \"returns input and context\"\n        inputs:\n          text: \"testing\"\n        expect:\n          type: json_match\n          match:\n            inputs:\n              text: \"testing\"\n\n  - name: hello-train\n    repo: replicate/cog-examples\n    path: hello-train\n    gpu: false\n    # NOTE: `cog train` in the RC may have input validation issues\n    # (validates against predict schema instead of train schema).\n    # The train_test below may fail — that's a real compatibility signal.\n    train_tests:\n      - description: \"train produces weights file\"\n        inputs:\n          prefix: \"custom\"\n        expect:\n          type: not_empty\n    tests:\n      - description: \"predict with default weights\"\n        inputs:\n          text: \"world\"\n        expect:\n          type: contains\n          value: \"world\"\n\n  # ── cog-examples (GPU required) ─────────────────────────────────────\n\n  - name: resnet\n    repo: replicate/cog-examples\n    path: resnet\n    gpu: true\n    tests:\n      - description: \"classify hotdog image\"\n        inputs:\n          image: \"@hotdog.png\"\n        expect:\n          type: json_keys\n\n  - name: z-image-turbo\n    repo: replicate/cog-examples\n    path: z-image-turbo\n    gpu: true\n    timeout: 600\n    tests:\n      - description: \"generate image from prompt\"\n        inputs:\n          prompt: \"a cat sitting on a windowsill\"\n        expect:\n          type: file_exists\n          mime: \"image/png\"\n\n  # ── cog-examples (requires external API, optional) ──────────────────\n\n  - name: hello-replicate\n    repo: replicate/cog-examples\n    path: hello-replicate\n    gpu: false\n    requires_env:\n      - REPLICATE_API_TOKEN\n    tests:\n      - description: \"round-trip through replicate API\"\n        inputs:\n          image: \"@test_image.png\"\n        expect:\n          type: file_exists\n\n  # ── External models (add your own below) ────────────────────────────\n  # - name: my-custom-model\n  #   repo: myorg/my-model-repo\n  #   path: \".\"\n  #   gpu: true\n  #   # sdk_version: \"0.16.12\"  # optional per-model override\n  #   env:\n  #     HF_TOKEN: \"${HF_TOKEN}\"\n  #   timeout: 600\n  #   tests:\n  #     - description: \"smoke test\"\n  #       inputs:\n  #         prompt: \"hello\"\n  #       expect:\n  #         type: contains\n  #         value: \"result\"\n"
  },
  {
    "path": "tools/test-harness/pyproject.toml",
    "content": "[project]\nname = \"cog-test-harness\"\nversion = \"0.1.0\"\ndescription = \"Test harness for validating cog models against new SDK versions\"\nrequires-python = \">=3.10\"\ndependencies = [\n    \"pyyaml>=6.0\",\n]\n\n[project.scripts]\ncog-test = \"harness.cli:main\"\n"
  },
  {
    "path": "tools/test-harness/results/.gitkeep",
    "content": ""
  },
  {
    "path": "tools/test-registry-util/README.md",
    "content": "# `test-registry-util`\n\nA tool for creating and inspecting a local registry for testing. \n\n## Purpose\n\nWe have a lot of intricate image manipulation code that needs to be tested. Mocks are't great for this because we need to make sure the code works with actual data. This tool helps setup real data for a test registry.\n\n## Usage\n\nImage data is stored in `pkg/registry_testhelpers/testdata` and matches the structore expected by `distribution/distribution`. \n\nDuring tests an ephemeral registry is spun up on a random local port, populated with the image data, and turn down when the test finishes.\n\n### Booting a registry in a test:\n\n```go\nimport \"github.com/replicate/cog/pkg/registry_testhelpers\"\n\nfunc TestMyFunction(t *testing.T) {\n\tregistryContainer := registry_testhelpers.StartTestRegistry(ctx)\n  image := registryContainer.ImageRef(\"alpine:latest\")\n\t\n  // use image as a real image reference\n}\n```\n### Inspect the current images in the registry:\n\n```bash\ngo run ./tools/test-registry-util catalog\n```\nwill print something like:\n\n```\nalpine:latest application/vnd.oci.image.index.v1+json\n  index -> sha256:9a0ff41dccad7a96f324a4655a715c623ed3511c7336361ffa9dadcecbdb99e5\n  linux/amd64 -> sha256:1c4eef651f65e2f7daee7ee785882ac164b02b78fb74503052a26dc061c90474\n  linux/arm64 -> sha256:757d680068d77be46fd1ea20fb21db16f150468c5e7079a08a2e4705aec096ac\npython:3.10 application/vnd.oci.image.manifest.v1+json\n  single platform image -> sha256:f33bb19d5a518ba7e0353b6da48d58a04ef674de0bab0810e4751230ea1d4b19\n```\n\nYou can then use these images in your tests using references like:\n\n- `localhost:/alpine:latest` to get a multi-platform index\n- `localhost:/alpine:latest` with platform `linux/amd64` to get a single image from a multi-platform index\n- `localhost:/alpine:latest@sha256:1c4eef651f65e2f7daee7ee785882ac164b02b78fb74503052a26dc061c90474` to get a specific image\n- `localhost:/python:3.10` to get a single-platform image\n\n\n### Initialize a new registry storage\n\nTo create a new directory of images, run:\n\n```\ngo run ./tools/test-registry-util init\n```\n\nThis will download all the images specified in `main.go` and save them to `pkg/registry_testhelpers/testdata`.\n\n### Run a registry\n\nThis is just a convenience to inspect a registry outside of a test.\n\n```\ngo run ./tools/test-registry-util run\n```\n"
  },
  {
    "path": "tools/test-registry-util/main.go",
    "content": "package main\n\nimport (\n\t\"context\"\n\t\"fmt\"\n\t\"os\"\n\t\"os/signal\"\n\t\"path/filepath\"\n\t\"strings\"\n\t\"time\"\n\n\t\"github.com/docker/docker/api/types/container\"\n\t\"github.com/docker/docker/api/types/mount\"\n\t\"github.com/google/go-containerregistry/pkg/authn\"\n\t\"github.com/google/go-containerregistry/pkg/name\"\n\tv1 \"github.com/google/go-containerregistry/pkg/v1\"\n\t\"github.com/google/go-containerregistry/pkg/v1/empty\"\n\t\"github.com/google/go-containerregistry/pkg/v1/mutate\"\n\t\"github.com/google/go-containerregistry/pkg/v1/remote\"\n\t\"github.com/google/go-containerregistry/pkg/v1/types\"\n\t\"github.com/spf13/cobra\"\n\t\"github.com/testcontainers/testcontainers-go\"\n\t\"github.com/testcontainers/testcontainers-go/modules/registry\"\n\t\"github.com/testcontainers/testcontainers-go/wait\"\n\n\t\"github.com/replicate/cog/pkg/util/files\"\n)\n\n// images to download and push to the registry. Keep the images sizes small since they're stored in git.\n// For reference, the `alpine:latest` image for `linux/amd64` ~3.5MB compressed.\nvar images = []struct {\n\tImage          string\n\tPlatforms      []string\n\tSinglePlatform string\n}{\n\t{\n\t\tImage: \"alpine:latest\",\n\t\tPlatforms: []string{\n\t\t\t\"linux/amd64\",\n\t\t\t\"linux/arm64\",\n\t\t},\n\t},\n}\n\n// relative to the root of the repo\nvar destinationDir string = \"pkg/registry_testhelpers/testdata\"\n\nfunc main() {\n\trootCmd := &cobra.Command{\n\t\tUse: \"test-registry-util\",\n\t}\n\trootCmd.PersistentFlags().StringVar(&destinationDir, \"storage-dir\", destinationDir, \"path to the directory where the registry will store its data\")\n\n\trootCmd.AddCommand(\n\t\t&cobra.Command{\n\t\t\tUse: \"init\",\n\t\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\t\treturn runAndInit(cmd.Context(), destinationDir)\n\t\t\t},\n\t\t},\n\t)\n\trootCmd.AddCommand(\n\t\t&cobra.Command{\n\t\t\tUse: \"catalog\",\n\t\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\t\treturn runAndCatalog(cmd.Context(), destinationDir)\n\t\t\t},\n\t\t},\n\t)\n\n\trootCmd.AddCommand(\n\t\t&cobra.Command{\n\t\t\tUse: \"run\",\n\t\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\t\tctx, cancel := signal.NotifyContext(cmd.Context(), os.Interrupt)\n\t\t\t\tdefer cancel()\n\n\t\t\t\tc, port, err := startRegistryTC(cmd.Context(), destinationDir)\n\t\t\t\tif err != nil {\n\t\t\t\t\treturn err\n\t\t\t\t}\n\t\t\t\tdefer func() {\n\t\t\t\t\tif err := c.Terminate(cmd.Context()); err != nil {\n\t\t\t\t\t\tfmt.Println(\"Failed to terminate registry:\", err)\n\t\t\t\t\t}\n\t\t\t\t}()\n\n\t\t\t\tfmt.Println(\"Registry running at\", fmt.Sprintf(\"localhost:%d\", port))\n\n\t\t\t\t<-ctx.Done()\n\t\t\t\treturn nil\n\t\t\t},\n\t\t},\n\t)\n\n\tif err := rootCmd.Execute(); err != nil {\n\t\tfmt.Println(\"Failed to run:\", err)\n\t\tos.Exit(1)\n\t}\n}\n\nfunc runAndInit(ctx context.Context, dstDir string) error {\n\tif empty, err := files.IsEmpty(dstDir); err != nil {\n\t\treturn fmt.Errorf(\"failed to check if destination directory is empty: %w\", err)\n\t} else if !empty {\n\t\treturn fmt.Errorf(\"destination directory %s is not empty\", dstDir)\n\t}\n\tif err := os.MkdirAll(dstDir, 0o755); err != nil {\n\t\treturn fmt.Errorf(\"failed to create destination directory: %w\", err)\n\t}\n\n\ttmpDir, err := os.MkdirTemp(\"\", \"test-registry-\")\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer os.RemoveAll(tmpDir)\n\n\treg, hostPort, err := startRegistryTC(ctx, tmpDir)\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer func() {\n\t\tif err := reg.Terminate(ctx); err != nil {\n\t\t\tfmt.Println(\"Failed to terminate registry:\", err)\n\t\t}\n\t}()\n\n\taddr := fmt.Sprintf(\"localhost:%d\", hostPort)\n\tfor _, src := range images {\n\t\tdestRepo := fmt.Sprintf(\"%s/%s\", addr, strings.Split(src.Image, \":\")[0]) // e.g. localhost:5000/alpine\n\t\ttagPart := strings.Split(src.Image, \":\")[1]\n\n\t\tif src.SinglePlatform != \"\" {\n\t\t\tosArch := strings.SplitN(src.SinglePlatform, \"/\", 2)\n\t\t\tplat := v1.Platform{OS: osArch[0], Architecture: osArch[1]}\n\n\t\t\t// Pull source image for specified platform\n\t\t\tsrcRef, err := name.ParseReference(src.Image)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"parse reference: %w\", err)\n\t\t\t}\n\t\t\tsrcImg, err := remote.Image(srcRef, remote.WithPlatform(plat), remote.WithContext(ctx))\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\n\t\t\t// Push with desired tag\n\t\t\tdestRef, err := name.ParseReference(fmt.Sprintf(\"%s:%s\", destRepo, tagPart), name.Insecure)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"parse reference: %w\", err)\n\t\t\t}\n\t\t\tif err := remote.Write(destRef, srcImg,\n\t\t\t\tremote.WithContext(ctx), remote.WithAuth(authn.Anonymous)); err != nil {\n\t\t\t\treturn fmt.Errorf(\"write %s: %w\", destRef, err)\n\t\t\t}\n\t\t\tfmt.Printf(\"✅ pushed single-platform image %s\\n\", destRef.Name())\n\t\t\tcontinue\n\t\t}\n\n\t\tidx := mutate.IndexMediaType(empty.Index, types.OCIImageIndex) // start empty\n\n\t\tfor _, platStr := range src.Platforms {\n\t\t\tosArch := strings.SplitN(platStr, \"/\", 2)\n\t\t\tplat := v1.Platform{OS: osArch[0], Architecture: osArch[1]}\n\n\t\t\t// 1. pull source manifest for this platform\n\t\t\tsrcRef, err := name.ParseReference(src.Image)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"parse reference: %w\", err)\n\t\t\t}\n\t\t\tsrcImg, err := remote.Image(srcRef, remote.WithPlatform(plat), remote.WithContext(ctx))\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\n\t\t\t// 2. push it *by digest* into the new registry\n\t\t\tdigest, _ := srcImg.Digest()\n\t\t\tdestDigestRef, err := name.ParseReference(fmt.Sprintf(\"%s@%s\", destRepo, digest.String()), name.Insecure)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"parse reference: %w\", err)\n\t\t\t}\n\t\t\tif err := remote.Write(destDigestRef, srcImg,\n\t\t\t\tremote.WithContext(ctx), remote.WithAuth(authn.Anonymous)); err != nil {\n\t\t\t\treturn fmt.Errorf(\"write %s: %w\", destDigestRef, err)\n\t\t\t}\n\n\t\t\t// 3. add it to the (soon‑to‑be) index\n\t\t\tidx = mutate.AppendManifests(idx,\n\t\t\t\tmutate.IndexAddendum{Add: srcImg, Descriptor: v1.Descriptor{Platform: &plat}})\n\n\t\t\tfmt.Printf(\"✅ pushed %s for %s/%s\\n\", destDigestRef.Name(), plat.OS, plat.Architecture)\n\t\t}\n\n\t\t// 4. push the assembled index and tag it\n\t\tindexTag, err := name.ParseReference(fmt.Sprintf(\"%s:%s\", destRepo, tagPart), name.Insecure)\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"parse reference: %w\", err)\n\t\t}\n\t\tif err := remote.WriteIndex(indexTag, idx,\n\t\t\tremote.WithContext(ctx), remote.WithAuth(authn.Anonymous)); err != nil {\n\t\t\treturn fmt.Errorf(\"write index %s: %w\", indexTag, err)\n\t\t}\n\t\tfmt.Printf(\"🏷️  tagged multi-arch index %s\\n\", indexTag.Name())\n\t}\n\n\tfmt.Println(\"Copying registry data to\", dstDir)\n\tif err := os.CopyFS(dstDir, os.DirFS(tmpDir)); err != nil {\n\t\treturn fmt.Errorf(\"failed to copy registry data: %w\", err)\n\t}\n\n\tif err := catalog(ctx, addr); err != nil {\n\t\treturn fmt.Errorf(\"catalog tree: %w\", err)\n\t}\n\n\treturn nil\n}\n\nfunc runAndCatalog(ctx context.Context, dir string) error {\n\tdir, err := filepath.Abs(dir)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"failed to get absolute path: %w\", err)\n\t}\n\n\treg, _, err := startRegistryTC(ctx, dir)\n\tif err != nil {\n\t\treturn err\n\t}\n\tdefer func() {\n\t\tif err := reg.Terminate(ctx); err != nil {\n\t\t\tfmt.Println(\"Failed to terminate registry:\", err)\n\t\t}\n\t}()\n\n\tif err := catalog(ctx, reg.RegistryName); err != nil {\n\t\treturn fmt.Errorf(\"catalog: %w\", err)\n\t}\n\n\treturn nil\n}\n\nfunc catalog(ctx context.Context, addr string) error {\n\topts := []remote.Option{\n\t\tremote.WithContext(ctx),\n\t\tremote.WithAuth(authn.Anonymous), // local registry\n\t}\n\n\treg, err := name.NewRegistry(addr, name.Insecure)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"new registry: %w\", err)\n\t}\n\n\t// first, list all repositories\n\trepos, err := remote.Catalog(ctx, reg, opts...)\n\tif err != nil {\n\t\treturn err\n\t}\n\n\tfor _, repoName := range repos {\n\t\trepo := reg.Repo(repoName)\n\n\t\t// second, list all tags\n\t\ttagNames, err := remote.List(repo, opts...)\n\t\tif err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tfor _, tagName := range tagNames {\n\t\t\t// third, get the manifest\n\t\t\tref, err := name.ParseReference(fmt.Sprintf(\"%s/%s:%s\", addr, repoName, tagName))\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"parse reference: %w\", err)\n\t\t\t}\n\t\t\tdesc, err := remote.Get(ref, opts...)\n\t\t\tif err != nil {\n\t\t\t\treturn err\n\t\t\t}\n\n\t\t\trepoTag := fmt.Sprintf(\"%s:%s\", ref.Context().RepositoryStr(), ref.Identifier())\n\n\t\t\tswitch mt := desc.MediaType; mt {\n\t\t\tcase types.OCIImageIndex, types.DockerManifestList:\n\n\t\t\t\tfmt.Printf(\"%s %s\\n  index -> %s\\n\", repoTag, mt, desc.Digest)\n\n\t\t\t\tidx, _ := desc.ImageIndex()\n\t\t\t\tim, _ := idx.IndexManifest()\n\t\t\t\tfor _, m := range im.Manifests {\n\t\t\t\t\tfmt.Printf(\"  %s -> %s\\n\",\n\t\t\t\t\t\tm.Platform.String(),\n\t\t\t\t\t\tm.Digest,\n\t\t\t\t\t)\n\t\t\t\t}\n\n\t\t\tdefault: // single‑platform image\n\t\t\t\tfmt.Printf(\"%s %s\\n  single platform image -> %s\\n\", repoTag, mt, desc.Digest)\n\t\t\t}\n\t\t}\n\t}\n\treturn nil\n\n}\n\nfunc startRegistryTC(ctx context.Context, dir string) (*registry.RegistryContainer, int, error) {\n\tdir, err := filepath.Abs(dir)\n\tif err != nil {\n\t\treturn nil, 0, fmt.Errorf(\"failed to get absolute path: %w\", err)\n\t}\n\n\treg, err := registry.Run(ctx,\n\t\t\"registry:3\",\n\t\ttestcontainers.WithHostConfigModifier(func(hostConfig *container.HostConfig) {\n\t\t\thostConfig.Mounts = []mount.Mount{\n\t\t\t\t{\n\t\t\t\t\tType:   \"bind\",\n\t\t\t\t\tSource: dir,\n\t\t\t\t\tTarget: \"/var/lib/registry\",\n\t\t\t\t},\n\t\t\t}\n\t\t}),\n\t\ttestcontainers.WithWaitStrategy(\n\t\t\twait.ForHTTP(\"/v2/\").WithPort(\"5000/tcp\").\n\t\t\t\tWithStartupTimeout(10*time.Second),\n\t\t),\n\t)\n\tif err != nil {\n\t\treturn nil, 0, fmt.Errorf(\"start registry: %w\", err)\n\t}\n\n\tport, err := reg.MappedPort(ctx, \"5000/tcp\")\n\tif err != nil {\n\t\tif err := reg.Terminate(ctx); err != nil {\n\t\t\tfmt.Println(\"Failed to terminate registry:\", err)\n\t\t}\n\t\treturn nil, 0, fmt.Errorf(\"mapped port: %w\", err)\n\t}\n\treturn reg, port.Int(), nil\n}\n"
  },
  {
    "path": "tools/weights-gen/README.md",
    "content": "# weights-gen\n\nA tool for generating random weight files and optionally a `weights.lock` file for testing.\n\n## Installation\n\n```bash\ngo install github.com/replicate/cog/tools/weights-gen@latest\n```\n\n## Usage\n\n```bash\n# If installed via go install\nweights-gen [flags]\n\n# Or run directly from the repository\ngo run ./tools/weights-gen [flags]\n```\n\n## Flags\n\n| Flag | Short | Default | Description |\n|------|-------|---------|-------------|\n| `--count` | `-n` | `3` | Number of random weight files to generate |\n| `--min-size` | | `25mb` | Minimum file size (e.g., `12mb`, `25MB`, `1gb`) |\n| `--max-size` | | `50mb` | Maximum file size (e.g., `50mb`, `100MB`, `1gb`) |\n| `--output-dir` | | temp dir | Directory to write generated weight files |\n| `--output` | `-o` | `weights.lock` | Output path for weights.lock file |\n| `--dest-prefix` | | `/cache/` | Prefix for destination paths in lock file |\n| `--no-lock` | | `false` | Skip generating the weights.lock file |\n\n## Examples\n\n```bash\n# Generate 3 random files (25-50MB each) with a weights.lock file\ngo run ./tools/weights-gen\n\n# Generate 5 files between 12-50MB\ngo run ./tools/weights-gen --count 5 --min-size 12mb --max-size 50mb\n\n# Generate files to a specific output directory\ngo run ./tools/weights-gen --output-dir ./my-weights/\n\n# Generate only weight files without a lock file\ngo run ./tools/weights-gen --output-dir ./my-weights/ --no-lock\n\n# Generate files with custom destination prefix\ngo run ./tools/weights-gen --output-dir ./my-weights/ --dest-prefix /models/\n```\n\n## Output\n\nThe tool generates:\n- Random binary weight files named `weights-001.bin`, `weights-002.bin`, etc.\n- A `weights.lock` file (unless `--no-lock` is specified) containing metadata about each file including SHA256 digests for both original and gzip-compressed content.\n\nThe path to the generated files is always printed to stdout.\n\n## How the lock file works\n\nThe `weights.lock` file contains a `dest` field for each weight file. By default, `dest` paths use the `/cache/` prefix, which is the standard location for weights in Cog containers.\n\nUse `--dest-prefix` to override this behavior if you need different paths in the lock file (e.g., `/models/` or local paths for testing).\n"
  },
  {
    "path": "tools/weights-gen/main.go",
    "content": "// tools/weights-gen/main.go\npackage main\n\nimport (\n\t\"crypto/sha256\"\n\t\"encoding/hex\"\n\t\"fmt\"\n\t\"math/rand\"\n\t\"os\"\n\t\"path/filepath\"\n\t\"strconv\"\n\t\"strings\"\n\t\"time\"\n\n\t\"github.com/spf13/cobra\"\n\n\t\"github.com/replicate/cog/pkg/model\"\n)\n\nfunc main() {\n\tvar (\n\t\tdestPrefix string\n\t\toutputPath string\n\t\toutputDir  string\n\t\tcount      int\n\t\tminSize    string\n\t\tmaxSize    string\n\t\tnoLock     bool\n\t)\n\n\tcmd := &cobra.Command{\n\t\tUse:   \"weights-gen\",\n\t\tShort: \"Generate random weight files and optionally a weights.lock file\",\n\t\tLong: `This tool generates random weight files and optionally a weights.lock file for testing.\n\nIt creates random binary files of configurable size and computes their digests,\nsimulating what a future \"cog weights\" command would do with real weight files.\n\nBy default, both weight files and a weights.lock file are generated. Use --no-lock\nto generate only the weight files without the lock file.\n\nThe lock file's dest paths default to /cache/ for container paths.\nUse --dest-prefix to override this.\n\nExamples:\n  # Generate 3 random files (25-50MB each) with defaults (includes weights.lock)\n  weights-gen\n\n  # Generate 5 files between 12-50MB\n  weights-gen --count 5 --min-size 12mb --max-size 50mb\n\n  # Generate files to a specific output directory\n  weights-gen --output-dir ./my-weights/\n\n  # Generate only weight files without a lock file\n  weights-gen --output-dir ./my-weights/ --no-lock\n\n  # Generate files with custom destination prefix\n  weights-gen --output-dir ./my-weights/ --dest-prefix /models/`,\n\t\tRunE: func(cmd *cobra.Command, args []string) error {\n\t\t\tminBytes, err := parseSize(minSize)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"invalid --min-size: %w\", err)\n\t\t\t}\n\t\t\tmaxBytes, err := parseSize(maxSize)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"invalid --max-size: %w\", err)\n\t\t\t}\n\t\t\tif minBytes > maxBytes {\n\t\t\t\treturn fmt.Errorf(\"--min-size (%s) cannot be greater than --max-size (%s)\", minSize, maxSize)\n\t\t\t}\n\t\t\tif count < 1 {\n\t\t\t\treturn fmt.Errorf(\"--count must be at least 1\")\n\t\t\t}\n\n\t\t\treturn generateWeights(outputDir, destPrefix, outputPath, count, minBytes, maxBytes, !noLock)\n\t\t},\n\t}\n\n\tcmd.Flags().StringVar(&destPrefix, \"dest-prefix\", \"/cache/\", \"Prefix for destination paths in lock file (default: /cache/)\")\n\tcmd.Flags().StringVarP(&outputPath, \"output\", \"o\", \"weights.lock\", \"Output path for weights.lock file\")\n\tcmd.Flags().StringVar(&outputDir, \"output-dir\", \"\", \"Directory to write generated weight files (default: temp dir)\")\n\tcmd.Flags().IntVarP(&count, \"count\", \"n\", 3, \"Number of random weight files to generate\")\n\tcmd.Flags().StringVar(&minSize, \"min-size\", \"25mb\", \"Minimum file size (e.g., 12mb, 25MB, 1gb)\")\n\tcmd.Flags().StringVar(&maxSize, \"max-size\", \"50mb\", \"Maximum file size (e.g., 50mb, 100MB, 1gb)\")\n\tcmd.Flags().BoolVar(&noLock, \"no-lock\", false, \"Skip generating the weights.lock file\")\n\n\tif err := cmd.Execute(); err != nil {\n\t\tos.Exit(1)\n\t}\n}\n\n// parseSize parses a size string like \"25mb\", \"50MB\", \"1gb\" into bytes.\nfunc parseSize(s string) (int64, error) {\n\ts = strings.TrimSpace(strings.ToLower(s))\n\tif s == \"\" {\n\t\treturn 0, fmt.Errorf(\"empty size string\")\n\t}\n\n\tvar multiplier int64 = 1\n\tvar numStr string\n\n\tswitch {\n\tcase strings.HasSuffix(s, \"gb\"):\n\t\tmultiplier = 1024 * 1024 * 1024\n\t\tnumStr = strings.TrimSuffix(s, \"gb\")\n\tcase strings.HasSuffix(s, \"mb\"):\n\t\tmultiplier = 1024 * 1024\n\t\tnumStr = strings.TrimSuffix(s, \"mb\")\n\tcase strings.HasSuffix(s, \"kb\"):\n\t\tmultiplier = 1024\n\t\tnumStr = strings.TrimSuffix(s, \"kb\")\n\tcase strings.HasSuffix(s, \"b\"):\n\t\tnumStr = strings.TrimSuffix(s, \"b\")\n\tdefault:\n\t\t// Assume bytes if no suffix\n\t\tnumStr = s\n\t}\n\n\tnum, err := strconv.ParseFloat(strings.TrimSpace(numStr), 64)\n\tif err != nil {\n\t\treturn 0, fmt.Errorf(\"invalid number: %s\", numStr)\n\t}\n\tif num < 0 {\n\t\treturn 0, fmt.Errorf(\"size cannot be negative\")\n\t}\n\n\treturn int64(num * float64(multiplier)), nil\n}\n\nfunc generateWeights(outputDir, destPrefix, outputPath string, count int, minSize, maxSize int64, generateLock bool) error {\n\t// Determine where to write files\n\tvar filesDir string\n\n\tif outputDir != \"\" {\n\t\t// User specified an output directory\n\t\tif err := os.MkdirAll(outputDir, 0o755); err != nil {\n\t\t\treturn fmt.Errorf(\"create output directory: %w\", err)\n\t\t}\n\t\tfilesDir = outputDir\n\t} else {\n\t\t// Use a temp directory (not cleaned up so user can access the files)\n\t\ttmpDir, err := os.MkdirTemp(\"\", \"weights-gen-\")\n\t\tif err != nil {\n\t\t\treturn fmt.Errorf(\"create temp directory: %w\", err)\n\t\t}\n\t\tfilesDir = tmpDir\n\t}\n\n\t// Seed random number generator\n\t// Using math/rand is fine for test data generation - we don't need crypto randomness\n\trng := rand.New(rand.NewSource(time.Now().UnixNano())) //nolint:gosec\n\n\t// Generate random files\n\tfmt.Printf(\"Generating %d random weight files (%s - %s each)...\\n\",\n\t\tcount, formatSize(minSize), formatSize(maxSize))\n\n\tvar files []model.WeightFile\n\tfor i := 1; i <= count; i++ {\n\t\t// Random size between min and max\n\t\tvar size int64\n\t\tif minSize == maxSize {\n\t\t\tsize = minSize\n\t\t} else {\n\t\t\tsize = minSize + rng.Int63n(maxSize-minSize+1)\n\t\t}\n\n\t\tfilename := fmt.Sprintf(\"weights-%03d.bin\", i)\n\t\tfilePath := filepath.Join(filesDir, filename)\n\n\t\tfmt.Printf(\"  Creating %s (%s)...\\n\", filename, formatSize(size))\n\n\t\tif err := generateRandomFile(filePath, size, rng); err != nil {\n\t\t\treturn fmt.Errorf(\"generate %s: %w\", filename, err)\n\t\t}\n\n\t\tif generateLock {\n\t\t\twf, err := processFile(filePath, filesDir, destPrefix)\n\t\t\tif err != nil {\n\t\t\t\treturn fmt.Errorf(\"process %s: %w\", filename, err)\n\t\t\t}\n\n\t\t\tfiles = append(files, *wf)\n\t\t\tfmt.Printf(\"  Processed: %s -> %s\\n\", wf.Name, wf.Dest)\n\t\t} else {\n\t\t\tfmt.Printf(\"  Created: %s\\n\", filename)\n\t\t}\n\t}\n\n\tif generateLock {\n\t\tlock := &model.WeightsLock{\n\t\t\tVersion: \"1\",\n\t\t\tCreated: time.Now().UTC(),\n\t\t\tFiles:   files,\n\t\t}\n\n\t\tif err := lock.Save(outputPath); err != nil {\n\t\t\treturn err\n\t\t}\n\n\t\tfmt.Printf(\"\\nGenerated %s with %d files\\n\", outputPath, len(files))\n\t} else {\n\t\tfmt.Printf(\"\\nGenerated %d weight files (no lock file)\\n\", count)\n\t}\n\n\tfmt.Printf(\"Weight files written to: %s\\n\", filesDir)\n\treturn nil\n}\n\n// generateRandomFile creates a file filled with random data of the specified size.\nfunc generateRandomFile(path string, size int64, rng *rand.Rand) error {\n\tf, err := os.Create(path)\n\tif err != nil {\n\t\treturn fmt.Errorf(\"create file: %w\", err)\n\t}\n\tdefer f.Close()\n\n\t// Write in chunks to avoid allocating huge buffers\n\tconst chunkSize = 1024 * 1024 // 1MB chunks\n\tchunk := make([]byte, chunkSize)\n\tremaining := size\n\n\tfor remaining > 0 {\n\t\ttoWrite := min(remaining, chunkSize)\n\n\t\t// Fill chunk with random data\n\t\t_, _ = rng.Read(chunk[:toWrite])\n\n\t\tif _, err := f.Write(chunk[:toWrite]); err != nil {\n\t\t\treturn fmt.Errorf(\"write: %w\", err)\n\t\t}\n\t\tremaining -= toWrite\n\t}\n\n\treturn nil\n}\n\n// formatSize formats bytes into a human-readable string.\nfunc formatSize(bytes int64) string {\n\tconst (\n\t\tkb = 1024\n\t\tmb = kb * 1024\n\t\tgb = mb * 1024\n\t)\n\n\tswitch {\n\tcase bytes >= gb:\n\t\treturn fmt.Sprintf(\"%.1fGB\", float64(bytes)/float64(gb))\n\tcase bytes >= mb:\n\t\treturn fmt.Sprintf(\"%.1fMB\", float64(bytes)/float64(mb))\n\tcase bytes >= kb:\n\t\treturn fmt.Sprintf(\"%.1fKB\", float64(bytes)/float64(kb))\n\tdefault:\n\t\treturn fmt.Sprintf(\"%dB\", bytes)\n\t}\n}\n\nfunc processFile(path, baseDir, destPrefix string) (*model.WeightFile, error) {\n\t// Read file\n\tdata, err := os.ReadFile(path)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"read file: %w\", err)\n\t}\n\n\t// Compute digest\n\thash := sha256.Sum256(data)\n\tdigest := \"sha256:\" + hex.EncodeToString(hash[:])\n\n\t// Compute relative path for dest\n\trelPath, err := filepath.Rel(baseDir, path)\n\tif err != nil {\n\t\treturn nil, fmt.Errorf(\"rel path: %w\", err)\n\t}\n\tdest := filepath.Join(destPrefix, relPath)\n\t// Normalize to forward slashes for container paths\n\tdest = strings.ReplaceAll(dest, \"\\\\\", \"/\")\n\n\t// Generate a simple identifier from the filename (without extension)\n\tbaseName := filepath.Base(path)\n\tname := baseName[:len(baseName)-len(filepath.Ext(baseName))]\n\n\tsize := int64(len(data))\n\n\treturn &model.WeightFile{\n\t\tName:             name,\n\t\tDest:             dest,\n\t\tDigest:           digest,\n\t\tDigestOriginal:   digest,\n\t\tSize:             size,\n\t\tSizeUncompressed: size,\n\t\tMediaType:        model.MediaTypeWeightLayer,\n\t}, nil\n}\n"
  }
]