diff --git a/documentation/algorithms-user-guide.md b/documentation/algorithms-user-guide.md index 9a5a4ac93..5874a44a2 100644 --- a/documentation/algorithms-user-guide.md +++ b/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 diff --git a/documentation/context/conventions.md b/documentation/context/conventions.md index fcc62b950..82613bd19 100644 --- a/documentation/context/conventions.md +++ b/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`. diff --git a/exaflow/algorithms/exareme3/preprocessing/missing_values_handler.py b/exaflow/algorithms/exareme3/preprocessing/missing_values_handler.py index 3b25f8070..a254c5109 100644 --- a/exaflow/algorithms/exareme3/preprocessing/missing_values_handler.py +++ b/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) diff --git a/exaflow/algorithms/exareme3/utils/preprocessing_step.py b/exaflow/algorithms/exareme3/utils/preprocessing_step.py index 6b1e7f777..2d2fc0d92 100644 --- a/exaflow/algorithms/exareme3/utils/preprocessing_step.py +++ b/exaflow/algorithms/exareme3/utils/preprocessing_step.py @@ -1,3 +1,4 @@ +import inspect from abc import ABC from abc import abstractmethod from typing import Dict @@ -5,9 +6,12 @@ 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 ) diff --git a/exaflow/controller/services/algorithm_execution_strategy_factory.py b/exaflow/controller/services/algorithm_execution_strategy_factory.py index a7e8d204c..c2c614eb5 100644 --- a/exaflow/controller/services/algorithm_execution_strategy_factory.py +++ b/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], diff --git a/exaflow/worker/exareme3/udf/udf_service.py b/exaflow/worker/exareme3/udf/udf_service.py index 942174496..dc12208de 100644 --- a/exaflow/worker/exareme3/udf/udf_service.py +++ b/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) diff --git a/tests/standalone_tests/controller/services/test_algorithm_execution_strategy_factory.py b/tests/standalone_tests/controller/services/test_algorithm_execution_strategy_factory.py new file mode 100644 index 000000000..5a0f46356 --- /dev/null +++ b/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 diff --git a/tests/standalone_tests/exareme3/test_udf_service.py b/tests/standalone_tests/exareme3/test_udf_service.py index 60fd6e72c..3124e7038 100644 --- a/tests/standalone_tests/exareme3/test_udf_service.py +++ b/tests/standalone_tests/exareme3/test_udf_service.py @@ -1,6 +1,8 @@ import pandas as pd import pytest +from exaflow.algorithms import specifications as specs +from exaflow.algorithms.exareme3.utils.preprocessing_step import PreprocessingStep from exaflow.worker.exareme3.udf import udf_service from exaflow.worker_communication import InsufficientDataError @@ -17,32 +19,97 @@ def close(self): self.closed = True -class _DropAllStep: +class _BaseStep: def __init__(self, *, params): self._params = params - def transform_data_and_metadata(self, *, data, metadata): - return data.iloc[0:0], metadata + @classmethod + def aggregation_server_required(cls): + return False + @classmethod + def aggregation_client_parameter_name(cls): + return "agg_client" + + +class _DropAllStep(_BaseStep): + def transform_data_and_metadata(self, *, data, metadata, agg_client=None): + return data.iloc[0:0], metadata -class _FirstOrderedStep: - def __init__(self, *, params): - self._params = params - def transform_data_and_metadata(self, *, data, metadata): +class _FirstOrderedStep(_BaseStep): + def transform_data_and_metadata(self, *, data, metadata, agg_client=None): self._params["calls"].append("first") return data, metadata -class _SecondOrderedStep: - def __init__(self, *, params): - self._params = params - - def transform_data_and_metadata(self, *, data, metadata): +class _SecondOrderedStep(_BaseStep): + def transform_data_and_metadata(self, *, data, metadata, agg_client=None): self._params["calls"].append("second") return data, metadata +class _AggAwareStep(_BaseStep): + @classmethod + def aggregation_server_required(cls): + return True + + def transform_data_and_metadata(self, *, data, metadata, agg_client): + self._params["received_client"].append(agg_client) + return data, metadata + + +class _BaseTransformAggStep(PreprocessingStep): + def __init__(self, *, params): + super().__init__(params=params) + + @classmethod + def get_specification(cls) -> specs.PreprocessingStepSpecification: + return specs.PreprocessingStepSpecification( + name="base_transform_agg", + desc="test", + documentation="test", + label="test", + enabled=True, + components=[specs.ComponentType.AGGREGATION_SERVER], + ) + + def validate_params(self, *, inputdata, metadata): + pass + + def transform_variables(self, *, variables): + return variables + + def transform_metadata(self, *, metadata): + return metadata + + def transform_data(self, *, data, agg_client): + self._params["received_client"].append(agg_client) + return data + + +class _CustomAggClientNameStep(_BaseTransformAggStep): + @classmethod + def aggregation_client_parameter_name(cls): + return "client" + + def transform_data(self, *, data, client): + self._params["received_client"].append(client) + return data + + +class _NoAggSpecStep(_BaseTransformAggStep): + @classmethod + def get_specification(cls) -> specs.PreprocessingStepSpecification: + return specs.PreprocessingStepSpecification( + name="no_agg", + desc="test", + documentation="test", + label="test", + enabled=True, + ) + + def test_check_min_rows_uses_pandas_row_count(monkeypatch): monkeypatch.setattr(udf_service.worker_config.privacy, "minimum_row_count", 3) data = pd.DataFrame({"x": [1, 2]}) @@ -112,3 +179,118 @@ def test_apply_preprocessing_uses_request_order(monkeypatch): ) assert calls == ["second", "first"] + + +def test_preprocessing_step_base_detects_aggregation_server_component(): + assert _BaseTransformAggStep.aggregation_server_required() is True + assert _NoAggSpecStep.aggregation_server_required() is False + + +def test_apply_preprocessing_passes_agg_client_only_to_steps_that_accept_it( + monkeypatch, +): + received_client = [] + agg_client = _DummyAggClient() + monkeypatch.setattr( + udf_service, + "exareme3_preprocessing_step_classes", + { + "first": _FirstOrderedStep, + "agg_aware": _AggAwareStep, + }, + ) + + udf_service._apply_preprocessing_steps_to_data_and_metadata( + data=pd.DataFrame({"x": [1, 2]}), + metadata={"x": {"is_categorical": False}}, + preprocessing=[ + {"name": "first", "parameters": {"calls": []}}, + {"name": "agg_aware", "parameters": {"received_client": received_client}}, + ], + check_min_rows=False, + agg_client=agg_client, + ) + + assert received_client == [agg_client] + + +def test_apply_preprocessing_supports_custom_agg_client_parameter_name(monkeypatch): + received_client = [] + agg_client = _DummyAggClient() + monkeypatch.setattr( + udf_service, + "exareme3_preprocessing_step_classes", + {"custom": _CustomAggClientNameStep}, + ) + + udf_service._apply_preprocessing_steps_to_data_and_metadata( + data=pd.DataFrame({"x": [1, 2]}), + metadata={"x": {"is_categorical": False}}, + preprocessing=[ + {"name": "custom", "parameters": {"received_client": received_client}} + ], + check_min_rows=False, + agg_client=agg_client, + ) + + assert received_client == [agg_client] + + +def test_base_transform_data_and_metadata_passes_agg_client_to_transform_data(): + received_client = [] + agg_client = _DummyAggClient() + step = _BaseTransformAggStep(params={"received_client": received_client}) + + data, _ = step.transform_data_and_metadata( + data=pd.DataFrame({"x": [1, 2]}), + metadata={"x": {"is_categorical": False}}, + agg_client=agg_client, + ) + + assert list(data["x"]) == [1, 2] + assert received_client == [agg_client] + + +def test_preprocessing_requirement_can_create_client_for_non_aggregation_udf( + monkeypatch, +): + created_clients = [] + + class _CreatedAggClient(_DummyAggClient): + def __init__(self, request_id, aggregator_dns=None): + super().__init__() + self.request_id = request_id + self.aggregator_dns = aggregator_dns + created_clients.append(self) + + monkeypatch.setattr(udf_service, "AggregationClient", _CreatedAggClient) + monkeypatch.setattr( + udf_service, + "exareme3_preprocessing_step_classes", + {"agg_aware": _AggAwareStep}, + ) + + client = udf_service._create_aggregation_client_if_required( + request_id="req", + udf_registry_key="non_agg_udf", + preprocessing=[{"name": "agg_aware", "parameters": {}}], + ) + + assert client is created_clients[0] + assert client.request_id == "req" + + +def test_execute_udf_does_not_pass_preprocessing_only_agg_client_to_udf(): + def udf(data): + return {"rows": len(data)} + + result = udf_service._execute_udf( + udf=udf, + kw_args={}, + data=pd.DataFrame({"x": [1, 2]}), + metadata={"x": {"is_categorical": False}}, + agg_client=_DummyAggClient(), + agg_client_name=None, + ) + + assert result == {"rows": 2}