Add CollaboratorSerialiser middleware for serialisation/deserialisation between collaborator and aggregator client

[theakshaypant]

Summary

Moving serialisation/deserialisation logic out of collaborator into a separate middleware.

Why?

By clearly separating concerns across layers, it aims to enhance code readability for contributors and improves debuggability in case of failures.
This middleware introduction is similar to "middleware chaining" paradigm followed widely in Go where the same function is implemented across multiple "layers", each handling a distinct aspect of the functionality. In the context of this PR, this is done such that

• Layer 1 (Collaborator): Performs operations central to FL.
• Layer 2 (CollaboratorSerialiser): Performs serialisation/desrialisation of numpy array.
• Layer 3 (Transport): Interacts with the aggregator server to send/receive data.

Type of Change (Mandatory)

• Feature enhancement
• Restructuring

Description (Mandatory)

• The collaborator uses TensorCodec to convert tensors to protobuf format and vice versa in order to communicate with the aggregator.
• This conversion logic is moved out from the collaborator and it focuses solely on the core functionality while the conversion and client call is abstracted from the user.
• With the introduction of other transport implementations in the future, the functionality can be retained by passing the new (say REST) client to the middleware.

Testing

Manual and CI.
Eden compression CI

Additional Information

• Will be adding a serialisation middleware to the aggregator in a subsequent PR.

[theakshaypant]

The changes in tests mostly include moving tests between the now appropriate modules and/or renaming them.

[theakshaypant]

This method has been moved to Collaborator after removing these lines

                if self.compression_pipeline.is_lossy():
                    new_tags = ("aggregated", "delta", "lossy_compressed")
                else:
                    new_tags = ("aggregated", "delta", "compressed")
                tensor_key_dependencies.append(
                    TensorKey(tensor_name, origin, round_number, report, new_tags)
                )

As can be seen from the usage of the returned dependencies in the Collaborator.get_data_for_tensorkey, only the first tensor dependency is used, which is always

TensorKey(tensor_name, origin, round_number - 1, report, tags)

Ref.

[psfoley]

I don't think that's quite right. When use_delta_updates is set to true, there are two dependencies: the prior version of the model, and the delta that should be applied (either compressed with a lossless or lossy pipeline) to that model to bring it to the current version. The 2nd tensor dependency (tensor_dependency[1]) is used later in the function that you reference here. If you're only seeing the first tensor dependency for every run, then what's likely happening is that db_store_rounds for the collaborator is set to 1 (it needs to be set to 2 in order to retain the prior version of the model locally). Before these lines are removed, you should add tests for use_delta_updates=True and use a lossy compression pipeline (such as this one).

[theakshaypant]

We can also consider removing client completely from the collaborator since the middleware already has the client object.
The client object is currently used by the secure aggregation callback which could potentially also work with the middleware given the changes proposed in #1518

[teoparvanov]

On the contrary, I'm wondering whether the collaborator shouldn't just have a client object, while any "middlewares" are added as decorators to the AggregatorGRPCClient class. Those decorators could receive certain configurations via the FL plan (such as the compression algo).

WDYT?

[theakshaypant]

Ah, I see your point, I hadn't considered using decorators here!
I went with "middleware" approach as I feel it is easier to "chain" - multiple layers can be used and each of them can be easily added/removed without affecting any of the other layers.
With the middleware approach, the only change required to switch between gRPC and REST would be in how the client is initialized. Similarly, in the future if we decide to make changes to the compression pipeline or such, we can look to make changes only the serialisation layer.

Gotta admit, I’m a bit biased toward this approach 😁 but totally open to other perspectives!

[MasterSkepticista]

I am leaning towards having entire approach encapsulated within the client (@teoparvanov's approach is one of the ways to do that).

The current implementation removes compression pipeline from the collaborator and it adds a new serializer object. From an abstraction standpoint it would be worth asking: does a collaborator need to know about serialization? Having it all in the client marks that separation very well.

I see the benefit in your approach of having a middleware when a different use case leverages it. Could you list any planned features that may benefit from this approach?

[theakshaypant]

My original point was that, once we have sorted out the secure aggregation serialisation, we can fully replace the client object with this proposed class object.
A key argument for keeping this logic separate from the client is the potential addition of a REST client. With the approach proposed in this PR, the serialiser would simply reference the enabled transport client, REST or gRPC.
An argument can be made for going with the decorator approach as mentioned by Teo but in my opinion, the middleware approach provides us more flexibility.

[theakshaypant]

These tests are commented out and planned to be removed, refer #1476 (comment)

[teoparvanov]

Hi @theakshaypant , thanks for starting this effort.

To begin with, could you add a "why" section in the PR description? In addition to the obvious encapsulation benefits, what framework capabilities do you expect this refactoring to enable down the road?

Additional comments based on the code itself:

[teoparvanov]

On the contrary, I'm wondering whether the collaborator shouldn't just have a client object, while any "middlewares" are added as decorators to the AggregatorGRPCClient class. Those decorators could receive certain configurations via the FL plan (such as the compression algo).

WDYT?

[teoparvanov]

This class indeed looks like a decorator of AggregatorGRPCClient, as it implements the same interface, but adds some orthogonal aspects like encoding/decoding via the tensor codec.

[theakshaypant]

To begin with, could you add a "why" section in the PR description? In addition to the obvious encapsulation benefits, what framework capabilities do you expect this refactoring to enable down the road?

@teoparvanov Thanks for the comments! Updated the PR description.
Adding a bit more descriptive answer here. This PR doesn't directly support any upcoming features/capabilities. However, by clearly separating concerns across layers, it enhances code readability for contributors and improves debuggability in case of failures. For instance, future developers won’t need to worry about low-level details like converting NumPy arrays to named tensors, as such operations are abstracted away and don’t impact the core functionality of the component.
This is a small but hopefully meaningful step toward the larger improvements we’re aiming for in the repo.

[psfoley]

I haven't gotten to the point of reviewing the architectural changes proposed, but I see some changes that appear to break the delta update functionality. I see some positive steps taken overall - reducing code within the collaborator is always a good thing. 2nd half of the review to come tomorrow.

[psfoley]

I don't think that's quite right. When use_delta_updates is set to true, there are two dependencies: the prior version of the model, and the delta that should be applied (either compressed with a lossless or lossy pipeline) to that model to bring it to the current version. The 2nd tensor dependency (tensor_dependency[1]) is used later in the function that you reference here. If you're only seeing the first tensor dependency for every run, then what's likely happening is that db_store_rounds for the collaborator is set to 1 (it needs to be set to 2 in order to retain the prior version of the model locally). Before these lines are removed, you should add tests for use_delta_updates=True and use a lossy compression pipeline (such as this one).

[MasterSkepticista]

While I understand the motivation behind adding a new component, this adds significant LOC for what is effectively adding deserialize/serialize calls into each RPC.

IMO let us migrate those calls to grpc client to land a minimum API change version.

When we realize the need for a serializer to do more than what is being done today, I clearly see value in having a dedicated component that forwards across RPC methods like we do here.

[MasterSkepticista]

I am leaning towards having entire approach encapsulated within the client (@teoparvanov's approach is one of the ways to do that).

The current implementation removes compression pipeline from the collaborator and it adds a new serializer object. From an abstraction standpoint it would be worth asking: does a collaborator need to know about serialization? Having it all in the client marks that separation very well.

I see the benefit in your approach of having a middleware when a different use case leverages it. Could you list any planned features that may benefit from this approach?