Custom protoc plugin to generate a grpclib wrapper

[freider]

Basically forked grpclib's grpclib.plugin.main module and strips most of the stuff there to just reference the original api stub from grpclib (assumes same output directory).

Output looks like this:

class ModalClientModal:

    def __init__(self, grpclib_stub: modal_proto.api_grpc.ModalClientStub) -> None:
        self.AppClientDisconnect = modal._utils.grpc_utils.UnaryUnaryWrapper(grpclib_stub.AppClientDisconnect)
        self.AppCreate = modal._utils.grpc_utils.UnaryUnaryWrapper(grpclib_stub.AppCreate)
        self.AppDeploy = modal._utils.grpc_utils.UnaryUnaryWrapper(grpclib_stub.AppDeploy)
        self.AppDeploySingleObject = modal._utils.grpc_utils.UnaryUnaryWrapper(grpclib_stub.AppDeploySingleObject)
        self.AppDeploymentHistory = modal._utils.grpc_utils.UnaryUnaryWrapper(grpclib_stub.AppDeploymentHistory)
        self.AppGetByDeploymentName = modal._utils.grpc_utils.UnaryUnaryWrapper(grpclib_stub.AppGetByDeploymentName)
        self.AppGetLogs = modal._utils.grpc_utils.UnaryStreamWrapper(grpclib_stub.AppGetLogs)
        self.AppGetObjects = modal._utils.grpc_utils.UnaryUnaryWrapper(grpclib_stub.AppGetObjects)
        ...

And new functionality can be added by modifying UnaryStreamWrapper and UnaryUnaryWrapper ( see https://github.com/modal-labs/modal-client/pull/2181/files#diff-355dd58f4cc15405c65e0df2b896f1be96f94dd5b08a304ebfdf5a3ab74f9b7dR125-R160)

The advantage of doing this vs doing it dynamically is that static typing is maintainted since the cardinality of each method gets its own wrapper + type vars references the original Send and Receive types.

This patch replaces all usages of the api stub with the custom wrapper, but at this point the wrappers don't do anything except forward the request to the underlying grpclib method. In a followup we could customize the wrapper code to do things like Client-closed tracking, retries and "unary_stream as async gen" convenience methods

[erikbern]

Super cool! I have a feeling this will help us in the long run

[mwaskom]

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.

[freider]

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

[mwaskom]

whoops didn't read carefully enough!

[ekzhang]

I’m new to the thread @freider but saw this because I authored the code this PR modifies in grpc_utils (notably the unary_stream stuff) when I moved this repo to grpclib. Could we do this with interceptors instead? Just curious why we need to fork the repo to add middleware

[erikbern]

We're not forking the repo. protoc already has the notion of a "plugin" which lets you define custom code for generating stubs. I think the nice thing about doing things "statically" is that type safety comes naturally and that tracebacks are simpler. I think it will also let us do things like stub.FunctionCreate(request, retries=3) or similar, which I think will lead to cleaner code. Some of those things are probably possible to do using interceptors though!

[freider]

Yeah "forking" was perhaps a bad choice of words - I just used grpclib's proto generator plugin as a template for implementing my own plugin - ripping out most of the original code, and the new one isn't a replacement but a complement to the original generated api stub (so the original grpclib is still emitted and used).

This allows us to still use grpclib's api stub but add a type-safe wrapper for the API stub that respects the cardinality of each underlying method and lets us add cardinality-specific utility methods directly on those wrappers. The same can be solved by always making calls via wrappers like retry_transient_errors(gprc_api.SomeMethod, ...) etc but it's pretty nice to not have the "default code path" be via our wrappers.
Sadly the python typing system is rather limited when it comes to modifying classes as a whole, so I couldn't find a dynamic way of wrapping the grpclib stub class and maintain cardinality info for each method.

Couldn't find any docs about interceptors in grpclib - there is some basic "event" support, but afaict it wouldn't allow us to modify the call signature or injecting a client-specific call context, which is the purpose of this refactor (coming in followup PR).