ray.serve.batch return type problems with pylance and mypy #45032
Assignees
Comments
ArthurBook
added
bug
ArthurBook
changed the title
7 tasks
|
Hi @ArthurBook Thanks for raising the issue and start to contribute. Please do let us know when your PR is ready for review, I can help to tag the right folks! |
GeneDer
added
P2
|
Hi @GeneDer! |
edoakes
pushed a commit
that referenced
this issue
May 17, 2024
#45033) This PR solves issue [#45032](#45032) by implementing overloads that handle the return type inference for synchronous and asynchronous functions and methods decorated with `@ray.serve.batch`. The change improves compatibility with pylance, mypy and other type checkers, enhancing user-friendlyness. NOTE: No runtime changes are expected. # Before PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(integers: list[int]) -> list[int] NOTE: still expecting a list[int] return type @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] NOTE: still expecting a list[int] return type @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) batch_fn_sync_w_args: G@batch NOTE: G@batch is un-inferred by pylance @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) batch_fn_async_w_args: G@batch NOTE: G@batch is un-inferred by pylance class Server: @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # Argument of type "Literal[1]" cannot be assigned to parameter "integers" of type "list[int]" in function "batch_meth_sync_no_args" # (method) def batch_meth_sync_no_args(integers: list[int]) -> list[int] f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] g = Server().batch_meth_sync_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_async_w_args() -> R ``` In summary, before the PR, the typechecker (Pylance in the above case) is unable to match the `@ray.serve.batch` function signature of the functions and methods that it is intended to decorate. This causes the input and return type to be inferred incorrectly: 1) The decorated function input type is inferred as as a list, when it in fact should be a scalar. 2) The decorated function return type is inferred as as a list, when it in fact should be a scalar. 3) For the decorated method, 0 positional input args are expected 4) For the decorated method, the return type is the unbound TypeVar R. # After PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(int) -> int @ray.serve.batch async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(int) -> Coroutine[Any, Any, int] @ray.serve.batch(max_batch_size=2) def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) def batch_fn_sync_w_args(int) -> int @ray.serve.batch async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) def batch_fn_async_w_args(int) -> Coroutine[Any, Any, int] class Server: @ray.serve.batch def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # (method) def batch_meth_sync_no_args(int) -> int f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(int) -> Coroutine[Any, Any, int] g = Server().batch_meth_sync_w_args(1) # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # (method) def batch_meth_async_w_args(int) -> Coroutine[Any, Any, int] ``` <!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. --> <!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. --> ## Why are these changes needed? The return-type for a `@ray.serve.batch` decorated function is not inferred correctly, making mypy / pylance (probably other checkers too) lost the function signature for the decorated function/method. <!-- Please give a short summary of the change and the problem this solves. --> ## Related issue number Closes #45032 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run `scripts/format.sh` to lint the changes in this PR. - [x] I've included any doc changes needed for https://docs.ray.io/en/master/. - [x] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [x] Release tests --------- Signed-off-by: Arthur <atte.book@gmail.com> Signed-off-by: Arthur Böök <49250723+ArthurBook@users.noreply.github.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this issue
Jun 6, 2024
ray-project#45033) This PR solves issue [ray-project#45032](ray-project#45032) by implementing overloads that handle the return type inference for synchronous and asynchronous functions and methods decorated with `@ray.serve.batch`. The change improves compatibility with pylance, mypy and other type checkers, enhancing user-friendlyness. NOTE: No runtime changes are expected. # Before PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(integers: list[int]) -> list[int] NOTE: still expecting a list[int] return type @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] NOTE: still expecting a list[int] return type @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) batch_fn_sync_w_args: G@batch NOTE: G@batch is un-inferred by pylance @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) batch_fn_async_w_args: G@batch NOTE: G@batch is un-inferred by pylance class Server: @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # Argument of type "Literal[1]" cannot be assigned to parameter "integers" of type "list[int]" in function "batch_meth_sync_no_args" # (method) def batch_meth_sync_no_args(integers: list[int]) -> list[int] f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] g = Server().batch_meth_sync_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_async_w_args() -> R ``` In summary, before the PR, the typechecker (Pylance in the above case) is unable to match the `@ray.serve.batch` function signature of the functions and methods that it is intended to decorate. This causes the input and return type to be inferred incorrectly: 1) The decorated function input type is inferred as as a list, when it in fact should be a scalar. 2) The decorated function return type is inferred as as a list, when it in fact should be a scalar. 3) For the decorated method, 0 positional input args are expected 4) For the decorated method, the return type is the unbound TypeVar R. # After PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(int) -> int @ray.serve.batch async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(int) -> Coroutine[Any, Any, int] @ray.serve.batch(max_batch_size=2) def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) def batch_fn_sync_w_args(int) -> int @ray.serve.batch async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) def batch_fn_async_w_args(int) -> Coroutine[Any, Any, int] class Server: @ray.serve.batch def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # (method) def batch_meth_sync_no_args(int) -> int f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(int) -> Coroutine[Any, Any, int] g = Server().batch_meth_sync_w_args(1) # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # (method) def batch_meth_async_w_args(int) -> Coroutine[Any, Any, int] ``` <!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. --> <!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. --> ## Why are these changes needed? The return-type for a `@ray.serve.batch` decorated function is not inferred correctly, making mypy / pylance (probably other checkers too) lost the function signature for the decorated function/method. <!-- Please give a short summary of the change and the problem this solves. --> ## Related issue number Closes ray-project#45032 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run `scripts/format.sh` to lint the changes in this PR. - [x] I've included any doc changes needed for https://docs.ray.io/en/master/. - [x] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [x] Release tests --------- Signed-off-by: Arthur <atte.book@gmail.com> Signed-off-by: Arthur Böök <49250723+ArthurBook@users.noreply.github.com> Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this issue
Jun 6, 2024
ray-project#45033) This PR solves issue [ray-project#45032](ray-project#45032) by implementing overloads that handle the return type inference for synchronous and asynchronous functions and methods decorated with `@ray.serve.batch`. The change improves compatibility with pylance, mypy and other type checkers, enhancing user-friendlyness. NOTE: No runtime changes are expected. # Before PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(integers: list[int]) -> list[int] NOTE: still expecting a list[int] return type @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] NOTE: still expecting a list[int] return type @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) batch_fn_sync_w_args: G@batch NOTE: G@batch is un-inferred by pylance @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) batch_fn_async_w_args: G@batch NOTE: G@batch is un-inferred by pylance class Server: @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # Argument of type "Literal[1]" cannot be assigned to parameter "integers" of type "list[int]" in function "batch_meth_sync_no_args" # (method) def batch_meth_sync_no_args(integers: list[int]) -> list[int] f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] g = Server().batch_meth_sync_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_async_w_args() -> R ``` In summary, before the PR, the typechecker (Pylance in the above case) is unable to match the `@ray.serve.batch` function signature of the functions and methods that it is intended to decorate. This causes the input and return type to be inferred incorrectly: 1) The decorated function input type is inferred as as a list, when it in fact should be a scalar. 2) The decorated function return type is inferred as as a list, when it in fact should be a scalar. 3) For the decorated method, 0 positional input args are expected 4) For the decorated method, the return type is the unbound TypeVar R. # After PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(int) -> int @ray.serve.batch async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(int) -> Coroutine[Any, Any, int] @ray.serve.batch(max_batch_size=2) def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) def batch_fn_sync_w_args(int) -> int @ray.serve.batch async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) def batch_fn_async_w_args(int) -> Coroutine[Any, Any, int] class Server: @ray.serve.batch def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # (method) def batch_meth_sync_no_args(int) -> int f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(int) -> Coroutine[Any, Any, int] g = Server().batch_meth_sync_w_args(1) # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # (method) def batch_meth_async_w_args(int) -> Coroutine[Any, Any, int] ``` <!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. --> <!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. --> ## Why are these changes needed? The return-type for a `@ray.serve.batch` decorated function is not inferred correctly, making mypy / pylance (probably other checkers too) lost the function signature for the decorated function/method. <!-- Please give a short summary of the change and the problem this solves. --> ## Related issue number Closes ray-project#45032 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run `scripts/format.sh` to lint the changes in this PR. - [x] I've included any doc changes needed for https://docs.ray.io/en/master/. - [x] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [x] Release tests --------- Signed-off-by: Arthur <atte.book@gmail.com> Signed-off-by: Arthur Böök <49250723+ArthurBook@users.noreply.github.com> Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this issue
Jun 7, 2024
ray-project#45033) This PR solves issue [ray-project#45032](ray-project#45032) by implementing overloads that handle the return type inference for synchronous and asynchronous functions and methods decorated with `@ray.serve.batch`. The change improves compatibility with pylance, mypy and other type checkers, enhancing user-friendlyness. NOTE: No runtime changes are expected. # Before PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(integers: list[int]) -> list[int] NOTE: still expecting a list[int] return type @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] NOTE: still expecting a list[int] return type @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) batch_fn_sync_w_args: G@batch NOTE: G@batch is un-inferred by pylance @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) batch_fn_async_w_args: G@batch NOTE: G@batch is un-inferred by pylance class Server: @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # Argument of type "Literal[1]" cannot be assigned to parameter "integers" of type "list[int]" in function "batch_meth_sync_no_args" # (method) def batch_meth_sync_no_args(integers: list[int]) -> list[int] f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] g = Server().batch_meth_sync_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_async_w_args() -> R ``` In summary, before the PR, the typechecker (Pylance in the above case) is unable to match the `@ray.serve.batch` function signature of the functions and methods that it is intended to decorate. This causes the input and return type to be inferred incorrectly: 1) The decorated function input type is inferred as as a list, when it in fact should be a scalar. 2) The decorated function return type is inferred as as a list, when it in fact should be a scalar. 3) For the decorated method, 0 positional input args are expected 4) For the decorated method, the return type is the unbound TypeVar R. # After PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(int) -> int @ray.serve.batch async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(int) -> Coroutine[Any, Any, int] @ray.serve.batch(max_batch_size=2) def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) def batch_fn_sync_w_args(int) -> int @ray.serve.batch async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) def batch_fn_async_w_args(int) -> Coroutine[Any, Any, int] class Server: @ray.serve.batch def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # (method) def batch_meth_sync_no_args(int) -> int f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(int) -> Coroutine[Any, Any, int] g = Server().batch_meth_sync_w_args(1) # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # (method) def batch_meth_async_w_args(int) -> Coroutine[Any, Any, int] ``` <!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. --> <!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. --> ## Why are these changes needed? The return-type for a `@ray.serve.batch` decorated function is not inferred correctly, making mypy / pylance (probably other checkers too) lost the function signature for the decorated function/method. <!-- Please give a short summary of the change and the problem this solves. --> ## Related issue number Closes ray-project#45032 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run `scripts/format.sh` to lint the changes in this PR. - [x] I've included any doc changes needed for https://docs.ray.io/en/master/. - [x] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [x] Release tests --------- Signed-off-by: Arthur <atte.book@gmail.com> Signed-off-by: Arthur Böök <49250723+ArthurBook@users.noreply.github.com>
GabeChurch
pushed a commit
to GabeChurch/ray
that referenced
this issue
Jun 11, 2024
ray-project#45033) This PR solves issue [ray-project#45032](ray-project#45032) by implementing overloads that handle the return type inference for synchronous and asynchronous functions and methods decorated with `@ray.serve.batch`. The change improves compatibility with pylance, mypy and other type checkers, enhancing user-friendlyness. NOTE: No runtime changes are expected. # Before PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(integers: list[int]) -> list[int] NOTE: still expecting a list[int] return type @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] NOTE: still expecting a list[int] return type @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) batch_fn_sync_w_args: G@batch NOTE: G@batch is un-inferred by pylance @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) batch_fn_async_w_args: G@batch NOTE: G@batch is un-inferred by pylance class Server: @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) # No overloads for "batch" match the provided argumentsPylancereportCallIssue async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # Argument of type "Literal[1]" cannot be assigned to parameter "integers" of type "list[int]" in function "batch_meth_sync_no_args" # (method) def batch_meth_sync_no_args(integers: list[int]) -> list[int] f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] # (method) def batch_meth_async_no_args(integers: list[int]) -> Coroutine[Any, Any, list[int]] g = Server().batch_meth_sync_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # Expected 0 positional argumentsPylancereportCallIssue # (method) def batch_meth_async_w_args() -> R ``` In summary, before the PR, the typechecker (Pylance in the above case) is unable to match the `@ray.serve.batch` function signature of the functions and methods that it is intended to decorate. This causes the input and return type to be inferred incorrectly: 1) The decorated function input type is inferred as as a list, when it in fact should be a scalar. 2) The decorated function return type is inferred as as a list, when it in fact should be a scalar. 3) For the decorated method, 0 positional input args are expected 4) For the decorated method, the return type is the unbound TypeVar R. # After PR: ```python '''Test the type infernece for return value of a batched function or method with pylance.''' import ray.serve @ray.serve.batch def batch_fn_sync_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] a = batch_fn_sync_no_args(1) # (function) def batch_fn_sync_no_args(int) -> int @ray.serve.batch async def batch_fn_async_no_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] b = batch_fn_async_no_args(1) # (function) def batch_fn_async_no_args(int) -> Coroutine[Any, Any, int] @ray.serve.batch(max_batch_size=2) def batch_fn_sync_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] c = batch_fn_sync_w_args(1) # (function) def batch_fn_sync_w_args(int) -> int @ray.serve.batch async def batch_fn_async_w_args(integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] d = batch_fn_async_w_args(1) # (function) def batch_fn_async_w_args(int) -> Coroutine[Any, Any, int] class Server: @ray.serve.batch def batch_meth_sync_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch async def batch_meth_async_no_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) def batch_meth_sync_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] @ray.serve.batch(max_batch_size=2) async def batch_meth_async_w_args(self, integers: list[int]) -> list[int]: "docs" return [i * 2 for i in integers] e = Server().batch_meth_sync_no_args(1) # (method) def batch_meth_sync_no_args(int) -> int f = Server().batch_meth_async_no_args(1) # (method) def batch_meth_async_no_args(int) -> Coroutine[Any, Any, int] g = Server().batch_meth_sync_w_args(1) # (method) def batch_meth_sync_w_args(int) -> int h = Server().batch_meth_async_w_args(1) # (method) def batch_meth_async_w_args(int) -> Coroutine[Any, Any, int] ``` <!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. --> <!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. --> ## Why are these changes needed? The return-type for a `@ray.serve.batch` decorated function is not inferred correctly, making mypy / pylance (probably other checkers too) lost the function signature for the decorated function/method. <!-- Please give a short summary of the change and the problem this solves. --> ## Related issue number Closes ray-project#45032 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run `scripts/format.sh` to lint the changes in this PR. - [x] I've included any doc changes needed for https://docs.ray.io/en/master/. - [x] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [x] Release tests --------- Signed-off-by: Arthur <atte.book@gmail.com> Signed-off-by: Arthur Böök <49250723+ArthurBook@users.noreply.github.com> Signed-off-by: gchurch <gabe1church@gmail.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
What happened + What you expected to happen
Hi! first of all, thank you for the work on this amazing project. I am a huge fan.
There is a type-hint nit that has been bothering me in the ray.serve.batch decorator. It is causing some issues with pylance and mypy.
The issue stems from that the bounds on the typevars F and G have TypeVar bound types, which aren't supported. Another related issue is the behavior of decorating a method on a class with
@batch. Since the method has an implicit self argument, the Callable typehint is unable to match the signature.In other words, the return-type for a
@batchdecorated function is not inferred correctly, making mypy / pylance (probably other checkers too) lost the function signature for the decorated function/method.Versions / Dependencies
dependencies
Reproduction script
pylance (vscode):
mypy:
Issue Severity
Low: It annoys or frustrates me.
The text was updated successfully, but these errors were encountered: