Strawman: FastAPI model server

[bfirsh]

This is potential design that could hit three interconnected birds with one stone:

• 🐣 Revisit design of Python type annotations #193
• 🐥 Design type system & signature #205
• 🐔 Review model API design #259

Those tickets contain user data and requirements. Please read them for context.

This is a strawman. It is most likely wrong. Please attack with your sharpest pitchforks!

Model definition

The model API is specified using Pydantic types in the predict function.

import cog
from pathlib import Path
from pydantic import Field
import torch


class Predictor(cog.Predictor):
    def setup(self):
        self.net = torch.load("weights.pth")

    def predict(
        self,
        input: Path = Field(..., title="Image to enlarge"),
        scale: float = Field(1.5, description="Factor to scale image by"),
    ) -> Path:
        """Run a single prediction on the model"""
        # ... pre-processing ...
        output = self.net(input)
        # ... post-processing ...
        return output

You could imagine some custom types like cog.Field if we need custom behavior, or don't want the extra import.

This predictor serves:

GET /predict

Takes a JSON object in request body:

{
  "input": "https://bucket.s3.amazon.com/...", # or "data:foo/bar;base64,..."
  "scale": 1.7,
}

Returns:

{
  "output": "https://bucket.s3.amazon.com/..."
}

Everything is JSON. Files are represented as URLs -- HTTP, cloud storage, data: URLs, or perhaps even file:// for mounted storage. This is for simplicity, so it can be represented sensibly in OpenAPI, and so more complex data structures can be represented.

If a more complex Pydantic object was defined as the return value, output would be a JSON object. The root object can also include other information about the run -- some prior art in #259 might be informative about what we want to include (metrics, etc).

Authentication and configuration for cloud storage needs some thought. It could specify a prefix for storing output files in the request.

GET /docs

A Swagger UI, generated automatically by FastAPI.

GET /openapi.json

An OpenAPI definition, generated automatically by FastAPI. This is used as the type definition. It contains a complete documentation of the HTTP API and entries like this to define the types for input/output:

{
  "components": {
    "schemas": {
      "Input": {
        "title": "Input",
        "required": [
          "input"
        ],
        "type": "object",
        "properties": {
          "input": {
            "title": "Image to enlarge",
            "type": "string",
            "format": "path"
          },
          "scale": {
            "title": "Scale",
            "type": "number",
            "description": "Factor to scale image by",
            "default": 1.5
          }
        }
      },
      "Output": ...
    }
  }
}

How this would be used for Replicate

We could either:

1. Integrate with existing redis queue worker.
2. Run a sidecar container that converts the current Redis API into HTTP.

Over time, we could define a queue API and include that in Cog. (Or leave it as an external system. See discussion in #259.)

Files

Files are represented as URLs. data:, http:, gcs:, s3:, etc.

This works well for input, but some way of specifying where output files go is required. Predictions could produce arbitrary output files (or should they be able to?) so some generic way of defining where multiple files can go is needed. For filesystem or bucket, this could be a prefix, and the prediction is responsible for naming files such that they don't collide.

Perhaps prediction requests could include the output path prefix:

{
  "input": {"image": "s3://predictions/1/input/image.jpg"},
  "output_path": "s3://predictions/1/output/",
}

For Replicate, files could be fetched and uploaded from Cloud Storage using signed URLs with a prefix. Unfortunately signed URLs can't be made for prefixes. You can assign IAM permissions for a prefix, but that would mean a service account for each prediction, which is unwieldy.

Some potential solutions:

1. We use serving's existing file storage system
2. We send files directly back to replicate-web, bypassing serving's file storage system. Perhaps if
3. We use data URIs...?!

Future

• An asynchronous API, where a job is started and you can query the status of the job. (This might be required for passing logs to Replicate?)
• This request/response can be passed over a queue. The same type definition could be used.
• How would do we support gRPC? Would we want to generate a generic protobuf schema that supports any kind of model, or generate one per-model somehow?
• We might want a more FastAPI-like API instead of the class-based Predictor.
• Custom URL routes on the FastAPI server. These might be namespaced under /extensions to preserve the stability of the top-level API.

Reflection

FastAPI is very HTTP + JSON centric. OpenAPI is too. This is optimal for simple use-cases, but using more advanced/efficient systems like gRPC or queues may feel like an impedance mismatch.

Implementation details

The predictor in the example above would be converted into a FastAPI method like this:

app = FastAPI()

class Input(BaseModel):
    input = Path = Field(..., title="Image to enlarge")
    scale = float = Field(1.5, description="Factor to scale image by")

@app.post("/predict")
async def predict(input: Input):
    # ...
    return Path(...)

We would need to come up with a custom Path handler that knows how to convert it into URLs or base64 encoded strings.

[andreasjansson]

Thanks for putting this together!

Throwing some forks at this strawman:

• FastAPI can return files as FileResponse or StreamingResponse. How would you return multiple files?
  • multipart/form-data seems like a reasonable approach (see Support for returning multiple outputs #328), but I'm not sure that's supported in FastAPI.
• I assume FastAPI's StreamingResponse can handle files of arbitrary size, and only keeps small chunks in memory?
• How would we implement progressive outputs?
• How would we capture logs from setup() and predict()?
• Would this API still allow you to import a cog.Predictor in a plain Python project and call p = Predictor(); p.setup(); p.predict(...)?
• How can models be adapted for AI Platform, Cloud Dataflow, Seldon, etc., where k8s sidecars aren't applicable?
  • Can we design our API with plugin capabilities from the start, so we don't have to do another major redesign when we want to deploy Cog models on non-k8s targets?
  • Perhaps the queue worker could be implemented using some adapter pattern that lives outside of the main Predictor class?

[bfirsh]

FastAPI can return files as FileResponse or StreamingResponse. How would you return multiple files?

Everything is JSON. Files are represented as URLs -- HTTP, cloud storage, data: URLs, or perhaps even file:// for mounted storage.

This is for simplicity, so it can be represented sensibly in OpenAPI, and so more complex data structures can be represented.

We have user data here -- we fetch and store files on cloud storage, and two of our users is requesting it works this way too.

The GET /predict has some more detail.

How would we implement progressive outputs?

An asynchronous API -- see "future". Lambdas could be used as they current are. In addition, we could represent progressive output as a list of outputs in the response so it is persistent and consistent.

How would we capture logs from setup() and predict()?

Like we currently do, with an asynchronous API I expect. See "future".

Would this API still allow you to import a cog.Predictor in a plain Python project and call p = Predictor(); p.setup(); p.predict(...)?

Yeah, I don't think anything's going to change there.

How can models be adapted for AI Platform, Cloud Dataflow, Seldon, etc., where k8s sidecars aren't applicable?

This is touched on in #259.

[andreasjansson]

Lets step back for a minute and consider this use case:

You're an image sharing website and you have lots of user-generated images stored on GCS. You now want to classify all those images, as well as new images that come in to your service.

You find a permissively licensed image classifier model on Replicate that is perfect for this feature. You now want to download that model and run a massive batch job over all your existing images, as well as process new images that come in from your users.

You decide to use Cloud Dataflow for batch processing of existing images and AI Platform Serving for new images. How do you adapt the Cog model on Replicate so that it can be deployed on Dataflow and AI Platform?

As a thought experiment, how would you do this if the model were a simple trained scikit-learn model? In that case you could make two Docker images, one for Dataflow and one for AI Platform, that wrapped the scikit-learn model in a bit of boilerplate to implement the APIs. That boilerplate could be abstracted out into some "sklearn-deployment" library where you just pass a predictor and it figures it all out: sklearn_deployment.serve_ai_platform(my_sklearn_predictor).

If you want to read from GCS and write to BigTable, you could provide input and output handlers:

sklearn_deployment.serve_ai_platform(my_sklearn_predictor) \
    .with_input_handler(sklearn_deployment.google.cloud_storage) \
    .with_output_handler(sklearn_deployment.google.bigtable)

Since Cog models are tied to their environments, it's more complicated. But perhaps we (and the community) could provide Cog-aware libraries for serving Cog models in different environments, of which a FastAPI server is one instance.

Here's a strawman for this kind of architecture where the user provides a "server" class and input and output handlers in cog.yaml:

build:
  ...
predict: predict.py:Predictor
serve:
  server: FastAPI
  input_handler: google.gcs
  output_handler: google.big_table

This means that we'd be forced to decouple the model interface from the serving interface.

[bfirsh]

Agreed we should figure out a way of supporting multiple deployment environments. This is mostly relevant to #259 and discussed a bit in there, but I'll continue this here.

At a high-level, model packaging should be independent of model deployment. A researcher should be able to package a model and it should be able to be deployed anywhere.

This implies that the deployment configuration should happen after a model is packaged. In an ideal world, this would simply be done with a sidecar, where the standard interface on the researcher's model was translated into whatever system they have running.

But as you suggest, there are platforms like AI Platform where things need to be shipped as a single Docker container. So, at some point, we need some way adding a different interface inside a model that has already been created.

The question in my mind: does this abstraction live at the Python level or the HTTP level? You could imagine two ways of implementing this:

1. An additional Python library is added that imports Predictor translates it to the other API.
2. An additional binary entrypoint is added to the container that translates the HTTP API to the other API.

There feel like a few trade-offs here we could dig into.

A few additional thoughts:

1. I wonder whether we can defer this, but keep it in mind. Right now our user needs are a queue-based API (for us and one user) and an asynchronous HTTP API, all of which are able to run sidecars. We could build for that, with some kind of more generic system in mind.
2. I wonder whether there is prior art here. What systems is this similar to?

[bfirsh]

Wheee we did this ages ago