Skip to content

Commit 11331a8

Browse files
authored
Integration of client & ml model id in end-to-end ml streaming flow (#196)
* `/register/` endpoint to register a new client * `/request_model_id/` endpoint to request a new model id * `/clients/` endpoint to get all clients * `/client/models/` endpoint to request a specific model within a specific client * integration with previous 4 endpoints in `RESTClientCommunicator` * update previous endpoints to apply client/model id check/filter * pass download information to `PyMiloServer` in `WebSocketServerCommunicator` * `uuid` used to generate universal unique ids * use `HTTPException` to handle erros * integrate id creation/allocation for model and client in PyMiloClient * integrate id creation/allocation for model and client in PyMiloServer * develop id creation/allocation functions for ml model and client in PyMiloServer * add client_id and model_id to scenario3 * add init client/model in PyMiloServer setup * add client/model initiation * fulfill docstring * run `autopep8.sh` * remove trailing whitespaces * update naming * restructure and massive update on `REST[Server|Client]Communicator` * update `ClientCommunicator` interface to support new features * update `PymiloClient` to make it compatible with latest updates * update `PyMiloServer` to make it compatible with latest updates + add `allowance` * add api version param * comment `websocket` test, to be fixed in the next PR. * update interface * remove trailing whitespaces * remove trailing whitespaces
1 parent 77c36e0 commit 11331a8

File tree

10 files changed

+753
-143
lines changed

10 files changed

+753
-143
lines changed

pymilo/streaming/communicator.py

Lines changed: 258 additions & 80 deletions
Large diffs are not rendered by default.

pymilo/streaming/interfaces.py

Lines changed: 140 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -63,35 +63,161 @@ class ClientCommunicator(ABC):
6363
"""
6464
ClientCommunicator Interface.
6565
66-
Each ClientCommunicator has methods to upload the local ML model, download the remote ML model and delegate attribute call to the remote server.
66+
Defines the contract for client-server communication. Each implementation is responsible for:
67+
- Registering and removing clients and models
68+
- Uploading and downloading ML models
69+
- Handling delegated attribute access
70+
- Managing model allowances between clients
6771
"""
6872

6973
@abstractmethod
70-
def upload(self, payload):
74+
def register_client(self):
7175
"""
72-
Upload the given payload to the remote server.
76+
Register the client in the remote server.
7377
74-
:param payload: request payload
75-
:type payload: dict
76-
:return: remote server response
78+
:return: newly allocated client ID
79+
:rtype: str
7780
"""
7881

7982
@abstractmethod
80-
def download(self, payload):
83+
def remove_client(self, client_id):
8184
"""
82-
Download the remote ML model to local.
85+
Remove the client from the remote server.
8386
84-
:param payload: request payload
85-
:type payload: dict
86-
:return: remote server response
87+
:param client_id: client ID to remove
88+
:type client_id: str
89+
:return: success status
90+
:rtype: bool
91+
"""
92+
93+
@abstractmethod
94+
def register_model(self, client_id):
95+
"""
96+
Register an ML model for the given client.
97+
98+
:param client_id: client ID
99+
:type client_id: str
100+
:return: newly allocated model ID
101+
:rtype: str
87102
"""
88103

89104
@abstractmethod
90-
def attribute_call(self, payload):
105+
def remove_model(self, client_id, model_id):
106+
"""
107+
Remove the specified ML model for the client.
108+
109+
:param client_id: client ID
110+
:type client_id: str
111+
:param model_id: model ID
112+
:type model_id: str
113+
:return: success status
114+
:rtype: bool
115+
"""
116+
117+
@abstractmethod
118+
def get_ml_models(self, client_id):
119+
"""
120+
Get the list of ML models for the given client.
121+
122+
:param client_id: client ID
123+
:type client_id: str
124+
:return: list of model IDs
125+
:rtype: list[str]
126+
"""
127+
128+
@abstractmethod
129+
def grant_access(self, allower_id, allowee_id, model_id):
130+
"""
131+
Grant access to a model from one client to another.
132+
133+
:param allower_id: client who owns the model
134+
:type allower_id: str
135+
:param allowee_id: client to be granted access
136+
:type allowee_id: str
137+
:param model_id: model ID
138+
:type model_id: str
139+
:return: success status
140+
:rtype: bool
141+
"""
142+
143+
@abstractmethod
144+
def revoke_access(self, revoker_id, revokee_id, model_id):
145+
"""
146+
Revoke model access from one client to another.
147+
148+
:param revoker_id: client who owns the model
149+
:type revoker_id: str
150+
:param revokee_id: client to be revoked
151+
:type revokee_id: str
152+
:param model_id: model ID
153+
:type model_id: str
154+
:return: success status
155+
:rtype: bool
156+
"""
157+
158+
@abstractmethod
159+
def get_allowance(self, allower_id):
160+
"""
161+
Get all clients and models this client has allowed.
162+
163+
:param allower_id: client who granted access
164+
:type allower_id: str
165+
:return: dictionary mapping allowee_id to list of model_ids
166+
:rtype: dict
167+
"""
168+
169+
@abstractmethod
170+
def get_allowed_models(self, allower_id, allowee_id):
171+
"""
172+
Get the list of model IDs that `allowee_id` is allowed to access from `allower_id`.
173+
174+
:param allower_id: model owner
175+
:type allower_id: str
176+
:param allowee_id: recipient
177+
:type allowee_id: str
178+
:return: list of allowed model IDs
179+
:rtype: list[str]
180+
"""
181+
182+
@abstractmethod
183+
def upload(self, client_id, model_id, model):
184+
"""
185+
Upload the local ML model to the remote server.
186+
187+
:param client_id: ID of the client
188+
:param model_id: ID of the model
189+
:param model: serialized model content
190+
:return: True if upload was successful, False otherwise
191+
"""
192+
193+
@abstractmethod
194+
def download(self, client_id, model_id):
195+
"""
196+
Download the remote ML model.
197+
198+
:param client_id: ID of the requesting client
199+
:param model_id: ID of the model to download
200+
:return: string serialized model
201+
"""
202+
203+
@abstractmethod
204+
def attribute_call(self, client_id, model_id, call_payload):
91205
"""
92206
Execute an attribute call on the remote server.
93207
94-
:param payload: request payload
95-
:type payload: dict
208+
:param client_id: ID of the client
209+
:param model_id: ID of the model
210+
:param call_payload: payload containing attribute name, args, and kwargs
211+
:return: remote server response
212+
"""
213+
214+
@abstractmethod
215+
def attribute_type(self, client_id, model_id, type_payload):
216+
"""
217+
Identify the attribute type (method or field) on the remote model.
218+
219+
:param client_id: client ID
220+
:param model_id: model ID
221+
:param type_payload: payload containing targeted attribute
96222
:return: remote server response
97223
"""

pymilo/streaming/param.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,5 @@
1010
PYMILO_SERVER_NON_EXISTENT_ATTRIBUTE = "The requested attribute doesn't exist in this model."
1111
PYMILO_INVALID_URL = "The given URL is not valid."
1212
PYMILO_CLIENT_WEBSOCKET_NOT_CONNECTED = "WebSocket is not connected."
13+
14+
REST_API_PREFIX = "/api/v1"

pymilo/streaming/pymilo_client.py

Lines changed: 101 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -60,9 +60,7 @@ def encrypt_compress(self, body):
6060
:return: the compressed and encrypted version of the body payload
6161
"""
6262
return self._encryptor.encrypt(
63-
self._compressor.compress(
64-
body
65-
)
63+
self._compressor.compress(body)
6664
)
6765

6866
def toggle_mode(self, mode=Mode.LOCAL):
@@ -83,12 +81,8 @@ def download(self):
8381
:return: None
8482
"""
8583
serialized_model = self._communicator.download(
86-
self.encrypt_compress(
87-
{
88-
"client_id": self.client_id,
89-
"ml_model_id": self.ml_model_id,
90-
}
91-
)
84+
self.client_id,
85+
self.ml_model_id
9286
)
9387
if serialized_model is None:
9488
print(PYMILO_CLIENT_FAILED_TO_DOWNLOAD_REMOTE_MODEL)
@@ -103,19 +97,100 @@ def upload(self):
10397
:return: None
10498
"""
10599
succeed = self._communicator.upload(
106-
self.encrypt_compress(
107-
{
108-
"client_id": self.client_id,
109-
"ml_model_id": self.ml_model_id,
110-
"model": Export(self.model).to_json(),
111-
}
112-
)
100+
self.client_id,
101+
self.ml_model_id,
102+
self.encrypt_compress({"model": Export(self.model).to_json()})
113103
)
114104
if succeed:
115105
print(PYMILO_CLIENT_LOCAL_MODEL_UPLOADED)
116106
else:
117107
print(PYMILO_CLIENT_LOCAL_MODEL_UPLOAD_FAILED)
118108

109+
def register(self):
110+
"""
111+
Register client in the remote server.
112+
113+
:return: None
114+
"""
115+
self.client_id = self._communicator.register_client()
116+
117+
def deregister(self):
118+
"""
119+
Deregister client in the remote server.
120+
121+
:return: None
122+
"""
123+
self._communicator.remove_client(self.client_id)
124+
self.client_id = "0x_client_id"
125+
126+
def register_ml_model(self):
127+
"""
128+
Register ML model in the remote server.
129+
130+
:return: None
131+
"""
132+
self.ml_model_id = self._communicator.register_model(self.client_id)
133+
134+
def deregister_ml_model(self):
135+
"""
136+
Deregister ML model in the remote server.
137+
138+
:return: None
139+
"""
140+
self._communicator.remove_model(self.client_id, self.ml_model_id)
141+
self.ml_model_id = "0x_ml_model_id"
142+
143+
def get_ml_models(self):
144+
"""
145+
Get all registered ml models in the remote server for this client.
146+
147+
:return: list of ml model ids
148+
"""
149+
return self._communicator.get_ml_models(self.client_id)
150+
151+
def grant_access(self, allowee_id):
152+
"""
153+
Grant access to one of this client's models to another client.
154+
155+
:param allowee_id: The client ID to grant access to
156+
:return: True if successful, False otherwise
157+
"""
158+
return self._communicator.grant_access(
159+
self.client_id,
160+
allowee_id,
161+
self.ml_model_id
162+
)
163+
164+
def revoke_access(self, revokee_id):
165+
"""
166+
Revoke access previously granted to another client.
167+
168+
:param revokee_id: The client ID to revoke access from
169+
:return: True if successful, False otherwise
170+
"""
171+
return self._communicator.revoke_access(
172+
self.client_id,
173+
revokee_id,
174+
self.ml_model_id
175+
)
176+
177+
def get_allowance(self):
178+
"""
179+
Get a dictionary of all clients who have access to this client's models.
180+
181+
:return: Dict of allowee_id -> list of model_ids
182+
"""
183+
return self._communicator.get_allowance(self.client_id)
184+
185+
def get_allowed_models(self, allower_id):
186+
"""
187+
Get a list of models you are allowed to access from another client.
188+
189+
:param allower_id: The client ID who owns the models
190+
:return: list of allowed model IDs
191+
"""
192+
return self._communicator.get_allowed_models(allower_id, self.client_id)
193+
119194
def __getattr__(self, attribute):
120195
"""
121196
Overwrite the __getattr__ default function to extract requested.
@@ -133,13 +208,13 @@ def __getattr__(self, attribute):
133208
elif self._mode == PymiloClient.Mode.DELEGATE:
134209
gdst = GeneralDataStructureTransporter()
135210
response = self._communicator.attribute_type(
136-
self.encrypt_compress(
137-
{
138-
"client_id": self.client_id,
139-
"ml_model_id": self.ml_model_id,
140-
"attribute": attribute,
141-
}
142-
)
211+
self.client_id,
212+
self.ml_model_id,
213+
self.encrypt_compress({
214+
"attribute": attribute,
215+
"client_id": self.client_id,
216+
"ml_model_id": self.ml_model_id,
217+
})
143218
)
144219
if response["attribute type"] == "field":
145220
return gdst.deserialize(response, "attribute value", None)
@@ -155,9 +230,9 @@ def relayer(*args, **kwargs):
155230
payload["args"] = gdst.serialize(payload, "args", None)
156231
payload["kwargs"] = gdst.serialize(payload, "kwargs", None)
157232
result = self._communicator.attribute_call(
158-
self.encrypt_compress(
159-
payload
160-
)
233+
self.client_id,
234+
self.ml_model_id,
235+
self.encrypt_compress(payload)
161236
)
162237
return gdst.deserialize(result, "payload", None)
163238
return relayer

0 commit comments

Comments
 (0)