Add new `FlexibleClient` for added flexibility and migrate pFL and mixins to use this new class by nerdai · Pull Request #398 · VectorInstitute/FL4Health

nerdai · 2025-06-09T14:05:32Z

PR Type

Feature

Short Description

This PR adds a new FlexibleClient that is similar to BasicClient but introduces methods for training, predicting, and evaluating that has dependency injection of models and optimizers. I've moved our existing mixins stuff to work with FlexibleClient rather than BasicClient.

Now BasicClient is completely untouched.

In creating new clients:

Encourage users to subclass FlexibleClient if they only have a single-model client; otherwise, users should subclass BasicClient.

Migration plan:

After this PR lands, we can move existing clients over from BasicClient to FlexibleClient.

Added internal methods in `FlexibleClient`

Subclasses of FlexibleClient will need to implement:

predict_with_model() (this is a public method; removes predict() completely)
_train_step_with_model_and_optimizer()
_apply_backwards_on_losses_and_take_step() — called within private train step helper
_compute_preds_and_losses() — called within private train step helper
_val_step_with_model()
_transform_gradients_with_model() — allowing for model injection for transforming gradients

Tests Added

unit tests for raising warnings if any of predict, train_step and val_step are overridden.

codecov · 2025-06-09T14:09:47Z

Codecov Report

Attention: Patch coverage is 87.59690% with 16 lines in your changes missing coverage. Please review.

Project coverage is 79.85%. Comparing base (9d6788d) to head (b2f6ac4).
Report is 38 commits behind head on main.

Files with missing lines	Patch %	Lines
fl4health/clients/flexible_client.py	78.57%	15 Missing ⚠️
fl4health/mixins/personalized/ditto.py	96.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #398      +/-   ##
==========================================
+ Coverage   79.37%   79.85%   +0.48%     
==========================================
  Files         157      158       +1     
  Lines        9496     9540      +44     
==========================================
+ Hits         7537     7618      +81     
+ Misses       1959     1922      -37

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

scarere

Overall I like how the models and optimizers are now configurable, however I wonder if there is a way to do the same thing without adding a large number of internal methods? I could very well be missing something here in how all the pieces fit together.

nerdai · 2025-06-18T01:38:19Z

I think the codecov bot is acting funky.

emersodb · 2025-06-18T14:43:48Z

I think the codecov bot is acting funky.

If this doesn't stabilize by the time we get full reviews in, let's just forgo the patch coverage stuff for this case.

emersodb

Overall, I think this PR is pretty much ready to go. However, I wanted to point out two things:

I had a potentially dumb idea that I laid out in a comment in FlexibleClient
I wonder if, rather than migrating clients over to FlexibleClient, we should create implementations using FlexibleClientand keep those that use BasicClient as well, since we'll be supporting both paradigms, at least for the foreseeable future. We also want to make sure we're careful about moving more complex clients over, so that we understand what "Ditto-fying" of "MR-MTL-fying" actually means. I.e. does it actually function in a sensical way. Just a thought.

nerdai · 2025-06-20T02:54:46Z

Overall, I think this PR is pretty much ready to go. However, I wanted to point out two things:

I had a potentially dumb idea that I laid out in a comment in FlexibleClient

I wonder if, rather than migrating clients over to FlexibleClient, we should create implementations using FlexibleClientand keep those that use BasicClient as well, since we'll be supporting both paradigms, at least for the foreseeable future. We also want to make sure we're careful about moving more complex clients over, so that we understand what "Ditto-fying" of "MR-MTL-fying" actually means. I.e. does it actually function in a sensical way. Just a thought.

I think we can be good to try for now, and if it becomes a bother to maintain, then we can separate them then.
I'm okay with migrating slowly on this as we get comfortable with this. We should maybe map out some day what we need to see in order to gain more confidence with this approach -- on top of what we've already done i.e., confirming ditto_example works as expected.

scarere

Overall great PR, it's definitely clear that the scope of implementing these mixins was a lot more work than I had imagined it would be initially, a lot under the hood here. A lot of my comments are just straight questions for learning purposes as I'm new to Protocols and Mixins. I felt that there might be some opportunity to simplify the ditto mixin, and perhaps the protocols, however thats contigent that I understood the code properly, and there are likely some constraints I'm not seeing that necessitated certain choices in implementation. Also I was a decently confused by the structure of some of the tests, not sure if thats due to a knowledge gap or if here is a way to simplify them.

scarere · 2025-06-20T15:02:30Z

+        """
+        return self._val_step_with_model(self.model, input, target)
+
+    def predict_with_model(


For all the other new "with model" methods, you have redefined the original method to call the new one with the default model. Why didn't you redefine the predict method? Ie.

def predict(self, input): return self.predict_with_model(input, self.model)

I think we might have discussed this a bit, or something related in the last PR. Mostly this was to not break the public interface that had priorly been established with BasicClient. Now that we have FlexibleClient that inherits BasicClient I'm still inclined not to overload the interface.

We couldn't add model as param to BasicClient because we can't pass a default value, and thus we would have to unfortunately introduce a breaking change.

scarere · 2025-06-20T15:14:50Z

+
+        return losses, preds
+
+    def _apply_backwards_on_losses_and_take_step(


Is it necessary for this to be it's own function. It's only 3 lines of code and as far as I can tell the only place it will be called is in _train_step_with_model_and_optimizer. Do you feel there is a need here for this functionality to be separate? If not we can reduce the number of methods in this class by 1, and the number of lines by a decent amount

Without trying to get this all to work with NNunetClient I would have to agree with you. But when trying to ditto-ify it, I learned that this client specializes this logic and I need to be able to pass the respective model, optimizers and losses.

So, this should be interpreted as us introducing flexibility that will be used in the subsequent PR when we move over NNunetClient to inherit FlexibleClient

scarere · 2025-06-20T15:15:10Z

+        NOTE: Subclasses should implement this method if there is a need to specialize
+        the train_step logic.
+
+        Args:


The model and optmizer arguments are undefined in this docstring

scarere · 2025-06-20T15:18:08Z

+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:


nit: Not sure if there is a standard here, but when the return signature for a function is a fixed length tuple containing different return values, I usually seperate each value of the tuple as separate return values with the tuple type being implicit. Ie

""" Returns: losses (TrainingLosses) : The losses object ... preds (TorchPredType): dictionary of any predictions ... """

Is there a preferred way here?

This occurs in several other methods in this class, so I'll just let this comment represent all instances of this. I recognize it was probably just copied from BasicClient's docstrings but if my suggested method is preferable then maybe a good idea to start implementing this practice here?

I like this suggestion in principle, but I don't think it adheres to the "Google docstring" style. It's hard to tell how they "want" you to document returns, but by default the docstrings that are generated follow the tuple typing. The style you're describing is part of the numpy standard I think though. So I'm not 100% sure.

In any case, the rest of the library doesn't break tuples out in the way you're describing at the moment.

In this case, we should keep as is to be standard with rest of library.

scarere · 2025-06-20T15:23:13Z

+        self, model: nn.Module, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[EvaluationLosses, TorchPredType]:
+        """
+        Helper method for val_step that allows for injection of model.


Docstring missing definition for model argument

scarere · 2025-06-20T17:22:21Z

@@ -0,0 +1,286 @@
+import datetime


I'm assuming that all these tests are just copied from the tests for basic client. It seems redundant to have the same tests twice. Is there not a way to modify the basic client tests to be run a second time but with a difference MockClient? In this case using the FlexibleMockClient instead, like just parametrize the client. I'm pretty sure pytest has like a decorator or something that does this. Then the flexible client tests can be specific to the functionality that's unique to flexible client. Also now that flexible client inherits from the basic client, won't most of these tests just be covering the same code?

We could parametrize maybe -- not going to do that here tho lol.

emersodb

I left a very minor comment or two. However, from my perspective, the code is ready to go.

It would be good to address Shawn's comments before merging. 🙂

emersodb · 2025-06-20T18:09:44Z

+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:


I like this suggestion in principle, but I don't think it adheres to the "Google docstring" style. It's hard to tell how they "want" you to document returns, but by default the docstrings that are generated follow the tuple typing. The style you're describing is part of the numpy standard I think though. So I'm not 100% sure.

In any case, the rest of the library doesn't break tuples out in the way you're describing at the moment.

emersodb · 2025-06-20T18:20:31Z

        This mixin implements the Ditto algorithm from Ditto: Fair and Robust Federated Learning Through
        Personalization. This mixin inherits from the `AdaptiveDriftConstrainedMixin`, and like that mixin,
-        this should be mixed with a `BasicClient` type in order to apply the Ditto personalization method
+        this should be mixed with a `FlexibleClient` type in order to apply the Ditto personalization method


I can answer (1) at least. It's presence as something separate makes it overridable. Which was mildly useful in implementing the FENDA + Ditto approach in the library.

nerdai · 2025-06-21T16:06:43Z

Alright I believe I addressed all CRs at this point. @scarere comment's re: better organization of test suites are good and valid. I think we should handle them more broadly in the library. And, so as to not drag on this PR any longer, am choosing not to handle tests clean up here.

nerdai marked this pull request as ready for review June 9, 2025 19:15

nerdai requested review from emersodb and scarere June 9, 2025 19:42

nerdai changed the title ~~Nerdai/refactor basic client~~ Refactor BasicClient to prepare for pFL mixins Jun 9, 2025

nerdai changed the title ~~Refactor BasicClient to prepare for pFL mixins~~ Refactor BasicClient for added flexibility (and to prepare for pFL mixins) Jun 9, 2025

scarere reviewed Jun 10, 2025

View reviewed changes

Comment thread fl4health/clients/basic_client.py Outdated

Comment thread fl4health/clients/ditto_client.py

Comment thread fl4health/clients/nnunet_client.py Outdated

Comment thread fl4health/mixins/core_protocols.py Outdated

scarere marked this pull request as draft June 16, 2025 14:40

nerdai force-pushed the nerdai/refactor-basic-client branch from 193098a to 803363a Compare June 16, 2025 14:57

nerdai changed the title ~~Refactor BasicClient for added flexibility (and to prepare for pFL mixins)~~ Add new FlexibleClient for added flexibility and migrate pFL and mixins to use this Jun 16, 2025

nerdai changed the title ~~Add new FlexibleClient for added flexibility and migrate pFL and mixins to use this~~ Add new FlexibleClient for added flexibility and migrate pFL and mixins to use this new class Jun 16, 2025

nerdai marked this pull request as ready for review June 16, 2025 15:17

emersodb reviewed Jun 18, 2025

View reviewed changes

Comment thread fl4health/mixins/personalized/ditto.py Outdated

Comment thread fl4health/mixins/personalized/ditto.py Outdated

Comment thread tests/smoke_tests/test_smoke_tests.py

Comment thread fl4health/clients/flexible_client.py Outdated

nerdai force-pushed the nerdai/refactor-basic-client branch from b6fc115 to b3ee469 Compare June 19, 2025 20:49

nerdai added 14 commits June 19, 2025 22:15

private methods

56f4982

implement subclass init method for validation

5799a34

pin jupyter-core due to pip audit vulnerability

e39be07

pin jupyter-core due to pip audit vulnerability

b901844

update perfcl

cec0e81

bring in updated nnunet client

7cf1713

update basic client protocol

4d839f3

add tests for warnings being raised

bb38df5

add flexible client

25bd640

revert ditto and perfcl clients

8667c45

wip

432a84a

repeated

9702b13

rm datasets

338c5f6

poetry lock from main

102e6e6

nerdai added 10 commits June 19, 2025 22:19

add smoke tests for mnist client dynamic

767dd3b

test train step adaptive

4d8d197

more coverage

2a78759

more coverage

85eb62d

no cover to test

5dea6d7

no cover to test

fad0d39

cover again

6bf65e7

dont cover but tests do

36ac8ce

add epochs smoke test:

2bcc803

wip

2317ba2

nerdai force-pushed the nerdai/refactor-basic-client branch from 73b5abd to 2317ba2 Compare June 20, 2025 02:23

revert back

5f67019

nerdai force-pushed the nerdai/refactor-basic-client branch from 90ca7e0 to 5f67019 Compare June 20, 2025 02:26

fmt

c6e827d

nerdai force-pushed the nerdai/refactor-basic-client branch from eafdd5b to c6e827d Compare June 20, 2025 02:38

fmt

29167c5

nerdai force-pushed the nerdai/refactor-basic-client branch from 9c7c457 to 29167c5 Compare June 20, 2025 02:40

nerdai added 2 commits June 19, 2025 22:44

cr

b9124f4

cr

f0fca19

nerdai added 3 commits June 20, 2025 09:56

dry

5969f4f

bring back transform_gradients_with_model

bf3f2a9

more dry

440b491

scarere reviewed Jun 20, 2025

View reviewed changes

emersodb approved these changes Jun 20, 2025

View reviewed changes

cr

351dd67

cr

b2f6ac4

nerdai merged commit 3f50d8f into main Jun 21, 2025
10 checks passed

nerdai deleted the nerdai/refactor-basic-client branch June 21, 2025 16:46


		return losses, preds

		def _apply_backwards_on_losses_and_take_step(

Uh oh!

Conversation

nerdai commented Jun 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Type

Short Description

Added internal methods in FlexibleClient

Tests Added

Uh oh!

codecov Bot commented Jun 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

scarere left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nerdai commented Jun 18, 2025

Uh oh!

emersodb commented Jun 18, 2025

Uh oh!

emersodb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nerdai commented Jun 20, 2025

Uh oh!

scarere left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

emersodb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nerdai commented Jun 21, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

nerdai commented Jun 9, 2025 •

edited

Loading

Added internal methods in `FlexibleClient`

codecov Bot commented Jun 9, 2025 •

edited

Loading

scarere left a comment •

edited

Loading