Add support for async predictors

docs/metrics.md

@@ -0,0 +1,14 @@
+# Metrics
+
+Prediction objects have a `metrics` field. This normally includes `predict_time` and `total_time`. Official language models have metrics like `input_token_count`, `output_token_count`, `tokens_per_second`, and `time_to_first_token`. Currently, custom metrics from Cog are ignored when running on Replicate. Official Replicate-published models are the only exception to this. When running outside of Replicate, you can emit custom metrics like this:
+
+
+```python
+import cog
+from cog import BasePredictor, Path
+
+class Predictor(BasePredictor):
+    def predict(self, width: int, height: int) -> Path:
+        """Run a single prediction on the model"""
+        cog.emit_metric(name="pixel_count", value=width * height)
+```

pkg/config/config.go

@@ -57,16 +57,21 @@ type Build struct {
 	pythonRequirementsContent []string
 }
 
+type Concurrency struct {
+	Max int `json:"max,omitempty" yaml:"max"`
+}
+
 type Example struct {
 	Input  map[string]string `json:"input" yaml:"input"`
 	Output string            `json:"output" yaml:"output"`
 }
 
 type Config struct {
-	Build   *Build `json:"build" yaml:"build"`
-	Image   string `json:"image,omitempty" yaml:"image"`
-	Predict string `json:"predict,omitempty" yaml:"predict"`
-	Train   string `json:"train,omitempty" yaml:"train"`
+	Build       *Build       `json:"build" yaml:"build"`
+	Image       string       `json:"image,omitempty" yaml:"image"`
+	Predict     string       `json:"predict,omitempty" yaml:"predict"`
+	Train       string       `json:"train,omitempty" yaml:"train"`
+	Concurrency *Concurrency `json:"concurrency,omitempty" yaml:"concurrency"`
 }
 
 func DefaultConfig() *Config {

pkg/config/data/config_schema_v1.0.json

@@ -154,11 +154,6 @@
           "$id": "#/properties/concurrency/properties/max",
           "type": "integer",
           "description": "The maximum number of concurrent predictions."
-        },
-        "default_target": {
-          "$id": "#/properties/concurrency/properties/default_target",
-          "type": "integer",
-          "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."
         }
       }
     }

pyproject.toml

@@ -10,11 +10,13 @@ authors = [{ name = "Replicate", email = "[email protected]" }]
 license.file = "LICENSE"
 urls."Source" = "https://github.com/replicate/cog"
 
-requires-python = ">=3.7"
+requires-python = ">=3.8"
 dependencies = [
   # intentionally loose. perhaps these should be vendored to not collide with user code?
   "attrs>=20.1,<24",
   "fastapi>=0.75.2,<0.99.0",
+  # we may not need http2
+  "httpx[http2]>=0.21.0,<1",
   "pydantic>=1.9,<2",
   "PyYAML",
   "requests>=2,<3",
@@ -27,14 +29,15 @@ dependencies = [
 optional-dependencies = { "dev" = [
   "black",
   "build",
-  "httpx",
   'hypothesis<6.80.0; python_version < "3.8"',
   'hypothesis; python_version >= "3.8"',
+  "respx",
   'numpy<1.22.0; python_version < "3.8"',
   'numpy; python_version >= "3.8"',
   "pillow",
   "pyright==1.1.347",
   "pytest",
+  "pytest-asyncio",
   "pytest-httpserver",
   "pytest-rerunfailures",
   "pytest-xdist",
@@ -66,6 +69,21 @@ reportUnusedExpression = "warning"
 [tool.setuptools]
 package-dir = { "" = "python" }
 
+[tool.pylint.main]
+disable = [
+  "C0114", # Missing module docstring
+  "C0115", # Missing class docstring
+  "C0116", # Missing function or method docstring
+  "C0301", # Line too long
+  "C0413", # Import should be placed at the top of the module
+  "R0903", # Too few public methods
+  "W0622", # Redefining built-in
+]
+good-names = ["id", "input"]
+
+ignore-paths = ["python/cog/_version.py", "python/tests"]
+
+
 [tool.ruff]
 lint.select = [
   "E",   # pycodestyle error

python/cog/__init__.py

@@ -1,7 +1,15 @@
 from pydantic import BaseModel
 
 from .predictor import BasePredictor
-from .types import ConcatenateIterator, File, Input, Path, Secret
+from .server.worker import emit_metric
+from .types import (
+    AsyncConcatenateIterator,
+    ConcatenateIterator,
+    File,
+    Input,
+    Path,
+    Secret,
+)
 
 try:
     from ._version import __version__
@@ -14,8 +22,10 @@
     "BaseModel",
     "BasePredictor",
     "ConcatenateIterator",
+    "AsyncConcatenateIterator",
     "File",
     "Input",
     "Path",
     "Secret",
+    "emit_metric",
 ]

python/cog/code_xforms.py

@@ -12,7 +12,7 @@ def load_module_from_string(
     if not source or not name:
         return None
     module = types.ModuleType(name)
-    exec(source, module.__dict__)  # noqa: S102
+    exec(source, module.__dict__)  # noqa: S102 # pylint: disable=exec-used
     return module
 
 
@@ -32,7 +32,7 @@ class ClassExtractor(ast.NodeVisitor):
         def __init__(self) -> None:
             self.class_source = None
 
-        def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        def visit_ClassDef(self, node: ast.ClassDef) -> None:  # pylint: disable=invalid-name
             if node.name in all_class_names:
                 self.class_source = ast.get_source_segment(source_code, node)
 
@@ -56,7 +56,7 @@ class FunctionExtractor(ast.NodeVisitor):
         def __init__(self) -> None:
             self.function_source = None
 
-        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # pylint: disable=invalid-name
             if node.name == function_name and not isinstance(node, ast.Module):
                 # Extract the source segment for this function definition
                 self.function_source = ast.get_source_segment(source_code, node)
@@ -79,14 +79,16 @@ def make_class_methods_empty(source_code: Union[str, ast.AST], class_name: str)
     """
 
     class MethodBodyTransformer(ast.NodeTransformer):
-        def visit_ClassDef(self, node: ast.ClassDef) -> Optional[ast.AST]:
+        def visit_ClassDef(self, node: ast.ClassDef) -> Optional[ast.AST]:  # pylint: disable=invalid-name
             if node.name == class_name:
                 for body_item in node.body:
                     if isinstance(body_item, ast.FunctionDef):
                         # Replace the body of the method with `return None`
                         body_item.body = [ast.Return(value=ast.Constant(value=None))]
                 return node
 
+            return None
+
     tree = source_code if isinstance(source_code, ast.AST) else ast.parse(source_code)
     transformer = MethodBodyTransformer()
     transformed_tree = transformer.visit(tree)
@@ -111,11 +113,11 @@ class MethodReturnTypeExtractor(ast.NodeVisitor):
         def __init__(self) -> None:
             self.return_type = None
 
-        def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        def visit_ClassDef(self, node: ast.ClassDef) -> None:  # pylint: disable=invalid-name
             if node.name == class_name:
                 self.generic_visit(node)
 
-        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # pylint: disable=invalid-name
             if node.name == method_name and node.returns:
                 self.return_type = ast.unparse(node.returns)
 
@@ -142,7 +144,7 @@ class FunctionReturnTypeExtractor(ast.NodeVisitor):
         def __init__(self) -> None:
             self.return_type = None
 
-        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # pylint: disable=invalid-name
             if node.name == function_name and node.returns:
                 # Extract and return the string representation of the return type
                 self.return_type = ast.unparse(node.returns)
@@ -166,12 +168,14 @@ def make_function_empty(source_code: Union[str, ast.AST], function_name: str) ->
     """
 
     class FunctionBodyTransformer(ast.NodeTransformer):
-        def visit_FunctionDef(self, node: ast.FunctionDef) -> Optional[ast.AST]:
+        def visit_FunctionDef(self, node: ast.FunctionDef) -> Optional[ast.AST]:  # pylint: disable=invalid-name
             if node.name == function_name:
                 # Replace the body of the function with `return None`
                 node.body = [ast.Return(value=ast.Constant(value=None))]
                 return node
 
+            return None
+
     tree = source_code if isinstance(source_code, ast.AST) else ast.parse(source_code)
     transformer = FunctionBodyTransformer()
     transformed_tree = transformer.visit(tree)
@@ -195,12 +199,12 @@ class ImportExtractor(ast.NodeVisitor):
         def __init__(self) -> None:
             self.imports = []
 
-        def visit_Import(self, node: ast.Import) -> None:
+        def visit_Import(self, node: ast.Import) -> None:  # pylint: disable=invalid-name
             for alias in node.names:
                 if alias.name in module_names:
                     self.imports.append(ast.unparse(node))
 
-        def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        def visit_ImportFrom(self, node: ast.ImportFrom) -> None:  # pylint: disable=invalid-name
             if node.module in module_names:
                 self.imports.append(ast.unparse(node))
 

python/cog/command/ast_openapi_schema.py

@@ -147,6 +147,24 @@
         "summary": "Healthcheck"
       }
     },
+    "/ready": {
+      "get": {
+        "summary": "Ready",
+        "operationId": "ready_ready_get",
+        "responses": {
+          "200": {
+            "description": "Successful Response",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "title": "Response Ready Ready Get"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
     "/predictions": {
       "post": {
         "description": "Run a single prediction on the model",
@@ -324,13 +342,12 @@ def find(obj: ast.AST, name: str) -> ast.AST:
 def to_serializable(val: "AstVal") -> "JSONObject":
     if isinstance(val, bytes):
         return val.decode("utf-8")
-    elif isinstance(val, list):
+    if isinstance(val, list):
         return [to_serializable(x) for x in val]
-    elif isinstance(val, complex):
+    if isinstance(val, complex):
         msg = "complex inputs are not supported"
         raise ValueError(msg)
-    else:
-        return val
+    return val
 
 
 def get_value(node: ast.AST) -> "AstVal":
@@ -372,7 +389,7 @@ def get_call_name(call: ast.Call) -> str:
 def parse_args(tree: ast.AST) -> "list[tuple[ast.arg, ast.expr | types.EllipsisType]]":
     """Parse argument, default pairs from a file with a predict function"""
     predict = find(tree, "predict")
-    assert isinstance(predict, ast.FunctionDef)
+    assert isinstance(predict, (ast.FunctionDef, ast.AsyncFunctionDef))
     args = predict.args.args  # [-len(defaults) :]
     # use Ellipsis instead of None here to distinguish a default of None
     defaults = [...] * (len(args) - len(predict.args.defaults)) + predict.args.defaults
@@ -449,7 +466,7 @@ def parse_return_annotation(
     tree: ast.AST, fn: str = "predict"
 ) -> "tuple[JSONDict, JSONDict]":
     predict = find(tree, fn)
-    if not isinstance(predict, ast.FunctionDef):
+    if not isinstance(predict, (ast.FunctionDef, ast.AsyncFunctionDef)):
         raise ValueError("Could not find predict function")
     annotation = predict.returns
     if not annotation:
@@ -472,8 +489,8 @@ def predict(
     name = resolve_name(annotation)
     if isinstance(annotation, ast.Subscript):
         # forget about other subscripts like Optional, and assume otherlib.File will still be an uri
-        slice = resolve_name(annotation.slice)
-        format = {"format": "uri"} if slice in ("Path", "File") else {}
+        slice = resolve_name(annotation.slice)  # pylint: disable=redefined-builtin
+        format = {"format": "uri"} if slice in ("Path", "File") else {}  # pylint: disable=redefined-builtin
         array_type = {"x-cog-array-type": "iterator"} if "Iterator" in name else {}
         display_type = (
             {"x-cog-array-display": "concatenate"} if "Concatenate" in name else {}
@@ -503,7 +520,7 @@ def predict(
 KEPT_ATTRS = ("description", "default", "ge", "le", "max_length", "min_length", "regex")
 
 
-def extract_info(code: str) -> "JSONDict":
+def extract_info(code: str) -> "JSONDict":  # pylint: disable=too-many-branches,too-many-locals
     """Parse the schemas from a file with a predict function"""
     tree = ast.parse(code)
     properties: JSONDict = {}
@@ -526,7 +543,7 @@ def extract_info(code: str) -> "JSONDict":
             kws = {}
         else:
             raise ValueError("Unexpected default value", default)
-        input: JSONDict = {"x-order": len(properties)}
+        input: JSONDict = {"x-order": len(properties)}  # pylint: disable=redefined-builtin
         # need to handle other types?
         arg_type = OPENAPI_TYPES.get(get_annotation(arg.annotation), "string")
         if get_annotation(arg.annotation) in ("Path", "File"):