Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Async setup #2089

Merged
merged 4 commits into from
Dec 18, 2024
Merged

Async setup #2089

merged 4 commits into from
Dec 18, 2024

Conversation

aron
Copy link
Contributor

@aron aron commented Dec 12, 2024

This commit introduces the ability to define an async setup function on your predictor. For simplicity an async setup() function is only supported alongside an async predict() function. An error will be raised during setup if this is not the case.

Various pieces of the code have been extracted into smaller methods in order to achieve this. A new _handle_setup_error context manager has been created to handle setup errors and send appropriate Done event over the worker channel.

The _setup() method has been split into two phases, first we perform validation on the requirements for async/concurrency support. Then we attempt to run the setup() method either as a direct call for the non-async path or as part of the event loop in the async path.

TODO

  • Add tests for async setup
  • Perform async validation at build time, we should be able to perform the concurrency checks as part of the app start-up which would mean errors are caught during build/publish.

@aron aron requested a review from philandstuff December 12, 2024 16:25
@aron aron force-pushed the async-setup branch 3 times, most recently from d355d91 to 9ae118a Compare December 12, 2024 21:31
@philandstuff philandstuff self-assigned this Dec 13, 2024
Copy link
Contributor

@philandstuff philandstuff left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great! I especially like the refactoring of ChildWorker into smaller methods. A couple of small comments.

python/cog/server/worker.py Outdated Show resolved Hide resolved
python/tests/server/test_worker.py Show resolved Hide resolved
Aron Carroll added 2 commits December 16, 2024 21:14
This is in preperation for supporting async setup functions we now
extract the weights for the setup function seperately from running the
setup function.
Aron Carroll added 2 commits December 18, 2024 13:19
This commit introduces the ability to define an async `setup` function
on your predictor. For simplicity an async `setup()` function is only
supported alongside an async `predict()` function. An error will be
raised during setup if this is not the case.

Various pieces of the code have been extracted into smaller methods
in order to achieve this. A new `_handle_setup_error` context manager
has been created to handle setup errors and send appropriate `Done`
event over the worker channel.

The `_setup()` method has been split into two phases, first we perform
validation on the requirements for async/concurrency support. Then we
attempt to run the `setup()` method either as a direct call for the
non-async path or as part of the event loop in the async path.
@aron aron merged commit e3e2ce8 into main Dec 18, 2024
19 checks passed
@aron aron deleted the async-setup branch December 18, 2024 15:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants