create event loop before predictor setup (#1366)

* conditionally create the event loop if predictor is async, and add a path for hypothetical async setup
* don't use async for predict loop if predict is not async
* add test cases for shared loop across setup and predict + asyncio.run in setup

(reverts commit b533c6b)

Signed-off-by: technillogue <technillogue@gmail.com>

Author: technillogue

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))
 
 
@@ -372,6 +378,7 @@ def parse_args(tree: ast.AST) -> "list[tuple[ast.arg, ast.expr | types.EllipsisT
     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,15 +437,17 @@ 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, ast.AsyncFunctionDef)):
         raise ValueError("Could not find predict function")
@@ -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, AsyncIterator
+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

python/cog/server/http.py

@@ -247,7 +247,10 @@ async def healthcheck() -> Any:
         response_model=PredictionResponse,
         response_model_exclude_unset=True,
     )
-    async def predict(request: PredictionRequest = Body(default=None), prefer: Union[str, None] = Header(default=None)) -> Any:  # type: ignore
+    async def predict(
+        request: PredictionRequest = Body(default=None),
+        prefer: Union[str, None] = Header(default=None),
+    ) -> Any:  # type: ignore
         """
         Run a single prediction on the model
         """

python/cog/server/webhook.py

@@ -37,7 +37,8 @@ def _get_version() -> str:
 
 
 def webhook_caller_filtered(
-    webhook: str, webhook_events_filter: Set[WebhookEvent],
+    webhook: str,
+    webhook_events_filter: Set[WebhookEvent],
 ) -> Callable[[Any, WebhookEvent], None]:
     upstream_caller = webhook_caller(webhook)