mithril-security / blindai

Confidential AI deployment with secure enclaves :lock:
https://www.mithrilsecurity.io/
Apache License 2.0
502 stars 36 forks source link

Use python context in API #40

Closed cchudant closed 2 years ago

cchudant commented 2 years ago

Motivation and Context

Current usage may leak an open socket if the user forgets to call close_connection:

client = BlindAiClient()
client.connect_server(addr="localhost", simulation=True)
# do something with client...
client.close_connection()

This is not an issue right now since our users are mostly testing the app, making jupyter notebooks and not actual production usage. As the project matures, this may become an issue.

Description

There is a way in python to make APIs that work like this:

with client = BlindAiClient.connect_server(addr="localhost", simulation=True):
  # do something with client...
  client.run_model("aaa")
  # implicitely close the connection when exiting the scope

using special __enter__ and __exit__ functions, iirc

What do you think? Is this a better API surface? This should be backward compatible with the current API.

Test plans

unit tests

Additional Information

None

Checklist

JoFrost commented 2 years ago

Good point.