Patch REST get_aggregated_tensor api for batched changes as per #1575

 - rename the function to get_aggregated_tensors
 - adjust both grpc and rest client for name and signature change in api
 - for the protobuf changes of TensorSpeca and Batched response adjust both client and server in line with grpc
 - fix the test cases
rebased
 29th.May.1
Signed-off-by: Shailesh Pant <[email protected]>

Author: ishaileshpant

openfl/protocols/aggregator_client_interface.py

@@ -24,13 +24,10 @@ def get_tasks(self) -> Tuple[List[Any], int, int, bool]:
         pass
 
     @abstractmethod
-    def get_aggregated_tensor(
+    def get_aggregated_tensors(
         self,
-        tensor_name: str,
-        round_number: int,
-        report: bool,
-        tags: List[str],
-        require_lossless: bool,
+        tensor_keys,
+        require_lossless: bool = True,
     ) -> Any:
         """
         Retrieves the aggregated tensor.

openfl/transport/grpc/aggregator_client.py

@@ -10,10 +10,8 @@
 
 import grpc
 
-from openfl.protocols import aggregator_pb2, aggregator_pb2_grpc, utils
-from openfl.protocols.aggregator_client_interface import AggregatorClientInterface
-from openfl.protocols import aggregator_pb2, aggregator_pb2_grpc
 from openfl.protocols import aggregator_pb2, aggregator_pb2_grpc, base_pb2, utils
+from openfl.protocols.aggregator_client_interface import AggregatorClientInterface
 from openfl.transport.grpc.common import create_header, create_insecure_channel, create_tls_channel
 
 logger = logging.getLogger(__name__)

openfl/transport/rest/aggregator_client.py

@@ -443,46 +443,75 @@ def get_tasks(self) -> Tuple[List[Any], int, int, bool]:
         )
         return tasks_resp.tasks, tasks_resp.round_number, tasks_resp.sleep_time, tasks_resp.quit
 
-    def get_aggregated_tensor(
+    def get_aggregated_tensors(
         self,
-        tensor_name: str,
-        round_number: int,
-        report: bool,
-        tags: List[str],
-        require_lossless: bool,
-    ) -> Any:
-        """Get aggregated tensor with proper security settings."""
-        params = {
-            "sender": self.collaborator_name,
-            "receiver": self.aggregator_uuid,
-            "federation_uuid": self.federation_uuid,
-            "tensor_name": tensor_name,
-            "round_number": round_number,
-            "report": report,
-            "tags": tags,
-            "require_lossless": require_lossless,
-            "collaborator_id": self.collaborator_name,
+        tensor_keys,
+        require_lossless: bool = True,
+    ) -> List[base_pb2.NamedTensor]:
+        """
+        Get aggregated tensors from the aggregator.
+
+        Args:
+            tensor_keys (list): A list of tensor keys to fetch from aggregator.
+            require_lossless (bool): Whether lossless compression is required.
+
+        Returns:
+            A list of `NamedTensor`s in the same order as requested.
+        """
+        logger.debug(f"Requesting {len(tensor_keys)} aggregated tensors")
+
+        # Build the request payload similar to gRPC implementation
+        tensor_specs = []
+        for k in tensor_keys:
+            tensor_specs.append(
+                {
+                    "tensor_name": k.tensor_name,
+                    "round_number": k.round_number,
+                    "report": k.report,
+                    "tags": k.tags,
+                    "require_lossless": require_lossless,
+                }
+            )
+
+        request_data = {
+            "header": {
+                "sender": self.collaborator_name,
+                "receiver": self.aggregator_uuid,
+                "federation_uuid": self.federation_uuid,
+                "single_col_cert_common_name": self.single_col_cert_common_name or "",
+            },
+            "tensor_specs": tensor_specs,
         }
-        headers = {"Accept": "application/json", "Sender": self.collaborator_name}
-        url = f"{self.base_url}/tensors/aggregated"
+
+        headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+            "Sender": self.collaborator_name,
+        }
+        url = f"{self.base_url}/tensors/aggregated/batch"
         extended_timeout = (30, 600)  # 30 seconds connect, 10 minutes read timeout
+
         try:
-            logger.debug(f"Requesting aggregated tensor {tensor_name} for round {round_number}")
+            logger.debug(f"Requesting batch of {len(tensor_keys)} aggregated tensors")
             response = self._make_request(
-                "GET", url, params=params, headers=headers, timeout=extended_timeout
+                "POST",
+                url,
+                data=json_format.MessageToJson(
+                    aggregator_pb2.GetAggregatedTensorsRequest(**request_data)
+                ),
+                headers=headers,
+                timeout=extended_timeout,
             )
             data = response.json()
-            resp = aggregator_pb2.GetAggregatedTensorResponse()
+            resp = aggregator_pb2.GetAggregatedTensorsResponse()
             json_format.ParseDict(data, resp, ignore_unknown_fields=True)
-            logger.debug(f"Successfully retrieved tensor {tensor_name} for round {round_number}")
-            return resp.tensor
+            logger.debug(f"Successfully retrieved {len(resp.tensors)} aggregated tensors")
+            return resp.tensors
         except requests.exceptions.HTTPError as e:
             if e.response.status_code == 404:
-                # This is expected during round 0 or when tensor hasn't been aggregated yet
-                logger.debug(
-                    f"No aggregated tensor found for {tensor_name} at round {round_number}"
-                )
-                return None
+                # This is expected during round 0 or when tensors haven't been aggregated yet
+                logger.debug("No aggregated tensors found for the requested tensor keys")
+                return []
             raise
 
     def send_local_task_results(

openfl/transport/rest/aggregator_server.py

@@ -741,61 +741,97 @@ def post_task_results():
     def _setup_tensor_route(self):
         """Set up the /tensors/aggregated endpoint."""
 
-        @self.app.route(f"/{self.api_prefix}/tensors/aggregated", methods=["GET"])
-        def get_aggregated_tensor():
-            """Endpoint for collaborators to retrieve an aggregated tensor."""
+        @self.app.route(f"/{self.api_prefix}/tensors/aggregated/batch", methods=["POST"])
+        def get_aggregated_tensors():
+            """Endpoint for collaborators to retrieve multiple aggregated tensors."""
             start_time = time.time()
 
             # Validate that this endpoint is not used in connector mode
             if self.use_connector:
-                abort(501, "GetAggregatedTensor not supported in connector mode")
+                abort(501, "GetAggregatedTensors not supported in connector mode")
 
-            # Get and validate collaborator identity
-            collaborator_id = request.args.get("collaborator_id")
-            federation_id = request.args.get("federation_uuid")
+            try:
+                # Parse the incoming JSON to a GetAggregatedTensorsRequest protobuf message
+                request_data = request.get_json()
+                if not request_data:
+                    abort(400, "Invalid JSON payload")
 
-            # Use the consolidated validation method
-            self._is_authorized(collaborator_id, federation_id)
+                tensors_request = aggregator_pb2.GetAggregatedTensorsRequest()
+                json_format.ParseDict(request_data, tensors_request, ignore_unknown_fields=True)
 
-            # Extract tensor request parameters
-            tensor_name = request.args.get("tensor_name")
-            try:
-                round_number = int(request.args.get("round_number", 0))
-            except (TypeError, ValueError):
-                abort(400, "Invalid round number")
-            report = request.args.get("report", "").lower() == "true"
-            tags = request.args.getlist("tags")
-            require_lossless = request.args.get("require_lossless", "").lower() == "true"
-
-            # Get the tensor from aggregator - direct delegation to the aggregator
-            named_tensor = self.aggregator.get_aggregated_tensor(
-                tensor_name,
-                round_number,
-                report=report,
-                tags=tuple(tags),
-                require_lossless=require_lossless,
-                requested_by=collaborator_id,
-            )
+                # Validate headers and get collaborator identity
+                collaborator_id = tensors_request.header.sender
+                federation_id = tensors_request.header.federation_uuid
 
-            # Create response header using the standardized method
-            header = create_header(
-                sender=str(self.aggregator.uuid),
-                receiver=collaborator_id,
-                federation_uuid=str(self.aggregator.federation_uuid),
-                single_col_cert_common_name=self.aggregator.single_col_cert_common_name or "",
-            )
+                # Use the consolidated validation method
+                self._is_authorized(collaborator_id, federation_id)
 
-            # Create response with empty tensor if not found
-            response_proto = aggregator_pb2.GetAggregatedTensorResponse(
-                header=header,
-                round_number=round_number,
-                tensor=named_tensor
-                if named_tensor is not None
-                else aggregator_pb2.NamedTensorProto(),
-            )
+                # Validate request header similar to gRPC implementation
+                assert tensors_request.header.receiver == str(self.aggregator.uuid), (
+                    f"Header receiver mismatch. Expected: {self.aggregator.uuid}, "
+                    f"Got: {tensors_request.header.receiver}"
+                )
 
-            logger.debug(f"Tensor retrieval completed in {time.time() - start_time:.2f} seconds")
-            return jsonify(json_format.MessageToDict(response_proto))
+                assert tensors_request.header.federation_uuid == str(
+                    self.aggregator.federation_uuid
+                ), (
+                    f"Federation UUID mismatch. Expected: {self.aggregator.federation_uuid}, "
+                    f"Got: {tensors_request.header.federation_uuid}"
+                )
+
+                expected_cn = self.aggregator.single_col_cert_common_name or ""
+                assert tensors_request.header.single_col_cert_common_name == expected_cn, (
+                    f"Single col cert CN mismatch. Expected: {expected_cn}, "
+                    f"Got: {tensors_request.header.single_col_cert_common_name}"
+                )
+
+                # Get tensors from aggregator - similar to gRPC implementation
+                logger.debug(
+                    f"Processing batch request for {len(tensors_request.tensor_specs)} tensors"
+                )
+
+                named_tensors = []
+                for ts in tensors_request.tensor_specs:
+                    named_tensor = self.aggregator.get_aggregated_tensor(
+                        ts.tensor_name,
+                        ts.round_number,
+                        ts.report,
+                        tuple(ts.tags),
+                        ts.require_lossless,
+                        collaborator_id,
+                    )
+                    # Add tensor to list (None tensors will be handled by the client)
+                    if named_tensor is not None:
+                        named_tensors.append(named_tensor)
+                    else:
+                        # Add empty tensor placeholder to maintain order
+                        named_tensors.append(aggregator_pb2.NamedTensorProto())
+
+                # Create response header using the standardized method
+                header = create_header(
+                    sender=str(self.aggregator.uuid),
+                    receiver=collaborator_id,
+                    federation_uuid=str(self.aggregator.federation_uuid),
+                    single_col_cert_common_name=self.aggregator.single_col_cert_common_name or "",
+                )
+
+                # Create response
+                response_proto = aggregator_pb2.GetAggregatedTensorsResponse(
+                    header=header, tensors=named_tensors
+                )
+
+                logger.debug(
+                    f"Batch tensor retrieval completed in {time.time() - start_time:.2f} seconds. "
+                    f"Returned {len(named_tensors)} tensors"
+                )
+                return jsonify(json_format.MessageToDict(response_proto))
+
+            except AssertionError as e:
+                logger.error(f"Header validation failed: {str(e)}")
+                abort(400, str(e))
+            except Exception as e:
+                logger.error(f"Error processing batch tensor request: {str(e)}")
+                abort(400, f"Error processing batch tensor request: {str(e)}")
 
     def _setup_relay_route(self):
         """Set up the /interop/relay endpoint."""