actually allow concurrent predictions and refactor runner

.github/workflows/ci.yaml

@@ -4,14 +4,17 @@ on:
   push:
     branches:
       - main
+      - async
     tags:
       - "**"
   pull_request:
     branches:
       - main
+      - async
   merge_group:
     branches:
       - main
+      - async
     types:
       - checks_requested
 jobs:

pyproject.toml

@@ -35,6 +35,7 @@ optional-dependencies = { "dev" = [
   "pillow",
   "pyright==1.1.347",
   "pytest",
+  "pytest-asyncio",
   "pytest-httpserver",
   "pytest-rerunfailures",
   "pytest-xdist",

python/cog/command/ast_openapi_schema.py

@@ -309,10 +309,15 @@ def find(obj: ast.AST, name: str) -> ast.AST:
     """Find a particular named node in a tree"""
     return next(node for node in ast.walk(obj) if getattr(node, "name", "") == name)
 
+
 if typing.TYPE_CHECKING:
-    AstVal: "typing.TypeAlias" = "int | float | complex | str | list[AstVal] | bytes | None"
+    AstVal: "typing.TypeAlias" = (
+        "int | float | complex | str | list[AstVal] | bytes | None"
+    )
     AstValNoBytes: "typing.TypeAlias" = "int | float | str | list[AstValNoBytes]"
-    JSONObject: "typing.TypeAlias" = "int | float | str | list[JSONObject] | JSONDict | None"
+    JSONObject: "typing.TypeAlias" = (
+        "int | float | str | list[JSONObject] | JSONDict | None"
+    )
     JSONDict: "typing.TypeAlias" = "dict[str, JSONObject]"
 
 
@@ -327,6 +332,7 @@ def to_serializable(val: "AstVal") -> "JSONObject":
     else:
         return val
 
+
 def get_value(node: ast.AST) -> "AstVal":
     """Return the value of constant or list of constants"""
     if isinstance(node, ast.Constant):
@@ -339,7 +345,7 @@ def get_value(node: ast.AST) -> "AstVal":
     if isinstance(node, (ast.List, ast.Tuple)):
         return [get_value(e) for e in node.elts]
     if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
-            return -typing.cast(typing.Union[int, float, complex], get_value(node.operand))
+        return -typing.cast(typing.Union[int, float, complex], get_value(node.operand))
     raise ValueError("Unexpected node type", type(node))
 
 
@@ -366,12 +372,13 @@ 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
     return list(zip(args, defaults))
 
+
 def parse_assignment(assignment: ast.AST) -> "None | tuple[str, JSONObject]":
     """Parse an assignment into an OpenAPI object property"""
     if isinstance(assignment, ast.AnnAssign):
@@ -403,7 +410,9 @@ def parse_class(classdef: ast.AST) -> "JSONDict":
     """Parse a class definition into an OpenAPI object"""
     assert isinstance(classdef, ast.ClassDef)
     properties = {
-        assignment[0]: assignment[1] for assignment in map(parse_assignment, classdef.body) if assignment
+        assignment[0]: assignment[1]
+        for assignment in map(parse_assignment, classdef.body)
+        if assignment
     }
     return {
         "title": classdef.name,
@@ -428,17 +437,19 @@ def resolve_name(node: ast.expr) -> str:
         return node.id
     if isinstance(node, ast.Index):
         # deprecated, but needed for py3.8
-        return resolve_name(node.value) # type: ignore
+        return resolve_name(node.value)  # type: ignore
     if isinstance(node, ast.Attribute):
         return node.attr
     if isinstance(node, ast.Subscript):
         return resolve_name(node.value)
     raise ValueError("Unexpected node type", type(node), ast.unparse(node))
 
 
-def parse_return_annotation(tree: ast.AST, fn: str = "predict") -> "tuple[JSONDict, JSONDict]":
+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:
@@ -550,7 +561,7 @@ def extract_info(code: str) -> "JSONDict":
         **return_schema,
     }
     # trust me, typechecker, I know BASE_SCHEMA
-    x: "JSONDict" = schema["components"]["schemas"] # type: ignore
+    x: "JSONDict" = schema["components"]["schemas"]  # type: ignore
     x.update(components)
     return schema
 

python/cog/predictor.py

@@ -1,15 +1,15 @@
 import enum
 import importlib.util
 import inspect
-import io
 import os.path
 import sys
 import types
 from abc import ABC, abstractmethod
-from collections.abc import Iterator
+from collections.abc import AsyncIterator, Iterator
 from pathlib import Path
 from typing import (
     Any,
+    Awaitable,
     Callable,
     Dict,
     List,
@@ -48,7 +48,9 @@
 
 
 class BasePredictor(ABC):
-    def setup(self, weights: Optional[Union[CogFile, CogPath]] = None) -> None:
+    def setup(
+        self, weights: Optional[Union[CogFile, CogPath]] = None
+    ) -> Optional[Awaitable[None]]:
         """
         An optional method to prepare the model so multiple predictions run efficiently.
         """
@@ -63,15 +65,25 @@ def predict(self, **kwargs: Any) -> Any:
 
 
 def run_setup(predictor: BasePredictor) -> None:
-    weights_type = get_weights_type(predictor.setup)
-
-    # No weights need to be passed, so just run setup() without any arguments.
-    if weights_type is None:
+    weights = get_weights_argument(predictor)
+    if weights:
+        predictor.setup(weights=weights)
+    else:
         predictor.setup()
-        return
 
-    weights: Union[io.IOBase, Path, None]
 
+async def run_setup_async(predictor: BasePredictor) -> None:
+    weights = get_weights_argument(predictor)
+    maybe_coro = predictor.setup(weights=weights) if weights else predictor.setup()
+    if maybe_coro:
+        return await maybe_coro
+
+
+def get_weights_argument(predictor: BasePredictor) -> Union[CogFile, CogPath, None]:
+    # by the time we get here we assume predictor has a setup method
+    weights_type = get_weights_type(predictor.setup)
+    if weights_type is None:
+        return None
     weights_url = os.environ.get("COG_WEIGHTS")
     weights_path = "weights"
 
@@ -81,30 +93,27 @@ def run_setup(predictor: BasePredictor) -> None:
     # TODO: CogFile/CogPath should have subclasses for each of the subtypes
     if weights_url:
         if weights_type == CogFile:
-            weights = cast(CogFile, CogFile.validate(weights_url))
-        elif weights_type == CogPath:
+            return cast(CogFile, CogFile.validate(weights_url))
+        if weights_type == CogPath:
             # TODO: So this can be a url. evil!
-            weights = cast(CogPath, CogPath.validate(weights_url))
-        else:
-            raise ValueError(
-                f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File and Path are supported"
-            )
-    elif os.path.exists(weights_path):
+            return cast(CogPath, CogPath.validate(weights_url))
+        raise ValueError(
+            f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File and Path are supported"
+        )
+    if os.path.exists(weights_path):
         if weights_type == CogFile:
-            weights = cast(CogFile, open(weights_path, "rb"))
-        elif weights_type == CogPath:
-            weights = CogPath(weights_path)
-        else:
-            raise ValueError(
-                f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File and Path are supported"
-            )
-    else:
-        weights = None
-
-    predictor.setup(weights=weights)
+            return cast(CogFile, open(weights_path, "rb"))
+        if weights_type == CogPath:
+            return CogPath(weights_path)
+        raise ValueError(
+            f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File and Path are supported"
+        )
+    return None
 
 
-def get_weights_type(setup_function: Callable[[Any], None]) -> Optional[Any]:
+def get_weights_type(
+    setup_function: Callable[[Any], Optional[Awaitable[None]]]
+) -> Optional[Any]:
     signature = inspect.signature(setup_function)
     if "weights" not in signature.parameters:
         return None
@@ -341,7 +350,7 @@ def predict(
         OutputType = signature.return_annotation
 
     # The type that goes in the response is a list of the yielded type
-    if get_origin(OutputType) is Iterator:
+    if get_origin(OutputType) in {Iterator, AsyncIterator}:
         # Annotated allows us to attach Field annotations to the list, which we use to mark that this is an iterator
         # https://pydantic-docs.helpmanual.io/usage/schema/#typingannotated-fields
         field = Field(**{"x-cog-array-type": "iterator"})  # type: ignore

python/cog/schema.py

@@ -1,3 +1,4 @@
+import secrets
 import typing as t
 from datetime import datetime
 from enum import Enum
@@ -36,7 +37,12 @@ class PredictionBaseModel(pydantic.BaseModel, extra=pydantic.Extra.allow):
 
 
 class PredictionRequest(PredictionBaseModel):
-    id: t.Optional[str]
+    # there's a problem here where the idempotent endpoint is supposed to
+    # let you pass id in the route and omit it from the input
+    # however this fills in the default
+    # maybe it should be allowed to be optional without the factory initially
+    # and be filled in later
+    id: str = pydantic.Field(default_factory=lambda: secrets.token_hex(4))
     created_at: t.Optional[datetime]
 
     # TODO: deprecate this
@@ -85,8 +91,10 @@ def with_types(cls, input_type: t.Type[t.Any], output_type: t.Type[t.Any]) -> t.
             output=(output_type, None),
         )
 
+
 class TrainingRequest(PredictionRequest):
     pass
 
+
 class TrainingResponse(PredictionResponse):
     pass

python/cog/server/eventtypes.py

@@ -1,13 +1,31 @@
-from typing import Any, Dict
+import secrets
+from typing import Any, Dict, Union
 
 from attrs import define, field, validators
 
+from .. import schema, types
+
 
 # From worker parent process
 #
 @define
 class PredictionInput:
     payload: Dict[str, Any]
+    id: str = field(factory=lambda: secrets.token_hex(4))
+
+    @classmethod
+    def from_request(cls, request: schema.PredictionRequest) -> "PredictionInput":
+        assert request.id, "PredictionRequest must have an id"
+        payload = request.dict()["input"]
+        for k, v in payload.items():
+            if isinstance(v, types.URLPath):
+                payload[k] = v.convert()
+        return cls(payload=payload, id=request.id)
+
+
+@define
+class Cancel:
+    id: str
 
 
 @define
@@ -43,3 +61,6 @@ class Done:
 @define
 class Heartbeat:
     pass
+
+
+PublicEventType = Union[Done, Heartbeat, Log, PredictionOutput, PredictionOutputType]