Integration of client & ml model id in end-to-end ml streaming flow

pymilo/streaming/interfaces.py

@@ -63,35 +63,161 @@ class ClientCommunicator(ABC):
     """
     ClientCommunicator Interface.
 
-    Each ClientCommunicator has methods to upload the local ML model, download the remote ML model and delegate attribute call to the remote server.
+    Defines the contract for client-server communication. Each implementation is responsible for:
+    - Registering and removing clients and models
+    - Uploading and downloading ML models
+    - Handling delegated attribute access
+    - Managing model allowances between clients
     """
 
     @abstractmethod
-    def upload(self, payload):
+    def register_client(self):
         """
-        Upload the given payload to the remote server.
+        Register the client in the remote server.
 
-        :param payload: request payload
-        :type payload: dict
-        :return: remote server response
+        :return: newly allocated client ID
+        :rtype: str
         """
 
     @abstractmethod
-    def download(self, payload):
+    def remove_client(self, client_id):
         """
-        Download the remote ML model to local.
+        Remove the client from the remote server.
 
-        :param payload: request payload
-        :type payload: dict
-        :return: remote server response
+        :param client_id: client ID to remove
+        :type client_id: str
+        :return: success status
+        :rtype: bool
+        """
+
+    @abstractmethod
+    def register_model(self, client_id):
+        """
+        Register an ML model for the given client.
+
+        :param client_id: client ID
+        :type client_id: str
+        :return: newly allocated model ID
+        :rtype: str
         """
 
     @abstractmethod
-    def attribute_call(self, payload):
+    def remove_model(self, client_id, model_id):
+        """
+        Remove the specified ML model for the client.
+
+        :param client_id: client ID
+        :type client_id: str
+        :param model_id: model ID
+        :type model_id: str
+        :return: success status
+        :rtype: bool
+        """
+
+    @abstractmethod
+    def get_ml_models(self, client_id):
+        """
+        Get the list of ML models for the given client.
+
+        :param client_id: client ID
+        :type client_id: str
+        :return: list of model IDs
+        :rtype: list[str]
+        """
+
+    @abstractmethod
+    def grant_access(self, allower_id, allowee_id, model_id):
+        """
+        Grant access to a model from one client to another.
+
+        :param allower_id: client who owns the model
+        :type allower_id: str
+        :param allowee_id: client to be granted access
+        :type allowee_id: str
+        :param model_id: model ID
+        :type model_id: str
+        :return: success status
+        :rtype: bool
+        """
+
+    @abstractmethod
+    def revoke_access(self, revoker_id, revokee_id, model_id):
+        """
+        Revoke model access from one client to another.
+
+        :param revoker_id: client who owns the model
+        :type revoker_id: str
+        :param revokee_id: client to be revoked
+        :type revokee_id: str
+        :param model_id: model ID
+        :type model_id: str
+        :return: success status
+        :rtype: bool
+        """
+
+    @abstractmethod
+    def get_allowance(self, allower_id):
+        """
+        Get all clients and models this client has allowed.
+
+        :param allower_id: client who granted access
+        :type allower_id: str
+        :return: dictionary mapping allowee_id to list of model_ids
+        :rtype: dict
+        """
+
+    @abstractmethod
+    def get_allowed_models(self, allower_id, allowee_id):
+        """
+        Get the list of model IDs that `allowee_id` is allowed to access from `allower_id`.
+
+        :param allower_id: model owner
+        :type allower_id: str
+        :param allowee_id: recipient
+        :type allowee_id: str
+        :return: list of allowed model IDs
+        :rtype: list[str]
+        """
+
+    @abstractmethod
+    def upload(self, client_id, model_id, model):
+        """
+        Upload the local ML model to the remote server.
+
+        :param client_id: ID of the client
+        :param model_id: ID of the model
+        :param model: serialized model content
+        :return: True if upload was successful, False otherwise
+        """
+
+    @abstractmethod
+    def download(self, client_id, model_id):
+        """
+        Download the remote ML model.
+
+        :param client_id: ID of the requesting client
+        :param model_id: ID of the model to download
+        :return: string serialized model
+        """
+
+    @abstractmethod
+    def attribute_call(self, client_id, model_id, call_payload):
         """
         Execute an attribute call on the remote server.
 
-        :param payload: request payload
-        :type payload: dict
+        :param client_id: ID of the client
+        :param model_id: ID of the model
+        :param call_payload: payload containing attribute name, args, and kwargs
+        :return: remote server response
+        """
+
+    @abstractmethod
+    def attribute_type(self, client_id, model_id, type_payload):
+        """
+        Identify the attribute type (method or field) on the remote model.
+
+        :param client_id: client ID
+        :param model_id: model ID
+        :param type_payload: payload containing targeted attribute
         :return: remote server response
         """

pymilo/streaming/param.py

@@ -10,3 +10,5 @@
 PYMILO_SERVER_NON_EXISTENT_ATTRIBUTE = "The requested attribute doesn't exist in this model."
 PYMILO_INVALID_URL = "The given URL is not valid."
 PYMILO_CLIENT_WEBSOCKET_NOT_CONNECTED = "WebSocket is not connected."
+
+REST_API_PREFIX = "/api/v1"

pymilo/streaming/pymilo_client.py

@@ -60,9 +60,7 @@ def encrypt_compress(self, body):
         :return: the compressed and encrypted version of the body payload
         """
         return self._encryptor.encrypt(
-            self._compressor.compress(
-                body
-            )
+            self._compressor.compress(body)
         )
 
     def toggle_mode(self, mode=Mode.LOCAL):
@@ -83,12 +81,8 @@ def download(self):
         :return: None
         """
         serialized_model = self._communicator.download(
-            self.encrypt_compress(
-                {
-                    "client_id": self.client_id,
-                    "ml_model_id": self.ml_model_id,
-                }
-            )
+            self.client_id,
+            self.ml_model_id
         )
         if serialized_model is None:
             print(PYMILO_CLIENT_FAILED_TO_DOWNLOAD_REMOTE_MODEL)
@@ -103,19 +97,100 @@ def upload(self):
         :return: None
         """
         succeed = self._communicator.upload(
-            self.encrypt_compress(
-                {
-                    "client_id": self.client_id,
-                    "ml_model_id": self.ml_model_id,
-                    "model": Export(self.model).to_json(),
-                }
-            )
+            self.client_id,
+            self.ml_model_id,
+            self.encrypt_compress({"model": Export(self.model).to_json()})
         )
         if succeed:
             print(PYMILO_CLIENT_LOCAL_MODEL_UPLOADED)
         else:
             print(PYMILO_CLIENT_LOCAL_MODEL_UPLOAD_FAILED)
 
+    def register(self):
+        """
+        Register client in the remote server.
+
+        :return: None
+        """
+        self.client_id = self._communicator.register_client()
+
+    def deregister(self):
+        """
+        Deregister client in the remote server.
+
+        :return: None
+        """
+        self._communicator.remove_client(self.client_id)
+        self.client_id = "0x_client_id"
+
+    def register_ml_model(self):
+        """
+        Register ML model in the remote server.
+
+        :return: None
+        """
+        self.ml_model_id = self._communicator.register_model(self.client_id)
+
+    def deregister_ml_model(self):
+        """
+        Deregister ML model in the remote server.
+
+        :return: None
+        """
+        self._communicator.remove_model(self.client_id, self.ml_model_id)
+        self.ml_model_id = "0x_ml_model_id"
+
+    def get_ml_models(self):
+        """
+        Get all registered ml models in the remote server for this client.
+
+        :return: list of ml model ids
+        """
+        return self._communicator.get_ml_models(self.client_id)
+
+    def grant_access(self, allowee_id):
+        """
+        Grant access to one of this client's models to another client.
+
+        :param allowee_id: The client ID to grant access to
+        :return: True if successful, False otherwise
+        """
+        return self._communicator.grant_access(
+            self.client_id,
+            allowee_id,
+            self.ml_model_id
+        )
+
+    def revoke_access(self, revokee_id):
+        """
+        Revoke access previously granted to another client.
+
+        :param revokee_id: The client ID to revoke access from
+        :return: True if successful, False otherwise
+        """
+        return self._communicator.revoke_access(
+            self.client_id,
+            revokee_id,
+            self.ml_model_id
+        )
+
+    def get_allowance(self):
+        """
+        Get a dictionary of all clients who have access to this client's models.
+
+        :return: Dict of allowee_id -> list of model_ids
+        """
+        return self._communicator.get_allowance(self.client_id)
+
+    def get_allowed_models(self, allower_id):
+        """
+        Get a list of models you are allowed to access from another client.
+
+        :param allower_id: The client ID who owns the models
+        :return: list of allowed model IDs
+        """
+        return self._communicator.get_allowed_models(allower_id, self.client_id)
+
     def __getattr__(self, attribute):
         """
         Overwrite the __getattr__ default function to extract requested.
@@ -133,13 +208,13 @@ def __getattr__(self, attribute):
         elif self._mode == PymiloClient.Mode.DELEGATE:
             gdst = GeneralDataStructureTransporter()
             response = self._communicator.attribute_type(
-                self.encrypt_compress(
-                    {
-                        "client_id": self.client_id,
-                        "ml_model_id": self.ml_model_id,
-                        "attribute": attribute,
-                    }
-                )
+                self.client_id,
+                self.ml_model_id,
+                self.encrypt_compress({
+                    "attribute": attribute,
+                    "client_id": self.client_id,
+                    "ml_model_id": self.ml_model_id,
+                })
             )
             if response["attribute type"] == "field":
                 return gdst.deserialize(response, "attribute value", None)
@@ -155,9 +230,9 @@ def relayer(*args, **kwargs):
                 payload["args"] = gdst.serialize(payload, "args", None)
                 payload["kwargs"] = gdst.serialize(payload, "kwargs", None)
                 result = self._communicator.attribute_call(
-                    self.encrypt_compress(
-                        payload
-                    )
+                    self.client_id,
+                    self.ml_model_id,
+                    self.encrypt_compress(payload)
                 )
                 return gdst.deserialize(result, "payload", None)
             return relayer