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.

Sign up for GitHub

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

Custom protoc plugin to generate a grpclib wrapper #2181

Merged
merged 17 commits into from
Sep 4, 2024
Merged

Custom protoc plugin to generate a grpclib wrapper #2181

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
e61803f
Adds custom proto generator plugin
freider Sep 3, 2024
19e3490
Replace all usages
freider Sep 3, 2024
911f2f2
Plugin
freider Sep 3, 2024
926d3ba
Move plugin out of modal
freider Sep 3, 2024
098ffab
Don't pollute client cli namespace
freider Sep 3, 2024
492485c
Don't install plugin, just reference script
freider Sep 4, 2024
8d5b26a
Merge remote-tracking branch 'origin/main' into freider/grpc-plugin
freider Sep 4, 2024
858dcf4
Clean up generation
freider Sep 4, 2024
823a992
oops
freider Sep 4, 2024
608d740
Wrapper gen in ci
freider Sep 4, 2024
60c8ac2
Install grpclib before generating proto stubs
freider Sep 4, 2024
8686bf6
Copyright
freider Sep 4, 2024
d8a34cd
windows
freider Sep 4, 2024
e171dea
no pause
freider Sep 4, 2024
a4c4962
Remove prints
freider Sep 4, 2024
3d1c41e
Cleanup task
freider Sep 4, 2024
96b2370
Fix multi-file-access on windows
freider Sep 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion modal_proto/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +0,0 @@
# Copyright Modal Labs 2022
8 changes: 3 additions & 5 deletions protoc_plugin/plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,6 @@
from google.protobuf.descriptor_pb2 import DescriptorProto, FileDescriptorProto
from grpclib import const

from modal._utils.grpc_utils import UnaryStreamWrapper, UnaryUnaryWrapper

_CARDINALITY = {
(False, False): const.Cardinality.UNARY_UNARY,
(True, False): const.Cardinality.STREAM_UNARY,
Expand Down Expand Up @@ -86,9 +84,9 @@ def render(
name, cardinality, request_type, reply_type = method
wrapper_cls: type
if cardinality is const.Cardinality.UNARY_UNARY:
wrapper_cls = UnaryUnaryWrapper
wrapper_cls = "modal._utils.grpc_utils.UnaryUnaryWrapper"
Copy link
Contributor

Choose a reason for hiding this comment

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

Do we still need modal._utils.grpc_utils imported under a TYPE_CHECKING guard to use a forward reference? Never been totally clear on that.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

In this case it's not a forward reference - it's a "real" reference in the generated code and it's added as an import on line 175: https://github.com/modal-labs/modal-client/pull/2181/files#diff-d721170dbb2b36a3f20394f9563415ebe89a0916e4957cf99e7b615cfd8c772fR175

In case of forward/str references I think the TYPE_CHECKING-guarded imports are still required for type checkers to know what it's dealing with

Copy link
Contributor

Choose a reason for hiding this comment

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

whoops didn't read carefully enough!

elif cardinality is const.Cardinality.UNARY_STREAM:
wrapper_cls = UnaryStreamWrapper
wrapper_cls = "modal._utils.grpc_utils.UnaryStreamWrapper"
# elif cardinality is const.Cardinality.STREAM_UNARY:
# wrapper_cls = StreamUnaryWrapper
# elif cardinality is const.Cardinality.STREAM_STREAM:
Expand All @@ -97,7 +95,7 @@ def render(
raise TypeError(cardinality)

original_method = f"grpclib_stub.{name}"
buf.add(f"self.{name} = {wrapper_cls.__module__}.{wrapper_cls.__name__}({original_method})")
buf.add(f"self.{name} = {wrapper_cls}({original_method})")

return buf.content()

Expand Down
18 changes: 12 additions & 6 deletions tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,15 +25,21 @@

@task
def protoc(ctx):
grpc_plugin_path = Path(__file__).parent / "protoc_plugin" / "plugin.py"
protoc_cmd = f"{sys.executable} -m grpc_tools.protoc"
input_files = "modal_proto/api.proto modal_proto/options.proto"
py_protoc = (
f"{sys.executable} -m grpc_tools.protoc"
+ " --python_out=. --grpclib_python_out=. --modal-grpclib-python_out=."
+ " --grpc_python_out=. --mypy_out=. --mypy_grpc_out=."
+ f" --plugin=protoc-gen-modal-grpclib-python={grpc_plugin_path}"
protoc_cmd + " --python_out=. --grpclib_python_out=." + " --grpc_python_out=. --mypy_out=. --mypy_grpc_out=."
)
print(py_protoc)
ctx.run(f"{py_protoc} -I . " "modal_proto/api.proto " "modal_proto/options.proto ")
# generate grpcio and grpclib proto files:
ctx.run(f"{py_protoc} -I . {input_files}")

# generate modal-specific wrapper around grpclib api stub using custom plugin:
grpc_plugin_path = Path(__file__).parent / "protoc_plugin" / "plugin.py"
ctx.run(
f"{protoc_cmd} --plugin=protoc-gen-modal-grpclib-python={grpc_plugin_path}"
+ " --modal-grpclib-python_out=. -I . {input_files}"
)


@task
Expand Down
Loading