Repository navigation
Conversation
c7cd58c to
9726da6
Compare
9726da6 to
9f2fe66
Compare
JulienVig
left a comment
JulienVig
left a comment
There was a problem hiding this comment.
I think this refactor is a good idea but I see a few blockers right now:
- The zod schema is too loose, it allows incorrect arguments without error, and doesn't let TypeScript narrow types correctly
- The CLI parsing is quite shaky, it looks like passing CLI argument as JSON would let us avoid reimplementing a possible incorrect CLI parser even if less convenient
- The CLI isn't clear about its API: what are the accepted parameters, and it doesn't fail clearly when passing incorrect arguments
Let me know what you think on keeping JSON arguments for the CLI
| const key = trimmedPair.slice(0, eqIdx).trim(); | ||
| const valStr = trimmedPair.slice(eqIdx + 1).trim(); | ||
|
|
||
| if (valStr === "true") { |
There was a problem hiding this comment.
| if (valStr === "true") { | |
| if (valStr.toLowerCase() === "true") { |
Just to avoid be more flexible
|
|
||
| if (valStr === "true") { | ||
| params[key] = true; | ||
| } else if (valStr === "false") { |
| task.trainingInformation.aggregationStrategy = | ||
| aggregationConfig.aggregationStrategy; | ||
| task.trainingInformation.aggregationParams = | ||
| aggregationConfig.aggregationParams; |
There was a problem hiding this comment.
If no aggregationParams is specified in the CLI I think you should fallback to the default task's values before falling back to zod's default values
| const aggregationConfig = aggregationConfigSchema.parse({ | ||
| aggregationStrategy: unsafeArgs.aggregator, | ||
| aggregationParams: unsafeArgs.aggregatorParams, | ||
| }); |
There was a problem hiding this comment.
if we mispell an argument or enter a wrong argument, parse will just drop it and fill with the default value. The CLI should raise an error if it doesn't recognize an argument
| aggregatorParams: { | ||
| type: parseAggregatorParams, | ||
| description: | ||
| "Momentum coefficient to smooth the aggregation over multiple rounds", | ||
| optional: true, | ||
| }, | ||
|
|
||
| // Secure aggregator | ||
| maxShareValue: { | ||
| type: Number, | ||
| description: "Maximum absolute value over all the weights", | ||
| 'Parameters for the chosen aggregator as comma-separated key=value pairs (e.g. "clippingRadius=1.0,beta=0.8")', |
There was a problem hiding this comment.
With this refactor users now have to look into training_information.ts to find what are possible arguments and their types I think a CLI should be very clear on what parameters it accepts or not
| aggregationParams: z | ||
| .record(z.string(), aggregationParamValueSchema) | ||
| .optional(), | ||
| }); |
There was a problem hiding this comment.
Setting a default loose type for aggregationParams to the baseSchema allows entering mispelled or wrong parameters that are just ignored without error. This also prevents TypeScript from narrowing the type so we can write aggregationParams?.clipingRadius without the compiler raising a type errorwk
I think you should remove the base optional aggregationParams and define it for each network type with strictObject. You could set aggregationStrategy to z.undefined() or mean for local
| // Aggregator | ||
| aggregator: AggregationStrategy; | ||
| // Byzantine aggregator | ||
| clippingRadius?: number; | ||
| maxIterations?: number; | ||
| beta?: number; | ||
| // Secure aggregator | ||
| maxShareValue?: number; | ||
| aggregatorParams?: Record<string, unknown>; | ||
| aggregationConfig: AggregationConfig; |
There was a problem hiding this comment.
aggregationConfig is deduced and validated from aggregator and aggregatorParams so I think we should only expose aggregationConfig to avoid exposing unvalidated arguments and reduce redundancy
| if (task.trainingInformation.aggregationStrategy === "byzantine") { | ||
| this.#bftClippingRadius = | ||
| task.trainingInformation.aggregationParams?.clippingRadius; | ||
| } |
There was a problem hiding this comment.
The clipping radius is initialized with a fallback value of 1.0 in the zod schema and the aggregator get.ts but here it's ignored if uninitialized. Could you unify the aggregationParams parsing that sets the default values in one place to avoid this situation?
| if isinstance(v, dict): | ||
| items.extend(flatten_params(v, prefix=key)) | ||
| else: | ||
| items.append(f"{key}={v}") |
There was a problem hiding this comment.
Python would set booleans to "True" and "False", currently the JS parser would not parse it correctly
| def flatten_params(params, prefix=""): | ||
| items = [] | ||
| for k, v in params.items(): | ||
| key = f"{prefix}.{k}" if prefix else str(k) |
There was a problem hiding this comment.
The CLI doesn't parse this dotted format and no experiment uses it
Why
Currently the local client trainer is decoupled from the aggregator, which is limiting to implementing more complex algorithms.
What
This PR unifies the configuration of aggregation strategies across
discojsand the CLI.Previously, aggregator parameters were scattered across ad-hoc fields like
privacy.byzantineFaultToleranceand rootmaxShareValue, causing CLI argument sprawl.This refactor introduces a typed
aggregationParamsschema discriminated byaggregationStrategyand a generic--aggregatorParamsCLI flag.Individual flags (
--clippingRadius,--beta,--maxIterations,--maxShareValue) are removed in favor of comma-separated key-value pairs, keeping argument parsing clean and extensible.TODO