Enhance federated KMeans reporting and preprocessing - #616
Enhance federated KMeans reporting and preprocessing#616AndreasKtenidis wants to merge 1 commit into
Conversation
b24bf56 to
b649403
Compare
KFilippopolitis
left a comment
There was a problem hiding this comment.
Review recommendation: Request changes.
The PR is a strong expansion of KMeans reporting and preprocessing, but I found a few correctness and integration issues that should be fixed before merge.
- [P1] Recompute cluster counts after final label assignment
In exaflow/algorithms/federated/cluster/kmeans.py, final labels are recomputed after the last center update, but count_global still reflects the previous iteration’s assignments. This means labels_ can disagree with cluster_counts_, empty_clusters_, privacy checks, and size intervals.
This is particularly risky for kmeans_cluster_creator: it emits categories from the final labels but validates privacy using stale counts. A cluster that falls below the privacy threshold after the final assignment could still be exposed as if it were large enough.
Suggested fix: after recomputing final labels, recompute final per-cluster counts through aggregation and use those final counts for cluster_counts_, empty_clusters_, privacy validation, and reporting.
Suggested regression test:
assert model.cluster_counts_ == np.bincount(
model.labels_,
minlength=model.n_clusters,
).tolist()A maxiter=1 or non-converged case would be useful to catch this.
- [P1 / High] Elbow mode returns incorrect categorical enumerations
In exaflow/algorithms/exareme3/preprocessing/kmeans_cluster_creator.py, runtime metadata for full output mode still enumerates cluster_0 through cluster_{k_max-1} when k_selection=elbow, even if the fitted model selected a smaller k.
Downstream algorithms that consume category lists from metadata, such as chi-squared and categorical Naive Bayes, can then treat non-existent clusters as real levels with zero counts. This can produce incorrect degrees of freedom, contingency tables, and model statistics, even though the row labels only use the selected clusters.
Suggested fix: validation-time metadata may need to use k_max as an upper bound, but the metadata returned from transform_data_and_metadata() should be based on the fitted model.n_clusters.
Suggested regression test: create an elbow-mode preprocessing case where selected_k < k_max, then assert runtime metadata only includes cluster_0..cluster_{selected_k-1}.
- [P2] Avoid passing preprocessing-only aggregation clients to plain UDFs
In exaflow/worker/exareme3/udf/udf_service.py, an aggregation client is created when either the target UDF requires aggregation or any preprocessing step requires aggregation. Later, _execute_udf unconditionally injects agg_client into kw_args whenever it is non-null.
This breaks the intended use case of kmeans_cluster_creator before a non-aggregation algorithm: preprocessing needs the aggregation server, but the downstream algorithm UDF may not accept an agg_client argument. In that case the UDF can fail with:
TypeError: got an unexpected keyword argument 'agg_client'
Suggested fix: only pass agg_client to the UDF when the UDF signature accepts it or when the registry says the UDF itself requires the aggregation server. The aggregation client can still be passed to preprocessing steps that need it.
Suggested regression test: run an aggregation-requiring preprocessing step before a plain Exareme3 UDF whose signature does not include agg_client, and assert the UDF still executes.
- [P2] Copy algorithm components before extending them
In exaflow/controller/services/algorithm_execution_strategy_factory.py, get_component_types() returns the mutable components list stored on the enabled algorithm specification. Extending that list with preprocessing components mutates global spec state for that algorithm.
After one request combines a plain algorithm with an aggregation-server preprocessing step, later requests for the same algorithm without preprocessing will still see AGGREGATION_SERVER in the algorithm components and can be routed through Exareme3WithAggregationServerStrategy unexpectedly.
Suggested fix:
components = list(specifications.get_component_types(algorithm_name))
components.extend(_get_preprocessing_component_types(analysis_request_dto))Optionally deduplicate the list before choosing the strategy.
Suggested regression test: call the strategy factory once with aggregation-requiring preprocessing, then call it again for the same algorithm without preprocessing and assert the second request still uses the plain Exareme3 strategy.
b649403 to
e763b0f
Compare
ThanKarab
left a comment
There was a problem hiding this comment.
GPT 5.5 review comments:
The new preprocessing step can accept non-numeric text variables and fail at worker runtime
rather than during request validation. This is a functional validation gap in the added code.
Review comment:
- [P2] Validate KMeans cluster variables by SQL type — /home/thanasis/exaflow/exaflow/
algorithms/exareme3/preprocessing/kmeans_cluster_creator.py:273-277
When cluster_variables contains a text CDE that is not marked categorical, this validation
passes even though the preprocessing step promises numerical inputs. The worker then
reaches FederatedKMeans.fit(), where np.isfinite on an object/string array raises a raw
TypeError instead of a user-facing validation error. Please also reject non-numeric
sql_type values here, not only categorical variables.
Other than that, a lot of system changes were added. Will pause this for now, until the changes are implemented in a separate PR and then continue this line of work.
Enhance federated KMeans reporting and preprocessing