Skip to content

Commit

Permalink
Add test to ensure setup and predict use same event loop
Browse files Browse the repository at this point in the history
  • Loading branch information
Aron Carroll committed Dec 16, 2024
1 parent 6bf3e61 commit 0bea77c
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 4 deletions.
3 changes: 0 additions & 3 deletions python/cog/server/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import asyncio


class Predictor:
async def setup(self) -> None:
self.loop = asyncio.get_running_loop()
Expand Down
6 changes: 6 additions & 0 deletions python/tests/server/test_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]"
Expand Down

0 comments on commit 0bea77c

Please sign in to comment.