Inspect predictor to determine if predict/train is async

Author: aron

python/cog/config.py

@@ -1,3 +1,4 @@
+import inspect
 import os
 import sys
 import uuid
@@ -16,7 +17,9 @@
 from .predictor import (
     get_input_type,
     get_output_type,
+    get_predict,
     get_predictor,
+    get_train,
     get_training_input_type,
     get_training_output_type,
     load_full_predictor_from_file,
@@ -152,16 +155,26 @@ def get_predictor_ref(self, mode: Mode) -> str:
 
     def get_predictor_types(
         self, mode: Mode
-    ) -> Tuple[Type[BaseInput], Type[BaseModel]]:
+    ) -> Tuple[Type[BaseInput], Type[BaseModel], bool]:
         """Find the input and output types of a predictor."""
         predictor_ref = self.get_predictor_ref(mode=mode)
         predictor = self._load_predictor_for_types(
             predictor_ref, _method_name_from_mode(mode=mode), mode
         )
+
+        def is_async(fn) -> bool:
+            return inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn)
+
         if mode == Mode.PREDICT:
-            return get_input_type(predictor), get_output_type(predictor)
+            return (
+                get_input_type(predictor),
+                get_output_type(predictor),
+                is_async(get_predict(predictor)),
+            )
         elif mode == Mode.TRAIN:
-            return get_training_input_type(predictor), get_training_output_type(
-                predictor
+            return (
+                get_training_input_type(predictor),
+                get_training_output_type(predictor),
+                is_async(get_train(predictor)),
             )
         raise ValueError(f"Mode {mode} not found for generating input/output types.")

python/cog/server/http.py

@@ -156,7 +156,9 @@ async def start_shutdown() -> Any:
         return JSONResponse({}, status_code=200)
 
     try:
-        InputType, OutputType = cog_config.get_predictor_types(mode=Mode.PREDICT)
+        InputType, OutputType, is_async = cog_config.get_predictor_types(
+            mode=Mode.PREDICT
+        )
     except Exception:  # pylint: disable=broad-exception-caught
         msg = "Error while loading predictor:\n\n" + traceback.format_exc()
         add_setup_failed_routes(app, started_at, msg)

python/cog/server/worker.py

@@ -350,6 +350,7 @@ def __init__(
         predictor_ref: str,
         events: Connection,
         tee_output: bool = True,
+        is_async: bool = False,
     ) -> None:
         self._predictor_ref = predictor_ref
         self._predictor: Optional[BasePredictor] = None
@@ -361,6 +362,7 @@ def __init__(
 
         # for synchronous predictors only! async predictors use _tag_var instead
         self._sync_tag: Optional[str] = None
+        self._is_async = is_async
 
         super().__init__()
 
@@ -373,30 +375,34 @@ def run(self) -> None:
         # Initially, we ignore SIGUSR1.
         signal.signal(signal.SIGUSR1, signal.SIG_IGN)
 
-        self._predictor = self._load_predictor()
-
-        # If setup didn't set the predictor, we're done here.
-        if not self._predictor:
-            return
-
-        predict = get_predict(self._predictor)
-        if inspect.iscoroutinefunction(predict) or inspect.isasyncgenfunction(predict):
-            async_redirector = AsyncStreamRedirector(
+        if self._is_async:
+            redirector = AsyncStreamRedirector(
                 callback=self._stream_write_hook,
                 tee=self._tee_output,
             )
-            with async_redirector:
-                self._setup(async_redirector)
-                asyncio.run(self._aloop(predict, async_redirector))
         else:
-            # We use SIGUSR1 to signal an interrupt for cancelation.
-            signal.signal(signal.SIGUSR1, self._signal_handler)
-
             redirector = StreamRedirector(
                 callback=self._stream_write_hook,
                 tee=self._tee_output,
             )
-            with redirector:
+
+        with redirector:
+            self._predictor = self._load_predictor()
+
+            # If setup didn't set the predictor, we're done here.
+            if not self._predictor:
+                return
+
+            predict = get_predict(self._predictor)
+            if self._is_async:
+                assert isinstance(redirector, AsyncStreamRedirector)
+                self._setup(redirector)
+                asyncio.run(self._aloop(predict, redirector))
+            else:
+                # We use SIGUSR1 to signal an interrupt for cancelation.
+                signal.signal(signal.SIGUSR1, self._signal_handler)
+
+                assert isinstance(redirector, StreamRedirector)
                 self._setup(redirector)
                 self._loop(
                     predict,
@@ -706,10 +712,15 @@ def _stream_write_hook(self, stream_name: str, data: str) -> None:
 
 
 def make_worker(
-    predictor_ref: str, tee_output: bool = True, max_concurrency: int = 1
+    predictor_ref: str,
+    tee_output: bool = True,
+    max_concurrency: int = 1,
+    is_async: bool = False,
 ) -> Worker:
     parent_conn, child_conn = _spawn.Pipe()
-    child = _ChildWorker(predictor_ref, events=child_conn, tee_output=tee_output)
+    child = _ChildWorker(
+        predictor_ref, events=child_conn, tee_output=tee_output, is_async=is_async
+    )
     parent = Worker(child=child, events=parent_conn, max_concurrency=max_concurrency)
     return parent
 

python/tests/server/fixtures/logging.py

@@ -9,13 +9,15 @@
 # a bunch of stuff at import time.
 libc.puts(b"writing some stuff from C at import time")
 libc.fflush(None)
+
 sys.stdout.write("writing to stdout at import time\n")
 sys.stderr.write("writing to stderr at import time\n")
 
 
 class Predictor:
     def setup(self):
         print("setting up predictor")
+        print("writing to stderr at setup time", file=sys.stderr)
         self.foo = "foo"
 
     def predict(self) -> str: