feat: Support aggregation client in preprocessing steps

documentation/algorithms-user-guide.md

@@ -363,6 +363,62 @@ Treat these helpers as building blocks: copy the patterns from
 `exaflow/algorithms/exareme3/logistic_regression.py`,
 `linear_regression.py`, or `describe.py` when you need similar logic.
 
+### Aggregation-aware preprocessing steps
+
+Preprocessing steps that need the aggregation server must declare it in their
+specification:
+
+```python
+components=[specs.ComponentType.AGGREGATION_SERVER]
+```
+
+The runtime then configures the aggregation server for requests that include the
+step and always passes the aggregation client to the base
+`transform_data_and_metadata()` entrypoint. Aggregation-aware preprocessing
+steps should not override `transform_data_and_metadata()`; put the explicit
+client parameter on `transform_data()` instead. The default parameter name is
+`agg_client`.
+
+```python
+from exaflow.algorithms import specifications as specs
+from exaflow.algorithms.exareme3.utils.preprocessing_step import PreprocessingStep
+
+
+class GlobalPreprocessingStep(PreprocessingStep):
+  @classmethod
+  def get_specification(cls):
+    return specs.PreprocessingStepSpecification(
+      name="global_preprocessing_step",
+      desc="Example aggregation-aware preprocessing step.",
+      documentation="Example aggregation-aware preprocessing step.",
+      label="Global Preprocessing Step",
+      enabled=True,
+      components=[specs.ComponentType.AGGREGATION_SERVER],
+    )
+
+  def transform_data(self, *, data, agg_client):
+    global_count = agg_client.sum([float(len(data))])
+    return data
+```
+
+If a step uses a different client parameter name for `transform_data()`, override
+`aggregation_client_parameter_name()` and keep the base
+`transform_data_and_metadata()` implementation:
+
+```python
+@classmethod
+def aggregation_client_parameter_name(cls) -> str:
+  return "client"
+
+
+def transform_data(self, *, data, client):
+  global_count = client.sum([float(len(data))])
+  return data
+```
+
+Do not rely on `**kwargs` for this injection path; the parameter must be
+declared explicitly.
+
 ## Aggregation server (SUM/MIN/MAX)
 
 The aggregation server provides centralized SUM/MIN/MAX operations across

documentation/context/conventions.md

@@ -52,6 +52,12 @@
   should not repeat schema shape, options, defaults, min/max bounds, or
   requiredness.
 - UDF helpers are registered through `@exareme3_udf`.
+- Preprocessing steps that need aggregation declare
+  `ComponentType.AGGREGATION_SERVER` in their specification components and must
+  explicitly declare the configured aggregation client parameter, usually
+  `agg_client`, on `transform_data()`. Aggregation-aware preprocessing steps
+  should use the base `transform_data_and_metadata()` implementation instead of
+  overriding it. Do not rely on `**kwargs` for aggregation-client injection.
 - New algorithm work should use `.agents/skills/exaflow-algorithm-scaffold` and
   `.agents/skills/exaflow-algorithm-validate`.
 

exaflow/algorithms/exareme3/preprocessing/missing_values_handler.py

@@ -196,6 +196,7 @@ def transform_data_and_metadata(
         *,
         data: pd.DataFrame,
         metadata: Dict[str, dict],
+        agg_client=None,
     ) -> tuple[pd.DataFrame, Dict[str, dict]]:
         # Enforce categorical constant-fill constraints at runtime too.
         self._validate_categorical_constant_fill_values(metadata=metadata)

exaflow/algorithms/exareme3/utils/preprocessing_step.py

@@ -1,13 +1,17 @@
+import inspect
 from abc import ABC
 from abc import abstractmethod
 from typing import Dict
 from typing import List
 
 import pandas as pd
 
+from exaflow.algorithms.specifications import ComponentType
 from exaflow.algorithms.specifications import PreprocessingStepSpecification
 from exaflow.algorithms.utils.inputdata_utils import Inputdata
 
+DEFAULT_AGGREGATION_CLIENT_PARAMETER_NAME = "agg_client"
+
 
 class PreprocessingStep(ABC):
     def __init__(
@@ -29,6 +33,14 @@ def required_input_variables(cls) -> List[str]:
         """
         return []
 
+    @classmethod
+    def aggregation_server_required(cls) -> bool:
+        return ComponentType.AGGREGATION_SERVER in cls.get_specification().components
+
+    @classmethod
+    def aggregation_client_parameter_name(cls) -> str:
+        return DEFAULT_AGGREGATION_CLIENT_PARAMETER_NAME
+
     @abstractmethod
     def validate_params(
         self,
@@ -67,8 +79,20 @@ def transform_data_and_metadata(
         *,
         data: pd.DataFrame,
         metadata: Dict[str, dict],
+        agg_client=None,
     ) -> tuple[pd.DataFrame, Dict[str, dict]]:
-        """Convenience wrapper for transform_data + transform_metadata."""
-        return self.transform_data(data=data), self.transform_metadata(
+        """
+        Convenience wrapper for transform_data + transform_metadata.
+
+        Preprocessing steps that declare ComponentType.AGGREGATION_SERVER can
+        receive the aggregation client by using this base method and explicitly
+        adding the configured aggregation client parameter to transform_data.
+        """
+        transform_data_kwargs = {"data": data}
+        agg_client_name = self.aggregation_client_parameter_name()
+        transform_data_signature = inspect.signature(self.transform_data)
+        if agg_client_name in transform_data_signature.parameters:
+            transform_data_kwargs[agg_client_name] = agg_client
+        return self.transform_data(**transform_data_kwargs), self.transform_metadata(
             metadata=metadata
         )

exaflow/controller/services/algorithm_execution_strategy_factory.py

@@ -1,6 +1,7 @@
 from typing import List
 from typing import Type
 
+from exaflow import exareme3_preprocessing_step_classes
 from exaflow.algorithms.specifications import AlgorithmType
 from exaflow.algorithms.specifications import ComponentType
 from exaflow.controller.services.api.algorithm_spec_dtos import specifications
@@ -29,7 +30,7 @@ def get_algorithm_execution_strategy(
 
     algorithm_name = analysis_request_dto.algorithm.name
     algo_type = specifications.get_algorithm_type(algorithm_name)
-    components = specifications.get_component_types(algorithm_name)
+    components = _get_required_component_types(analysis_request_dto)
     controller = _get_algorithm_controller(algo_type)
     strategy_type = _get_algorithm_strategy_type(algo_type, components)
 
@@ -47,6 +48,20 @@ def _get_algorithm_controller(algo_type: AlgorithmType) -> ControllerI:
     )
 
 
+def _get_required_component_types(
+    analysis_request_dto: AnalysisRequestDTO,
+) -> List[ComponentType]:
+    algorithm_name = analysis_request_dto.algorithm.name
+    components = list(specifications.get_component_types(algorithm_name))
+    for preprocessing_step in analysis_request_dto.preprocessing or []:
+        preprocessing_step_cls = exareme3_preprocessing_step_classes[
+            preprocessing_step.name
+        ]
+        if preprocessing_step_cls.aggregation_server_required():
+            components.append(ComponentType.AGGREGATION_SERVER)
+    return components
+
+
 def _get_algorithm_strategy_type(
     algo_type: AlgorithmType,
     algo_component_types: List[ComponentType],

exaflow/worker/exareme3/udf/udf_service.py

@@ -31,8 +31,12 @@ def run_udf(
 ):
     udf = _get_udf_or_raise(udf_registry_key)
     udf = _wrap_udf_with_lazy_aggregation_if_enabled(udf_registry_key, udf)
-    agg_client = _create_aggregation_client_if_required(request_id, udf_registry_key)
     preprocessing = system_args.preprocessing
+    agg_client = _create_aggregation_client_if_required(
+        request_id=request_id,
+        udf_registry_key=udf_registry_key,
+        preprocessing=preprocessing,
+    )
     metadata = enforce_enum_order(system_args.metadata)
     data = _load_worker_data_for_udf(
         inputdata=system_args.inputdata,
@@ -54,6 +58,7 @@ def run_udf(
             data=data,
             metadata=metadata,
             agg_client=agg_client,
+            agg_client_name=_get_udf_aggregation_client_name(udf_registry_key),
         )
     except BadInputError as e:
         logger = get_logger()
@@ -87,13 +92,36 @@ def _wrap_udf_with_lazy_aggregation_if_enabled(udf_registry_key: str, udf):
 def _create_aggregation_client_if_required(
     request_id,
     udf_registry_key: str,
+    preprocessing: Optional[List[Dict[str, Any]]],
 ) -> Optional[AggregationClient]:
-    if not exareme3_registry.aggregation_server_required(udf_registry_key):
+    if not (
+        exareme3_registry.aggregation_server_required(udf_registry_key)
+        or _preprocessing_requires_aggregation_server(preprocessing)
+    ):
         return None
     agg_dns = worker_config.aggregation_server.dns
     return AggregationClient(request_id, aggregator_dns=agg_dns)
 
 
+def _get_udf_aggregation_client_name(udf_registry_key: str) -> Optional[str]:
+    if not exareme3_registry.aggregation_server_required(udf_registry_key):
+        return None
+    return exareme3_registry.agg_client_name(udf_registry_key)
+
+
+def _preprocessing_requires_aggregation_server(
+    preprocessing: Optional[List[Dict[str, Any]]],
+) -> bool:
+    for preprocessing_step in preprocessing or []:
+        preprocessing_step_name = preprocessing_step["name"]
+        preprocessing_step_cls = exareme3_preprocessing_step_classes[
+            preprocessing_step_name
+        ]
+        if preprocessing_step_cls.aggregation_server_required():
+            return True
+    return False
+
+
 def _collect_extra_columns(preprocessing: Optional[List[Dict[str, Any]]]) -> set[str]:
     extra_columns = set()
     for preprocessing_step in preprocessing or []:
@@ -171,6 +199,7 @@ def _apply_preprocessing_steps_to_data_and_metadata(
         data, metadata = preprocessing_step.transform_data_and_metadata(
             data=data,
             metadata=step_metadata,
+            agg_client=agg_client,
         )
         _check_min_rows_or_raise(
             data=data,
@@ -187,10 +216,12 @@ def _execute_udf(
     data,
     metadata: dict,
     agg_client: Optional[AggregationClient],
+    agg_client_name: Optional[str],
 ):
-    if agg_client:
-        kw_args["agg_client"] = agg_client
+    udf_signature = inspect.signature(udf)
+    if agg_client and agg_client_name and agg_client_name in udf_signature.parameters:
+        kw_args[agg_client_name] = agg_client
     kw_args["data"] = data
-    if "metadata" in inspect.signature(udf).parameters:
+    if "metadata" in udf_signature.parameters:
         kw_args["metadata"] = metadata
     return udf(**kw_args)

tests/standalone_tests/controller/services/test_algorithm_execution_strategy_factory.py

@@ -0,0 +1,92 @@
+import sys
+import types
+
+from exaflow.algorithms.specifications import ComponentType
+
+dns_module = types.ModuleType("dns")
+dns_resolver_module = types.ModuleType("dns.resolver")
+dns_module.resolver = dns_resolver_module
+sys.modules.setdefault("dns", dns_module)
+sys.modules.setdefault("dns.resolver", dns_resolver_module)
+
+from exaflow.controller.services import algorithm_execution_strategy_factory as factory
+from exaflow.controller.services.api.analysis_request_dtos import AnalysisAlgorithmDTO
+from exaflow.controller.services.api.analysis_request_dtos import AnalysisInputDataDTO
+from exaflow.controller.services.api.analysis_request_dtos import (
+    AnalysisPreprocessingStepDTO,
+)
+from exaflow.controller.services.api.analysis_request_dtos import AnalysisRequestDTO
+from exaflow.controller.services.exareme3.strategies import Exareme3Strategy
+from exaflow.controller.services.exareme3.strategies import (
+    Exareme3WithAggregationServerStrategy,
+)
+
+
+class _NoAggregationPreprocessing:
+    @classmethod
+    def aggregation_server_required(cls):
+        return False
+
+
+class _AggregationPreprocessing:
+    @classmethod
+    def aggregation_server_required(cls):
+        return True
+
+
+class _Specifications:
+    def get_component_types(self, algo_name):
+        return []
+
+
+def _request(preprocessing=None):
+    return AnalysisRequestDTO(
+        inputdata=AnalysisInputDataDTO(
+            data_model="dementia",
+            datasets=["ppmi0"],
+            variables=["x"],
+        ),
+        preprocessing=preprocessing,
+        algorithm=AnalysisAlgorithmDTO(name="sample_algo", y=["x"]),
+    )
+
+
+def test_preprocessing_aggregation_requirement_contributes_component(monkeypatch):
+    monkeypatch.setattr(factory, "specifications", _Specifications())
+    monkeypatch.setattr(
+        factory,
+        "exareme3_preprocessing_step_classes",
+        {
+            "local": _NoAggregationPreprocessing,
+            "global": _AggregationPreprocessing,
+        },
+    )
+
+    components = factory._get_required_component_types(
+        _request(
+            preprocessing=[
+                AnalysisPreprocessingStepDTO(name="local", parameters={}),
+                AnalysisPreprocessingStepDTO(name="global", parameters={}),
+            ]
+        )
+    )
+
+    assert ComponentType.AGGREGATION_SERVER in components
+
+
+def test_preprocessing_aggregation_component_selects_aggregation_strategy():
+    strategy_type = factory._get_algorithm_strategy_type(
+        algo_type=factory.AlgorithmType.EXAREME3,
+        algo_component_types=[ComponentType.AGGREGATION_SERVER],
+    )
+
+    assert strategy_type is Exareme3WithAggregationServerStrategy
+
+
+def test_no_aggregation_component_selects_plain_exareme3_strategy():
+    strategy_type = factory._get_algorithm_strategy_type(
+        algo_type=factory.AlgorithmType.EXAREME3,
+        algo_component_types=[],
+    )
+
+    assert strategy_type is Exareme3Strategy