mypy-protobuf 3.7.0 causes enormous mypy processing time

[subnix]

After upgrading mypy-protobuf to version 3.7.0, we observed enormous memory consumption and processing time by mypy. Upon further analysis, we discovered that this issue is primarily caused by a _pb2_grpc.pyi file. It appears that each RPC method within a single service contributes to the increased processing time, and this becomes noticeable when there are 11 RPC methods.

Here's a simple example to illustrate this: mypy takes more than 5 minutes to process stubs of this proto file.

test.proto:

syntax = "proto3";

package test;

message Simple {
  string a_string = 1;
}

service Test {
  rpc Test1(Simple) returns (Simple);
  rpc Test2(Simple) returns (Simple);
  rpc Test3(Simple) returns (Simple);
  rpc Test4(Simple) returns (Simple);
  rpc Test5(Simple) returns (Simple);
  rpc Test6(Simple) returns (Simple);
  rpc Test7(Simple) returns (Simple);
  rpc Test8(Simple) returns (Simple);
  rpc Test9(Simple) returns (Simple);
  rpc Test10(Simple) returns (Simple);
  rpc Test11(Simple) returns (Simple);
  rpc Test12(Simple) returns (Simple);
  rpc Test13(Simple) returns (Simple);
  //rpc Test14(Simple) returns (Simple);
  //rpc Test15(Simple) returns (Simple);
}

$ pip install mypy 'mypy-protobuf==3.7.0' grpcio-tools types-grpcio
$ python3 -m grpc_tools.protoc --python_out=test --mypy_out=test --grpc_python_out=test --mypy_grpc_out=test --proto_path=test test/test.proto
$ mypy -v --no-incremental test

Apparently, this issue arises because of the introduction of generics in 3.7.0. Some combination of generics and typevars with constraints significantly impacts the processing of __init__ methods. Therefore, removing __init__ methods restores normal processing.

[aidandj]

What version of mypy

[subnix]

@aidandj The latest one - 1.19.0. I also tried 1.18.2 and 1.17.1, but it didn't help.

[ilevkivskyi]

If this is dues to python/mypy#19622, then there is not much can be done from mypy side: I am quite sure any correct algorithm has exponential complexity for this.

If this is something else, you can open an issue on mypy tracker (with the generated stubs), and we can check if anything can be done performance-wise from mypy side.

[aidandj]

I'll work on reproducing and see what I can find.

Fwiw this isn't a problem on ty (which may or may not even be fully resolving these overloads, I haven't checked) or pyright (which is where I did most of the development on this feature).

[subnix]

I've identified another problem and created a pull request to mypy (python/mypy#20378). This change should enhance the processing of overloads. However, due to the issue mentioned (python/mypy#19622), we may still encounter performance issues if a service has numerous RPC methods.

[aidandj]

@WouldYouKindly does this fix your issue?

[subnix]

Compiled mypy builds with the fix are available here: https://github.com/mypyc/mypy_mypyc-wheels/releases/tag/v1.20.0%2Bdev.88cf8e08f3e04f797d216869df0f8d5e030aa07c

Now, mypy takes 2 seconds to process the file with 13 RPCs, but for 15 RPCs, it takes about 30 seconds. Even if we optimize something else, it will never be enough to process dozens of constrained parameters.

[aidandj]

Ok, I might have another solution. But I may need a 2nd set of eyes on it.

Instead of overriding init and using generics, I can override __new__ and return different stub classes.

This passes all our existing tests, and I can run 100+ methods per service with no real change in performance.

class _DummyServiceStub(DummyServiceStub):
    def __init__(self, channel: grpc.Channel) -> None: ...
    UnaryUnary: grpc.UnaryUnaryMultiCallable[testproto.grpc.dummy_pb2.DummyRequest, testproto.grpc.dummy_pb2.DummyReply]
    UnaryStream: grpc.UnaryStreamMultiCallable[testproto.grpc.dummy_pb2.DummyRequest, testproto.grpc.dummy_pb2.DummyReply]
    StreamUnary: grpc.StreamUnaryMultiCallable[testproto.grpc.dummy_pb2.DummyRequest, testproto.grpc.dummy_pb2.DummyReply]
    StreamStream: grpc.StreamStreamMultiCallable[testproto.grpc.dummy_pb2.DummyRequest, testproto.grpc.dummy_pb2.DummyReply]

class _DummyServiceStubAsync(DummyServiceStub):
    def __init__(self, channel: grpc.aio.Channel) -> None: ...
    UnaryUnary: grpc.aio.UnaryUnaryMultiCallable[testproto.grpc.dummy_pb2.DummyRequest, testproto.grpc.dummy_pb2.DummyReply]
    UnaryStream: grpc.aio.UnaryStreamMultiCallable[testproto.grpc.dummy_pb2.DummyRequest, testproto.grpc.dummy_pb2.DummyReply]
    StreamUnary: grpc.aio.StreamUnaryMultiCallable[testproto.grpc.dummy_pb2.DummyRequest, testproto.grpc.dummy_pb2.DummyReply]
    StreamStream: grpc.aio.StreamStreamMultiCallable[testproto.grpc.dummy_pb2.DummyRequest, testproto.grpc.dummy_pb2.DummyReply]

class DummyServiceStub:
    """DummyService"""

    @typing.overload
    def __new__(cls, channel: grpc.Channel) -> _DummyServiceStub: ...
    @typing.overload
    def __new__(cls, channel: grpc.aio.Channel) -> _DummyServiceStubAsync: ...

There is 1 downside. Which is where I could use some input.

The __init__ method is no longer typed, and stubtest complains about it as well. I tried adding __init__ to the main class with a Union[grpc.aio.Channel, grpc.Channel] but that started down a rabbit hole of mypy not being able to type narrow on __new__. I'm still spending some time on that, but i'm not finding any clear answers. Pyright seemed to work fine and narrow things.

I'm going to keep exploring the mypy errors, but I think this might work better.

[aidandj]

#710

@subnix can you test this?

[ilevkivskyi]

Hm, btw if the problem is specific to overloads, there is a chance python/mypy#18159 may fix this. (That PR is however quite old, and not yet ready to be merged.)

[aidandj]

The time is spent almost exclusively in here. Because of the multiple generics and the cross product check the number of calls is astronomical.

I'm not sure a reduction from O(N^2V^2K) to O(N^2V^K) is actually enough, as when I started testing with 40-50 rpc's it probably needs to somehow be constant time instead. I would need to look at how pyright does it, unless someone else has looked at that already.

There are only 2 overloads per class, and each type var only has 2 bounds. But each typevar is actually another generic, so that may play into it as well, I have not fully groked the algorithm at this point.