coglet support

.github/workflows/ci.yaml

@@ -26,6 +26,20 @@ jobs:
           fetch-depth: 0 # needed for setuptools_scm version determination
       - uses: hynek/build-and-inspect-python-package@v2
         id: baipp
+      # Build coglet wheel with embedded Go binary
+      - uses: actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+      - uses: astral-sh/setup-uv@v7
+      - name: Build coglet-server binary
+        run: script/build-coglet-server
+      - name: Build coglet wheel
+        run: uv build --wheel --out-dir=dist ./coglet
+      - name: Upload coglet wheel
+        uses: actions/upload-artifact@v4
+        with:
+          name: CogletWheel
+          path: dist/coglet-*.whl
     outputs:
       python-versions: ${{ steps.baipp.outputs.supported_python_classifiers_json_array }}
 
@@ -34,13 +48,11 @@ jobs:
     needs: build-python
     runs-on: ubuntu-latest-8-cores
     steps:
-      - name: Download pre-built packages
+      - name: Download built wheels
         uses: actions/download-artifact@v6
         with:
-          name: Packages
           path: dist
-      - name: Set COG_WHEEL
-        run: echo COG_WHEEL=$(ls dist/*.whl) >>"$GITHUB_ENV"
+          merge-multiple: true
       - name: Extract source distribution
         run: tar xf dist/*.tar.gz --strip-components=1
       - uses: actions/setup-go@v5
@@ -81,13 +93,11 @@ jobs:
       - uses: actions/checkout@v6
         with:
           fetch-depth: 0 # needed for goreleaser version determination
-      - name: Download pre-built packages
+      - name: Download built wheels
         uses: actions/download-artifact@v6
         with:
-          name: Packages
           path: dist
-      - name: Set COG_WHEEL
-        run: echo COG_WHEEL=$(ls dist/*.whl) >>"$GITHUB_ENV"
+          merge-multiple: true
       - uses: actions/setup-go@v5
         with:
           go-version-file: go.mod
@@ -106,13 +116,11 @@ jobs:
         pydantic: ["1", "2"]
         python-version: ${{ fromJson(needs.build-python.outputs.python-versions) }}
     steps:
-      - name: Download pre-built packages
+      - name: Download built wheels
         uses: actions/download-artifact@v6
         with:
-          name: Packages
           path: dist
-      - name: Set COG_WHEEL
-        run: echo COG_WHEEL=$(ls dist/*.whl) >>"$GITHUB_ENV"
+          merge-multiple: true
       - name: Extract source distribution
         run: tar xf dist/*.tar.gz --strip-components=1
       - uses: actions/setup-python@v6
@@ -127,7 +135,7 @@ jobs:
       - name: Remove src to ensure tests run against wheel
         run: rm -rf python/cog
       - name: Test
-        run: python -Im tox run --installpkg "$COG_WHEEL" -e ${{ env.TOX_PYTHON }}-pydantic${{ matrix.pydantic }}-tests
+        run: python -Im tox run --installpkg dist/cog-*.whl -e ${{ env.TOX_PYTHON }}-pydantic${{ matrix.pydantic }}-tests
 
   # ==========================================================================
   # Coglet tests (Go server + Python SDK for the coglet runtime)
@@ -184,10 +192,14 @@ jobs:
 
   # cannot run this on mac due to licensing issues: https://github.com/actions/virtual-environments/issues/2150
   test-integration-matrix:
-    name: "Test integration Matrix"
+    name: "Test integration (${{ matrix.runtime }})"
     needs: build-python
     runs-on: ubuntu-latest-16-cores
-    timeout-minutes: 25
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        runtime: [cog, coglet, coglet-alpha]
     steps:
       - uses: actions/checkout@v6
         with:
@@ -198,13 +210,11 @@ jobs:
           registry: index.docker.io
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
-      - name: Download pre-built packages
+      - name: Download built wheels
         uses: actions/download-artifact@v6
         with:
-          name: Packages
           path: dist
-      - name: Set COG_WHEEL
-        run: echo COG_WHEEL=$(ls dist/*.whl) >>"$GITHUB_ENV"
+          merge-multiple: true
       - uses: actions/setup-go@v5
         with:
           go-version-file: go.mod
@@ -215,7 +225,24 @@ jobs:
       - name: Prepare tox
         run: uv pip install --system tox tox-uv
       - name: Test
-        run: make test-integration
+        env:
+          COG_WHEEL: ${{ matrix.runtime }}
+        # coglet/coglet-alpha tests may fail due to error message differences - don't block PR
+        continue-on-error: ${{ matrix.runtime != 'cog' }}
+        run: |
+          # Build cog binary
+          make cog
+          # Run full test suite with reruns for flaky tests (matches main branch)
+          # - cog: full suite
+          # - coglet-alpha: stop after 5 failures (passes currently)
+          # - coglet: exit on first failure (needs work, just want to see what fails)
+          if [ "$COG_WHEEL" = "cog" ]; then
+            PATH="$(pwd):$PATH" uv run tox -e integration -- -n auto -vv --reruns 3
+          elif [ "$COG_WHEEL" = "coglet" ]; then
+            PATH="$(pwd):$PATH" uv run tox -e integration -- -n auto -vv -x
+          else
+            PATH="$(pwd):$PATH" uv run tox -e integration -- -n auto -vv --reruns 3 --maxfail=5
+          fi
 
   release:
     name: "Release"
@@ -231,13 +258,11 @@ jobs:
       - uses: actions/checkout@v6
         with:
           fetch-depth: 0 # needed for goreleaser version determination
-      - name: Download pre-built packages
+      - name: Download built wheels
         uses: actions/download-artifact@v6
         with:
-          name: Packages
           path: dist
-      - name: Set COG_WHEEL
-        run: echo COG_WHEEL=$(ls dist/*.whl) >>"$GITHUB_ENV"
+          merge-multiple: true
       - uses: actions/setup-go@v5
         with:
           go-version-file: go.mod

.gitignore

@@ -8,7 +8,7 @@ __pycache__
 build
 dist
 *.egg-info
-pkg/dockerfile/embed/*
+pkg/wheels/*.whl
 # Used by a vim plugin (projectionist)
 .projections.json
 .tox/
@@ -24,6 +24,9 @@ flag_file
 bin/*
 .beads/
 
-# Auto-generated version files from setuptools-scm
+# Auto-:d version files from setuptools-scm
 python/cog/_version.py
 coglet/python/coglet/_version.py
+
+# Built coglet-server binaries
+coglet/python/cog/bin/

.goreleaser.yaml

@@ -2,7 +2,7 @@ version: 2
 before:
   hooks:
     - go mod tidy
-    - make wheel
+    - go generate ./pkg/wheels/...
 dist: ./dist/go
 env:
   - CGO_ENABLED=0

CONTRIBUTING.md

@@ -146,6 +146,28 @@ As much as possible, this is attempting to follow the [Standard Go Project Layou
 - `pkg/util/` - Various packages that aren't part of Cog. They could reasonably be separate re-usable projects.
 - `python/` - The Cog Python library.
 - `test-integration/` - High-level integration tests for Cog.
+- `tools/compatgen/` - Tool for generating CUDA/PyTorch/TensorFlow compatibility matrices.
+
+## Updating compatibility matrices
+
+The 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.
+
+To regenerate the compatibility matrices, run:
+
+```sh
+# Regenerate all matrices
+script/generate-compat
+
+# Or regenerate specific matrices
+script/generate-compat cuda
+script/generate-compat torch
+script/generate-compat tensorflow
+```
+
+The generated files are:
+- `pkg/config/cuda_base_images.json` - Available NVIDIA CUDA base images
+- `pkg/config/torch_compatibility_matrix.json` - PyTorch/CUDA/Python compatibility
+- `pkg/config/tf_compatibility_matrix.json` - TensorFlow/CUDA/Python compatibility
 
 ## Concepts
 

Makefile

@@ -15,10 +15,7 @@ GOLINT := $(GO) tool golangci-lint
 UV ?= uv
 TOX := $(UV) run tox
 
-COG_GO_SOURCE := $(shell find cmd pkg -type f)
-COG_PYTHON_SOURCE := $(shell find python/cog -type f -name '*.py')
-COGLET_GO_SOURCE := $(shell find coglet -type f -name '*.go')
-COGLET_PYTHON_SOURCE := $(shell find coglet/python -type f -name '*.py')
+COG_GO_SOURCE := $(shell find cmd pkg -type f -name '*.go')
 
 COG_BINARIES := cog base-image
 
@@ -27,30 +24,7 @@ default: all
 .PHONY: all
 all: cog
 
-.PHONY: wheel
-wheel: pkg/dockerfile/embed/.wheel
-
-ifdef COG_WHEEL
-pkg/dockerfile/embed/.wheel: $(COG_WHEEL)
-	@mkdir -p pkg/dockerfile/embed
-	@rm -f pkg/dockerfile/embed/*.whl # there can only be one embedded wheel
-	@echo "Using prebuilt COG_WHEEL $<"
-	cp $< pkg/dockerfile/embed/
-	@touch $@
-else
-pkg/dockerfile/embed/.wheel: $(COG_PYTHON_SOURCE)
-	@mkdir -p pkg/dockerfile/embed
-	@rm -f pkg/dockerfile/embed/*.whl # there can only be one embedded wheel
-	$(UV) build --wheel --out-dir=pkg/dockerfile/embed .
-	@touch $@
-
-define COG_WHEEL
-    $(shell find pkg/dockerfile/embed -type f -name '*.whl')
-endef
-
-endif
-
-$(COG_BINARIES): $(COG_GO_SOURCE) pkg/dockerfile/embed/.wheel
+$(COG_BINARIES): $(COG_GO_SOURCE) generate
 	@echo Building $@
 	@if git name-rev --name-only --tags HEAD | grep -qFx undefined; then \
 		GOFLAGS=-buildvcs=false $(GORELEASER) build --clean --snapshot --single-target --id $@ --output $@; \
@@ -63,13 +37,17 @@ install: $(COG_BINARIES)
 	$(INSTALL) -d $(DESTDIR)$(BINDIR)
 	$(INSTALL) $< $(DESTDIR)$(BINDIR)/$<
 
+.PHONY: wheel
+wheel:
+	script/build-wheels
+
 .PHONY: clean
 clean: clean-coglet
-	rm -rf .tox build dist pkg/dockerfile/embed
+	rm -rf .tox build dist pkg/wheels/*.whl
 	rm -f $(COG_BINARIES)
 
 .PHONY: test-go
-test-go: pkg/dockerfile/embed/.wheel
+test-go: generate
 	$(GO) tool gotestsum -- -short -timeout 1200s -parallel 5 $$(go list ./... | grep -v 'coglet/') $(ARGS)
 
 .PHONY: test-integration
@@ -78,8 +56,8 @@ test-integration: $(COG_BINARIES)
 	PATH="$(PWD):$(PATH)" $(TOX) -e integration
 
 .PHONY: test-python
-test-python: pkg/dockerfile/embed/.wheel
-	$(TOX) run --installpkg $(COG_WHEEL) -f tests
+test-python: generate
+	$(TOX) run --installpkg $$(ls dist/cog-*.whl) -f tests
 
 .PHONY: test
 test: test-go test-python test-integration
@@ -94,7 +72,7 @@ generate:
 	$(GO) generate ./...
 
 .PHONY: vet
-vet: pkg/dockerfile/embed/.wheel
+vet: generate
 	$(GO) vet ./...
 
 .PHONY: check-fmt
@@ -103,9 +81,9 @@ check-fmt:
 	@test -z $$($(GOIMPORTS) -l .)
 
 .PHONY: lint
-lint: pkg/dockerfile/embed/.wheel check-fmt vet
+lint: generate check-fmt vet
 	$(GOLINT) run ./...
-	$(TOX) run --installpkg $(COG_WHEEL) -e lint,typecheck-pydantic2
+	$(TOX) run --installpkg $$(ls dist/cog-*.whl) -e lint,typecheck-pydantic2
 
 .PHONY: run-docs-server
 run-docs-server:
@@ -122,30 +100,6 @@ gen-mocks:
 # Coglet targets
 # =============================================================================
 
-COGLET_BINARY_DIR := coglet/python/cog
-
-# Build coglet-server binary for a specific OS/ARCH
-# Usage: make coglet-server-binary GOOS=linux GOARCH=amd64
-.PHONY: coglet-server-binary
-coglet-server-binary: $(COGLET_GO_SOURCE)
-	CGO_ENABLED=0 $(GO) build -o $(COGLET_BINARY_DIR)/cog-$(GOOS)-$(GOARCH) ./coglet/cmd/coglet-server
-
-# Build all coglet-server binaries (for wheel embedding)
-.PHONY: coglet-server-binaries
-coglet-server-binaries: $(COGLET_GO_SOURCE)
-	@rm -f $(COGLET_BINARY_DIR)/cog-*
-	@for os in darwin linux; do \
-		for arch in amd64 arm64; do \
-			echo "Building coglet-server for $$os/$$arch"; \
-			CGO_ENABLED=0 GOOS=$$os GOARCH=$$arch $(GO) build -o $(COGLET_BINARY_DIR)/cog-$$os-$$arch ./coglet/cmd/coglet-server; \
-		done; \
-	done
-
-# Build coglet wheel (includes embedded Go binaries)
-.PHONY: coglet-wheel
-coglet-wheel: coglet-server-binaries
-	cd coglet && $(UV) build --wheel --out-dir=dist .
-
 # Run coglet Go tests
 .PHONY: test-coglet-go
 test-coglet-go:
@@ -163,5 +117,4 @@ test-coglet: test-coglet-go test-coglet-python
 # Clean coglet build artifacts
 .PHONY: clean-coglet
 clean-coglet:
-	rm -rf coglet/dist coglet/build
-	rm -f $(COGLET_BINARY_DIR)/cog-*
+	rm -rf coglet/dist coglet/build coglet/python/cog/bin

coglet/pyproject.toml

@@ -74,7 +74,7 @@ where = ['python']
 exclude = ['tests*']
 
 [tool.setuptools.package-data]
-'cog' = ['cog-*', '**/*.pyi']
+'cog' = ['bin/coglet-server-*', '**/*.pyi']
 'coglet' = ['*.pyi']
 
 [tool.setuptools_scm]

coglet/python/cog/command/go_cog.py

@@ -21,9 +21,9 @@ def run(subcmd: str, args: List[str]) -> None:
         print(f'Unsupported architecture: {goarch}')
         sys.exit(1)
 
-    # Binaries are bundled in python/cog
-    cmd = f'cog-{goos}-{goarch}'
-    exe = os.path.join(Path(__file__).parent.parent, cmd)
+    # Binaries are bundled in python/cog/bin
+    cmd = f'coglet-server-{goos}-{goarch}'
+    exe = os.path.join(Path(__file__).parent.parent, 'bin', cmd)
     args = [exe, subcmd] + args
     # Replicate Go logger logs to stdout in production mode
     # Use stderr instead to be consistent with legacy Cog

coglet/python/tests/schemas/cog.yaml

@@ -1,3 +1,3 @@
 build:
-  python_version: "3.8"
+  python_version: "3.9"
 predict: "predict.py:Predictor"

docs/llms.txt

@@ -2405,7 +2405,7 @@ build:
   python_version: "3.11.1"
 ```
 
-Cog supports all active branches of Python: 3.8, 3.9, 3.10, 3.11, 3.12, 3.13. If you don't define a version, Cog will use the latest version of Python 3.12 or a version of Python that is compatible with the versions of PyTorch or TensorFlow you specify.
+Cog supports all active branches of Python: 3.9, 3.10, 3.11, 3.12, 3.13. If you don't define a version, Cog will use the latest version of Python 3.12 or a version of Python that is compatible with the versions of PyTorch or TensorFlow you specify.
 
 Note 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.
 

docs/yaml.md

@@ -115,7 +115,7 @@ build:
   python_version: "3.11.1"
 ```
 
-Cog supports all active branches of Python: 3.8, 3.9, 3.10, 3.11, 3.12, 3.13. If you don't define a version, Cog will use the latest version of Python 3.12 or a version of Python that is compatible with the versions of PyTorch or TensorFlow you specify.
+Cog supports all active branches of Python: 3.9, 3.10, 3.11, 3.12, 3.13. If you don't define a version, Cog will use the latest version of Python 3.12 or a version of Python that is compatible with the versions of PyTorch or TensorFlow you specify.
 
 Note 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.