Skip to content

Enhance federated KMeans reporting and preprocessing - #616

Open
AndreasKtenidis wants to merge 1 commit into
masterfrom
k_means_enhancements
Open

Enhance federated KMeans reporting and preprocessing#616
AndreasKtenidis wants to merge 1 commit into
masterfrom
k_means_enhancements

Conversation

@AndreasKtenidis

Copy link
Copy Markdown
Collaborator

Enhance federated KMeans reporting and preprocessing

  • Add privacy KMeans reporting with cluster size intervals, profile suppression for small clusters, convergence warnings, and clinical limitations.
  • Add manual and elbow-based K selection, plus multi start random initialization for more stable federated fitting.
  • Add kmeans_cluster_creator preprocessing step with full, binary, and subset output modes for downstream algorithms.
  • Wire preprocessing aggregation requirements into controller/worker execution so KMeans can be created safely.
  • Replace stale raw-centers fixtures with input-only API/prod validation cases.
  • Add focused unit, API-validation, and runtime integration coverage for KMeans reporting, privacy, selection, and preprocessing.
  • Add local KMeans experiment script for comparing report and preprocessing variants

@AndreasKtenidis
AndreasKtenidis force-pushed the k_means_enhancements branch 2 times, most recently from b24bf56 to b649403 Compare June 25, 2026 07:00

@KFilippopolitis KFilippopolitis left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

  1. [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.

  1. [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}.

  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.

  1. [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.

@KFilippopolitis KFilippopolitis left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great Job!

@ThanKarab ThanKarab left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants