Initial working version using Pydantic for type annotations

[bfirsh]

The exploration in #343 was sufficiently explored and getting very messy, so I thought I would create a new pull request for the real implementation.

This is an implementation of #193, #205, and #259 using Pydantic for the type system, FastAPI for the model server, and OpenAPI for the schema.

Models now look a bit like this:

from cog import Predictor, Input, Path
import torch

class ColorizationPredictor(Predictor):
    def setup(self):
        self.model = torch.load("./weights.pth")

    def predict(self,
          image: Path = Input(title="Grayscale input image")
    ) -> Path:
        processed_input = preprocess(image)
        output = self.model(processed_input)
        return postprocess(output)

The HTTP model server is a work in progress and will be finalized as part of future work.

The Redis queue worker remains in place and the API is unchanged. As part of future work we will implement an AMQP API with a better API.

Commits are significant and this should be reviewed commit by commit. This is an intentionally large and long-running branch -- main is currently used as the documentation for Replicate.

Todo

• cog predict with new API
• Input ordering?
• Ensure HTTP server is single threaded.
• HTTP URLs for input and output
• Add choices option to Input to generate an Enum.
• cog init template
• Update getting started docs
• Update reference docs
• Rename Predictor to BasePredictor so we can do from cog import BasePredictor, Input, Path without colliding with a Predictor subclass.
• Better errors if you get output type wrong.
• Reliably set Cog version. In development, the version is set to 0.0.1 which is pretty useless. It should probably be hardcoded somewhere, and the release process needs updating. We need this to be able to switch between old and new Cog API in Replicate. We're just going to support version dev everywhere which means "latest".
• Update examples Update examples for new Cog types cog-examples#8
• Update the OpenAPI Outputs to be a consistent place. (Simple outputs currently end up in Request object schema)
• Test returning object as output
• Test arbitrary JSON objects in output
• Is go.sum up to date?

Future (Pydantic + OpenAPI)

• Tell users to upgrade if they use @cog.input() with new Cog.
• Document nouns to make sure we're all on the same page. "Input"/"Inputs"?
• Document how to migrate @input() to Pydantic type annotations.
• Go through and fix any TODOs
• Make input validation errors work for cog predict in a simple way until we lock down HTTP API
• Make input errors work for Redis API & Replicate
• Integrate with Replicate - https://github.com/replicate/replicate-web/pull/907
• Force users to pick an output type -- throw an error if it's not set. (Couldn't just be Any, but they need to be explicit!)
• Clean up temporary files -- both input and output
• Do we want to make cog predict with an existing image backwards compatible? ('cos there are old models on Replicate)
• What happens if you use pathlib.Path as an input type? Does it work? Should we throw an error? (This might happen if you import wrong package / collide imports)
• What happens if you return pathlib.Path from a predict function? Does this work? Should it throw an error?
• Should we have a fixed set of supported inputs/outputs somehow?
• Do we want to force users to put extra properties in additionalProperties? Should we even allow extra properties to cog.Input()?
• Merge examples PR Update examples for new Cog types cog-examples#8

Far-flung future

• Vendor Pydantic and FastAPI
• Add back support for PyTorch and Tensorflow tensors as output. This was removed because it was fiddly to get working with Pydantic. You now need to convert them to numpy arrays.
• Add back timing information to server.
• Better error if you forget an output type

  pydantic.error_wrappers.ValidationError: 1 validation error for Prediction
  output
    unexpected value; permitted: None (type=value_error.const; given=hello foo; permitted=(None,))
  

• Multiple file outputs for cog predict. It currently just outputs the plain JSON if there isn't a single file output.
• Logging, as described in comment in Design exploration for new Cog schema and API #343
• Review request/response objects and make sure we're happy with them.
  • This might be an opportunity to rename the statuses and ensure they're consistent ("success" vs "failed").
  • Do we want the input type to have a second input level so we can have top-level request options?
• Do we need separate URLs for input/output schema, or can they be reliably fetched from OpenAPI?
• Finalize & document HTTP API
• Do we want to version the prediction API, or let the client handle it? Design prediction API versioning #94
• Test arbitrary objects in input
• What if people run new models from a directory of code with old Cog? What happens if we do that? Do we care? Display warning or throw error if user tries to run remote model with outdated local cog #286
• What if people run new models from an image with old Cog? Maybe we throw a simple error from /predict?

Based on (& blocked by)

• Upgrade to redis-py 4 #383
• Run redis integration tests in Docker network #385

Refs

Closes #58
Closes #113
Closes #193
Closes #205
Closes #212
Closes #246
Closes #259
Closes #262
Closes #263
Closes #304
Closes #306
Closes #327
Closes #328

[zeke]

❤️

[zeke]

Why is this commented out?

[bfirsh]

Just tinkering with this. Perhaps it's more clear what Cog is for (packaging models into Docker containers) when there isn't the diversion into development environments.

[zeke]

If this argument is not passed, the input is required. If it is explicitly set to None, the input is optional.

👍🏼

[bfirsh]

There could be an argument for a more explicit, easier to understand design here. #114 Not sure what is best.

[bfirsh]

👍 To summarize an IRL conversion: we chose the extremely cool name "future" for the branch.

I have:

1. Cleaned up the commits a little (squashed some "fix tests" commits etc)
2. Rebased on Run redis integration tests in Docker network #385 and main
3. Updated go.mod
4. Based this PR on our new future branch.

We're ready to get this into our future branch, I think!

Next steps, once this is in:

• Create a new Mega Pull Request for the "future" branch. This branch will be an initial version of Cog that we can use to start testing it with Replicate internally. Use Pydantic for type annotations #407
• Migrate "Future (Pydantic + OpenAPI)" tasks to that pull request.
• Create issues for the "Far flung future" and add them to the "next" in the project to remind us to do them before we ship new cog.

[zeke]

The future looks bright. 😎

Ship it (to a feature branch)! 🚢

[bfirsh]

I have added a few tickets for the far flung future that can be seen in the mentions above. They're in the "next" status in our project to remind us to look at them before we ship new Cog.