More resilient & long-lasting connections

[divymonads]

• ml-metadata==1.14.0
• mysql
• running using direct db connection & minio

I've seen this issue in my environment where I'm using a long lasting MetadataStore database connection. I'm using mysql database and connecting without the grpc server.

This connection disconnects after some idle timeout frequently. I've already checked configurations across the stack, likely something in my k8s runtime is killing the connection. I can see an issue with very simple code such as this. The second call to the client errors out with an Internal Error. Works fine if I re-instantiate it.

if __name__ == "__main__":
    mlmd_client = get_mlmetadata_client() # returns instantiated MetadataStore object
    context_types = mlmd_client.get_context_types_by_external_ids(["Test"])
    print(context_types)
    time.sleep(1800)
    context_types = mlmd_client.get_context_types_by_external_ids(["Test"])
    print(context_types)

For my use-case, I've implemented a simple retry in my code to workaround the issue by overriding the private _call method in the class. I felt this was most optimal for my use-case.

I think it would be a good feature to implement something that tries re-establishes the connection. If needed to be done in a later version I think that makes sense. Would be good to know if there is more stability with a workflow that's different from the way I'm doing it.