Initial working version using Pydantic for type annotations

CONTRIBUTING.md

@@ -96,7 +96,11 @@ Common contribution types include: `doc`, `code`, `bug`, and `ideas`. See the fu
 
 You'll need to [install Go 1.16](https://golang.org/doc/install). If you're using a newer Mac with an M1 chip, be sure to download the `darwin-arm64` installer package. Alternatively you can run `brew install go` which will automatically detect and use the appropriate installer for your system architecture.
 
-Once you have Go installed, then run:
+Install the Python dependencies:
+
+    pip install -r requirements-dev.txt
+
+Once you have Go installed, run:
 
     make install
 
@@ -127,6 +131,39 @@ As much as possible, this is attempting to follow the [Standard Go Project Layou
 - `python/` - The Cog Python library.
 - `test-integration/` - High-level integration tests for Cog.
 
+## Runnings tests
+
+To run the entire test suite:
+
+```sh
+make test
+```
+
+To run just the Golang tests:
+
+```sh
+make test-go
+```
+
+To run just the Python tests:
+
+```sh
+make test-python
+```
+
+To stand up a server for one of the integration tests:
+
+```sh
+make install
+pip install -r requirements-dev.txt
+make test
+cd test-integration/test_integration/fixtures/file-project
+cog build
+docker run -p 5001:5000 --init --platform=linux/amd64 cog-file-project
+```
+
+Then visit [localhost:5001](http://localhost:5001) in your browser.
+
 ## Publishing a release
 
 This project has a [GitHub Actions workflow](https://github.com/replicate/cog/blob/39cfc5c44ab81832886c9139ee130296f1585b28/.github/workflows/ci.yaml#L107) that uses [goreleaser](https://goreleaser.com/quick-start/#quick-start) to facilitate the process of publishing new releases. The release process is triggered by manually creating and pushing a new git tag.

README.md

@@ -1,16 +1,14 @@
 # Cog: Containers for machine learning
 
-Use Docker for machine learning, without all the pain.
+Put your machine learning model in a standard, production-ready Docker container without having to know how Docker works.
 
-Cog gives you a consistent environment to run your model in – for developing on your laptop, training on GPU machines, and for other people working on the model. Then, when the model is trained and you want to share or deploy it, you can bake the model into a Docker image that serves a standard HTTP API.
+Cog does a few things for you:
 
-Cog does a few handy things beyond normal Docker:
+- **Automatic Docker image.** Define your environment with a simple configuration file, and Cog generates a `Dockerfile` with all the best practices.
+- **Standard, production-ready HTTP and AMQP interface.** Automatically generate APIs for integrating with production systems, battle hardened on Replicate.
+- **No more CUDA hell.** Cog knows which CUDA/cuDNN/PyTorch/Tensorflow/Python combos are compatible and will set it all up correctly for you.
 
-- **Automatic Docker image.** Define your environment with a simple configuration file, then Cog will generate Dockerfiles with best practices and do all the GPU configuration for you.
-- **Automatic HTTP service.** Cog will generate an HTTP service from the definition of your model, so you don't need to write a Flask server in the right way.
-- **No more CUDA hell.** Cog knows which CUDA/cuDNN/PyTorch/Tensorflow/Python combos are compatible and will pick the right versions for you.
-
-## Develop and train in a consistent environment
+## How it works
 
 Define the Docker environment your model runs in with `cog.yaml`:
 
@@ -23,46 +21,24 @@ build:
   python_version: "3.8"
   python_packages:
     - "torch==1.8.1"
+predict: "predict.py:Predictor"
 ```
 
-Now, you can run commands inside this environment:
-
-```
-$ cog run python train.py
-...
-```
-
-This will:
-
-- Generate a `Dockerfile` with best practices
-- Pick the right CUDA version
-- Build an image
-- Run `python train.py` in the image with the current directory mounted as a volume and GPUs hooked up correctly
-
-<!-- TODO: this doesn't work yet (needs ports etc)
-Or, spin up a Jupyter notebook:
-
-```
-$ cog run jupyter notebook
-```
--->
-
-## Put a trained model in a Docker image
-
-First, you define how predictions are run on your model:
+And define how predictions are run on your model with `predict.py`:
 
 ```python
-import cog
+from cog import BasePredictor, Input, Path
 import torch
 
-class ColorizationPredictor(cog.Predictor):
+class Predictor(BasePredictor):
     def setup(self):
         """Load the model into memory to make running multiple predictions efficient"""
         self.model = torch.load("./weights.pth")
-    
+
     # The arguments and types the model takes as input
-    @cog.input("input", type=cog.Path, help="Grayscale input image")
-    def predict(self, input):
+    def predict(self,
+          input: Path = Input(title="Grayscale input image")
+    ) -> Path:
         """Run a single prediction on the model"""
         processed_input = preprocess(input)
         output = self.model(processed_input)
@@ -87,10 +63,34 @@ $ cog build -t my-colorization-model
 
 $ docker run -d -p 5000:5000 --gpus all my-colorization-model
 
-$ curl http://localhost:5000/predict -X POST -F [email protected]
+$ curl http://localhost:5000/predictions -X POST \
+   --data '{"input": "https://.../input.jpg"}'
 ```
 
-That's it! Your model will now run forever in this reproducible Docker environment.
+<!-- In development, you can also run arbitrary commands inside the Docker environment:
+
+```
+$ cog run python train.py
+...
+``` -->
+
+<!-- TODO: this doesn't work yet (needs ports etc)
+Or, spin up a Jupyter notebook:
+
+```
+$ cog run jupyter notebook
+```
+-->
+
+## Deploying models to production
+
+Cog does a number of things out of the box to help you deploy models to production:
+
+- **Standard interface.** Put models inside Cog containers, and they'll run anywhere that runs Docker containers.
+- **HTTP prediction server, based on FastAPI.**
+- **Type checking, based on Pydantic.** Cog models define their input and output with JSON Schema, and the HTTP server is defined with OpenAPI.
+- **AMQP RPC interface.** Long-running deep learning models or batch processing is best architected with a queue. Cog models can do this out of the box.
+- **Read/write files from cloud storage.** Files can be read and written directly on Amazon S3 and Google Cloud Storage for efficiency.
 
 ## Why are we building this?
 
@@ -129,10 +129,10 @@ sudo chmod +x /usr/local/bin/cog
 - [Get started with your own model](docs/getting-started-own-model.md)
 - [Take a look at some examples of using Cog](https://github.com/replicate/cog-examples)
 - [`cog.yaml` reference](docs/yaml.md) to learn how to define your model's environment
-- [Prediction interface reference](docs/python.md) to learn how the `cog.Predictor` interface works
+- [Prediction interface reference](docs/python.md) to learn how the `Predictor` interface works
 
 ## Need help?
- 
+
 [Join us in #cog on Discord.](https://discord.gg/QmzJApGjyE)
 
 ## Contributors ✨
@@ -166,4 +166,4 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/d
 
 <!-- ALL-CONTRIBUTORS-LIST:END -->
 
-This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
+This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!

docs/getting-started-own-model.md

@@ -66,19 +66,18 @@ With `cog.yaml`, you can also install system packages and other things. [Take a
 The 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:
 
 ```python
-import cog
-from pathlib import Path
+from cog import BasePredictor, Path, Input
 import torch
 
-class Predictor(cog.Predictor):
+class Predictor(BasePredictor):
     def setup(self):
         """Load the model into memory to make running multiple predictions efficient"""
         self.net = torch.load("weights.pth")
 
-    # Define the input types for a prediction
-    @cog.input("input", type=Path, help="Image to enlarge")
-    @cog.input("scale", type=float, default=1.5, help="Factor to scale image by")
-    def predict(self, input, scale):
+    def predict(self,
+            image: Path = Input(description="Image to enlarge"),
+            scale: float = Input(description="Factor to scale image by", default=1.5)
+    ) -> Path:
         """Run a single prediction on the model"""
         # ... pre-processing ...
         output = self.net(input)
@@ -88,16 +87,26 @@ class Predictor(cog.Predictor):
 
 Edit 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.
 
-You also need to define the inputs to your model using the `@cog.input()` decorator, as demonstrated above. The first argument maps to the name of the argument in the `predict()` function, and it also takes these other arguments:
+You 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:
 
-- `type`: Either `str`, `int`, `float`, `bool`, or `Path` (be sure to add the import, as in the example above). `Path` is used for files. For more complex inputs, save it to a file and use `Path`.
-- `help`: A description of what to pass to this input for users of the model
+- `str`: a string
+- `int`: an integer
+- `float`: a floating point number
+- `bool`: a boolean
+- `cog.File`: a file-like object representing a file
+- `cog.Path`: a path to a file on disk
+
+You can provide more information about the input with the `Input()` function, as shown above. It takes these basic arguments:
+
+- `description`: A description of what to pass to this input for users of the model
 - `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.
-- `min`: A minimum value for `int` or `float` types.
-- `max`: A maximum value for `int` or `float` types.
-- `options`: A list of values to limit the input to. It can be used with `str`, `int`, and `float` inputs.
+- `gt`: For `int` or `float` types, the value should be greater than this number.
+- `ge`: For `int` or `float` types, the value should be greater than or equal to this number.
+- `lt`: For `int` or `float` types, the value should be less than this number.
+- `le`: For `int` or `float` types, the value should be less than or equal to this number.
+- `choices`: A list of possible values for this input.
 
-For more details about writing your model interface, [take a look at the prediction interface documentation](python.md).
+There are some more advanced options you can pass, too. For more details, [take a look at the prediction interface documentation](python.md).
 
 Next, add the line `predict: "predict.py:Predictor"` to your `cog.yaml`, so it looks something like this:
 
@@ -122,7 +131,7 @@ Written output to output.png
 To pass more inputs to the model, you can add more `-i` options:
 
 ```
-$ cog predict -i input=@input.jpg -i scale=2.0
+$ cog predict -i image=@image.jpg -i scale=2.0
 ```
 
 In this case it is just a number, not a file, so you don't need the `@` prefix.

docs/getting-started.md

@@ -59,32 +59,31 @@ First, run this to get some pre-trained model weights:
 Then, we need to write some code to describe how predictions are run on the model. Save this to `predict.py`:
 
 ```python
-import cog
-from pathlib import Path
+from typing import Any
+from cog import BasePredictor, Input, Path
 from tensorflow.keras.applications.resnet50 import ResNet50
-from tensorflow.keras.preprocessing import image
+from tensorflow.keras.preprocessing import image as keras_image
 from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
 import numpy as np
 
 
-class ResNetPredictor(cog.Predictor):
+class ResNetPredictor(BasePredictor):
     def setup(self):
         """Load the model into memory to make running multiple predictions efficient"""
         self.model = ResNet50(weights='resnet50_weights_tf_dim_ordering_tf_kernels.h5')
 
     # Define the arguments and types the model takes as input
-    @cog.input("input", type=Path, help="Image to classify")
-    def predict(self, input):
+    def predict(self, image: Path = Input(description="Image to classify")) -> Any:
         """Run a single prediction on the model"""
         # Preprocess the image
-        img = image.load_img(input, target_size=(224, 224))
-        x = image.img_to_array(img)
+        img = keras_image.load_img(image, target_size=(224, 224))
+        x = keras_image.img_to_array(img)
         x = np.expand_dims(x, axis=0)
         x = preprocess_input(x)
         # Run the prediction
         preds = self.model.predict(x)
         # Return the top 3 predictions
-        return str(decode_predictions(preds, top=3)[0])
+        return decode_predictions(preds, top=3)[0]
 ```
 
 We also need to point Cog at this, and tell it what Python dependencies to install. Update `cog.yaml` to look like this: