don't use async for predict loop if predict is not async

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

Author: technillogue

python/cog/predictor.py

@@ -51,7 +51,7 @@
 class BasePredictor(ABC):
     def setup(
         self, weights: Optional[Union[CogFile, CogPath]] = None
-    ) -> Union[Awaitable[None], None]:
+    ) -> Optional[Awaitable[None]]:
         """
         An optional method to prepare the model so multiple predictions run efficiently.
         """
@@ -68,8 +68,9 @@ def predict(self, **kwargs: Any) -> Any:
 def run_setup(predictor: BasePredictor) -> None:
     weights = get_weights_argument(predictor)
     if weights:
-        return predictor.setup()
-    return predictor.setup(weights=weights)
+        predictor.setup(weights=weights)
+    else:
+        predictor.setup()
 
 
 async def run_setup_async(predictor: BasePredictor) -> None:
@@ -80,6 +81,7 @@ async def run_setup_async(predictor: BasePredictor) -> None:
 
 
 def get_weights_argument(predictor: BasePredictor) -> Union[io.IOBase, 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
@@ -110,7 +112,7 @@ def get_weights_argument(predictor: BasePredictor) -> Union[io.IOBase, CogPath,
     return None
 
 
-def get_weights_type(predictor: BasePredictor) -> Optional[Any]:
+def get_weights_type(setup_function: Callable) -> Optional[Any]:
     signature = inspect.signature(setup_function)
     if "weights" not in signature.parameters:
         return None

python/cog/server/worker.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import inspect
 import multiprocessing
 import os
@@ -8,7 +9,7 @@
 import types
 from enum import Enum, auto, unique
 from multiprocessing.connection import Connection
-from typing import Any, Dict, Iterable, Optional, TextIO, Union
+from typing import Any, Dict, Iterator, Optional, TextIO, Union
 
 from ..json import make_encodeable
 from ..predictor import (
@@ -58,7 +59,7 @@ def __init__(self, predictor_ref: str, tee_output: bool = True) -> None:
         self._child = _ChildWorker(predictor_ref, child_events, tee_output)
         self._terminating = False
 
-    def setup(self) -> Iterable[_PublicEventType]:
+    def setup(self) -> Iterator[_PublicEventType]:
         self._assert_state(WorkerState.NEW)
         self._state = WorkerState.STARTING
         self._child.start()
@@ -67,7 +68,7 @@ def setup(self) -> Iterable[_PublicEventType]:
 
     def predict(
         self, payload: Dict[str, Any], poll: Optional[float] = None
-    ) -> Iterable[_PublicEventType]:
+    ) -> Iterator[_PublicEventType]:
         self._assert_state(WorkerState.READY)
         self._state = WorkerState.PROCESSING
         self._allow_cancel = True
@@ -108,7 +109,7 @@ def _assert_state(self, state: WorkerState) -> None:
 
     def _wait(
         self, poll: Optional[float] = None, raise_on_error: Optional[str] = None
-    ) -> Iterable[_PublicEventType]:
+    ) -> Iterator[_PublicEventType]:
         done = None
 
         if poll:
@@ -178,15 +179,21 @@ def run(self) -> None:
             [ws_stdout, ws_stderr], self._stream_write_hook
         )
         self._stream_redirector.start()
-
         self._setup()
-        asyncio.run(self._loop())
+        self._loop()
         self._stream_redirector.shutdown()
 
     def _setup(self) -> None:
-        done = Done()
-        try:
+        with self._handle_setup_error():
+            # we need to load the predictor to know if setup is async
             self._predictor = load_predictor_from_ref(self._predictor_ref)
+            # if users want to access the same event loop from setup and predict,
+            # both have to be async. if setup isn't async, it doesn't matter if we
+            # create the event loop here or after setup
+            #
+            # otherwise, if setup is sync and the user does new_event_loop to use a ClientSession,
+            # then tries to use the same session from async predict, they would get an error.
+            # that's significant if connections are open and would need to be discarded
             if is_async_predictor(self._predictor):
                 self.loop = get_loop()
             # Could be a function or a class
@@ -195,6 +202,12 @@ def _setup(self) -> None:
                     self.loop.run_until_complete(run_setup_async(self._predictor))
                 else:
                     run_setup(self._predictor)
+
+    @contextlib.contextmanager
+    def _handle_setup_error(self) -> Iterator[None]:
+        done = Done()
+        try:
+            yield
         except Exception as e:
             traceback.print_exc()
             done.error = True
@@ -210,50 +223,76 @@ def _setup(self) -> None:
             self._stream_redirector.drain()
             self._events.send(done)
 
-    async def _loop(self) -> None:
+    def _loop_sync(self) -> None:
         while True:
             ev = self._events.recv()
             if isinstance(ev, Shutdown):
                 break
             if isinstance(ev, PredictionInput):
-                await self._predict(ev.payload)
+                self._predict_sync(ev.payload)
             else:
                 print(f"Got unexpected event: {ev}", file=sys.stderr)
 
-    async def _predict(self, payload: Dict[str, Any]) -> None:
+    async def _loop_async(self) -> None:
+        while True:
+            ev = self._events.recv()
+            if isinstance(ev, Shutdown):
+                break
+            if isinstance(ev, PredictionInput):
+                await self._predict_async(ev.payload)
+            else:
+                print(f"Got unexpected event: {ev}", file=sys.stderr)
+
+    def _loop(self) -> None:
+        if is_async(get_predict(self._predictor)):
+            self.loop.run_until_complete(self._loop_async())
+        else:
+            self._loop_sync()
+
+    @contextlib.contextmanager
+    def _handle_predict_error(self) -> Iterator[None]:
         assert self._predictor
         done = Done()
         self._cancelable = True
         try:
+            yield
+        except CancelationException:
+            done.canceled = True
+        except Exception as e:
+            traceback.print_exc()
+            done.error = True
+            done.error_detail = str(e)
+        finally:
+            self._cancelable = False
+        self._stream_redirector.drain()
+        self._events.send(done)
+
+    async def _predict_async(self, payload: Dict[str, Any]) -> None:
+        with self._handle_predict_error():
             predict = get_predict(self._predictor)
             result = predict(**payload)
-
             if result:
                 if inspect.isasyncgen(result):
                     self._events.send(PredictionOutputType(multi=True))
                     async for r in result:
                         self._events.send(PredictionOutput(payload=make_encodeable(r)))
-                elif inspect.isgenerator(result):
-                    self._events.send(PredictionOutputType(multi=True))
-                    for r in result:
-                        self._events.send(PredictionOutput(payload=make_encodeable(r)))
                 elif inspect.isawaitable(result):
                     result = await result
                     self._events.send(PredictionOutputType(multi=False))
                     self._events.send(PredictionOutput(payload=make_encodeable(result)))
+
+    def _predict_sync(self, payload: Dict[str, Any]) -> None:
+        with self._handle_predict_error():
+            predict = get_predict(self._predictor)
+            result = predict(**payload)
+            if result:
+                if inspect.isgenerator(result):
+                    self._events.send(PredictionOutputType(multi=True))
+                    for r in result:
+                        self._events.send(PredictionOutput(payload=make_encodeable(r)))
                 else:
                     self._events.send(PredictionOutputType(multi=False))
                     self._events.send(PredictionOutput(payload=make_encodeable(result)))
-        except CancelationException:
-            done.canceled = True
-        except Exception as e:
-            traceback.print_exc()
-            done.error = True
-            done.error_detail = str(e)
-        finally:
-            self._cancelable = False
-        self._stream_redirector.drain()
-        self._events.send(done)
 
     def _signal_handler(self, signum: int, frame: Optional[types.FrameType]) -> None:
         if signum == signal.SIGUSR1 and self._cancelable:
@@ -270,13 +309,16 @@ def _stream_write_hook(
 
 def get_loop() -> asyncio.AbstractEventLoop:
     try:
+        # just in case something else created an event loop already
         return asyncio.get_running_loop()
     except RuntimeError:
         return asyncio.new_event_loop()
 
 
+def is_async(fn: Any) -> bool:
+    return inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn)
+
+
 def is_async_predictor(predictor: BasePredictor) -> bool:
-    predict = get_predict(predictor)
-    if inspect.iscoroutinefunction(predict) or inspect.isasyncgenfunction(predict):
-        return True
-    return inspect.iscoroutinefunction(getattr(predictor, "setup", None))
+    setup = getattr(predictor, "setup", None)
+    return is_async(setup) or is_async(get_predict(predictor))