diff --git a/python/cog/server/worker.py b/python/cog/server/worker.py index cf8b72bf2a..a91ca85d92 100644 --- a/python/cog/server/worker.py +++ b/python/cog/server/worker.py @@ -479,9 +479,6 @@ def _validate_predictor( with self._handle_setup_error(redirector): assert self._predictor - predict = get_predict(self._predictor) - - # Async models require python >= 3.11 so we can use asyncio.TaskGroup # We should check for this before getting to this point if self._has_async_predictor and sys.version_info < (3, 11): diff --git a/python/tests/server/fixtures/async_setup_uses_same_loop_as_predict.py b/python/tests/server/fixtures/async_setup_uses_same_loop_as_predict.py index f454034036..92501755d9 100644 --- a/python/tests/server/fixtures/async_setup_uses_same_loop_as_predict.py +++ b/python/tests/server/fixtures/async_setup_uses_same_loop_as_predict.py @@ -1,6 +1,5 @@ import asyncio - class Predictor: async def setup(self) -> None: self.loop = asyncio.get_running_loop() diff --git a/python/tests/server/test_worker.py b/python/tests/server/test_worker.py index bf7bfa619c..004f21d653 100644 --- a/python/tests/server/test_worker.py +++ b/python/tests/server/test_worker.py @@ -606,6 +606,12 @@ def test_graceful_shutdown(worker): assert fut.result() == Done() +@uses_worker("async_setup_uses_same_loop_as_predict", min_python=(3, 11), is_async=True) +def test_async_setup_uses_same_loop_as_predict(worker: Worker): + result = _process(worker, lambda: worker.predict({}), tag=None) + assert result, "Expected worker to return True to assert same event loop" + + @frozen class SetupState: fut: "Future[Done]"