refactor(core): do not change signature of functions decorated with custom_experiment (#261)

* refactor(core): custom_experiment decorator does not change signature

Implementation of #252

* chore(core): apply suggestions from code review

* feat(core): Validate custom_experiment parameters in wrapper

Author: michael-johnston

orchestrator/modules/actuators/custom_experiments.py

@@ -5,7 +5,6 @@
 import logging
 import typing
 import uuid
-from functools import wraps
 
 import pydantic
 import ray
@@ -31,6 +30,7 @@
     ObservedProperty,
     ObservedPropertyValue,
 )
+from orchestrator.schema.point import SpacePoint
 from orchestrator.schema.property import (
     AbstractPropertyDescriptor,
     ConstitutiveProperty,
@@ -313,47 +313,6 @@ def calculate_density(mass, volume):
     logger = logging.getLogger("custom_experiment_decorator")
 
     def decorator(func):
-
-        @wraps(func)
-        def wrapper(
-            entity: Entity, experiment: Experiment
-        ) -> list[ObservedPropertyValue]:
-            """
-            Wrapper function that converts Entity+Experiment to dict and calls the wrapped function.
-            """
-            input_values = experiment.propertyValuesFromEntity(entity)
-            result_dict = func(**input_values)
-            observed_property_values = []
-            for property_identifier, value in result_dict.items():
-                observed_property = experiment.observedPropertyForTargetIdentifier(
-                    property_identifier
-                )
-                if not observed_property:
-                    raise ValueError(
-                        f"{experiment.identifier} returned a property called {property_identifier}, however "
-                        f"the experiment definition does not define an output property with this name"
-                    )
-
-                observed_property_value = ObservedPropertyValue(
-                    property=observed_property, value=value
-                )
-                observed_property_values.append(observed_property_value)
-
-            return observed_property_values
-
-        # Set the wrapper's __signature__  so it is (entity,experiment)
-        # This is required for the wrapped function to be used with ray.remote
-        import inspect
-
-        wrapper.__signature__ = inspect.Signature(
-            [
-                inspect.Parameter("entity", inspect.Parameter.POSITIONAL_OR_KEYWORD),
-                inspect.Parameter(
-                    "experiment", inspect.Parameter.POSITIONAL_OR_KEYWORD
-                ),
-            ]
-        )
-
         # If we were not given information on required/optional properties
         # or parameterization try to infer it
         # This function will log a critical error message and raise exception
@@ -380,12 +339,12 @@ def wrapper(
             )
             raise
 
-        # Store decorator arguments as function attributes
-        wrapper._decorator_required_properties = _required_properties
-        wrapper._decorator_optional_properties = _optional_properties
-        wrapper._decorator_parameterization = _parameterization
-        wrapper._original_func = func
-        wrapper._is_custom_experiment = True
+        # Store decorator arguments as function attributes (on func itself)
+        func._decorator_required_properties = _required_properties
+        func._decorator_optional_properties = _optional_properties
+        func._decorator_parameterization = _parameterization
+        func._original_func = func
+        func._is_custom_experiment = True
 
         # Create an ExperimentModuleConf instance describing where the function is
         metadata["module"] = ExperimentModuleConf(
@@ -415,12 +374,44 @@ def wrapper(
             deprecated=False,
             metadata=metadata,
         )
-        wrapper._experiment = experiment
+        func._experiment = experiment
 
         # Add the experiment to the module-level catalog
         _custom_experiments_catalog.addExperiment(experiment)
 
-        return wrapper
+        from functools import wraps
+
+        @wraps(func)
+        def validated_func(*args, **kwargs):
+            # Build property dict from either kwargs or args
+            # Prefer kwargs, but support positional for backwards compatibility
+            import inspect
+
+            sig = inspect.signature(func)
+            bound_args = sig.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+            param_dict = dict(bound_args.arguments)
+
+            # Validate using SpacePoint and Experiment.validate_entity
+            spoint = SpacePoint(entity=param_dict)
+            entity = spoint.to_entity()
+            if not experiment.validate_entity(entity, verbose=True):
+                raise ValueError(
+                    f"Arguments {param_dict} do not match required/optional properties for experiment '{experiment.identifier}'. "
+                    f"See logs/stderr for reasons, or check experiment.requiredProperties/optionalProperties."
+                )
+            # Call the original with the unpacked arguments
+            return func(*args, **kwargs)
+
+        # Attach metadata to validated_func not func, so end users get the right attributes
+        validated_func._decorator_required_properties = _required_properties
+        validated_func._decorator_optional_properties = _optional_properties
+        validated_func._decorator_parameterization = _parameterization
+        validated_func._original_func = func
+        validated_func._is_custom_experiment = True
+        validated_func._experiment = experiment
+
+        return validated_func
 
     return decorator
 
@@ -513,7 +504,51 @@ def get_custom_experiments_catalog() -> (
     return _custom_experiments_catalog
 
 
-async def custom_experiment_wrapper(
+def _call_decorated_custom_experiment(
+    function: typing.Callable, target_experiment: Experiment, entity: Entity
+) -> list[ObservedPropertyValue]:
+
+    # Build input dict using experiment values from entity
+    input_values = target_experiment.propertyValuesFromEntity(entity)
+    # Call function with unpacked parameters
+    result_dict = function(**input_values)
+
+    # Create observed property values
+    observed_property_values = []
+    for property_identifier, value in result_dict.items():
+        observed_property = target_experiment.observedPropertyForTargetIdentifier(
+            property_identifier
+        )
+        if not observed_property:
+            raise ValueError(
+                f"{target_experiment.identifier} returned a property called {property_identifier}, however "
+                f"the experiment definition does not define an output property with this name"
+            )
+        observed_property_value = ObservedPropertyValue(
+            property=observed_property, value=value
+        )
+        observed_property_values.append(observed_property_value)
+
+    return observed_property_values
+
+
+def _call_legacy_custom_experiment(
+    function: typing.Callable,
+    target_experiment: Experiment,
+    entity: Entity,
+    parameters: dict | None = None,
+) -> list[ObservedPropertyValue]:
+    # For legacy case or other functions, check for parameters kwarg else pass entity/experiment
+    func_signature = inspect.signature(function)
+    if "parameters" in func_signature.parameters:
+        values = function(entity, target_experiment, parameters=parameters)
+    else:
+        values = function(entity, target_experiment)
+
+    return values
+
+
+async def custom_experiment_executor(
     function: typing.Callable,
     parameters: dict,
     measurement_request: MeasurementRequest,
@@ -532,16 +567,22 @@ async def custom_experiment_wrapper(
 
     measurement_results = []
     for entity in measurement_request.entities:
-        # Inspect function to see if it has a keyword parameter "parameters"
-        func_signature = inspect.signature(function)
-        func_param_names = set(func_signature.parameters.keys())
         try:
-            if "parameters" in func_param_names:
-                values = function(entity, target_experiment, parameters=parameters)
+            # Check if this is a custom experiment decorated function
+            if getattr(function, "_is_custom_experiment", False):
+                values = _call_decorated_custom_experiment(
+                    function=function,
+                    target_experiment=target_experiment,
+                    entity=entity,
+                )
             else:
-                values = function(entity, target_experiment)
+                values = _call_legacy_custom_experiment(
+                    function=function,
+                    target_experiment=target_experiment,
+                    entity=entity,
+                    parameters=parameters,
+                )
 
-            # Record the results in the entity
             if len(values) > 0:
                 measurement_result = ValidMeasurementResult(
                     entityIdentifier=entity.identifier, measurements=values
@@ -683,7 +724,7 @@ async def submit(
                 **targetExperiment.model_dump(),
             )
 
-        await custom_experiment_wrapper(
+        await custom_experiment_executor(
             self._functionImplementations[
                 request.experimentReference.experimentIdentifier
             ],

website/docs/actuators/creating-custom-experiments.md

@@ -41,10 +41,10 @@ my_experiment = "my_custom_package.experiments"
 
 In the simplest case:
 
-- type the parameters (using python `typing`)
+- type your parameters using python `typing`
 - return the output in a dictionary of key value pairs
-- define the keys of this dictionary in the `output_property_identifiers`
-parameter of the decorator
+- define the set of output property keys in the
+`output_property_identifiers` parameter of the decorator
 
 ```python
 from typing import Dict, Any
@@ -286,15 +286,36 @@ This is illustrated in the above example.
 
 ## Using your decorated function in code
 
-The decorated function is wrapped to take `ado` internal
-data structures, and you would not typically need to
-call it directly. However, the decorated experiment function is
-still regular Python and can be called:
+The decorated function can be called
+directly in Python as normal e.g.,
 
 ```python
-# Access the original function (undecorated)
-original = calculate_density._original_func
-print(original(8, 4))  # {'density': 2}
+result = calculate_density(8, 4)  # {'density': 2}
+```
+
+The `custom_experiment` decorator attaches the
+ado `Experiment` object generated from the decoration as an attribute e.g.
+
+```python
+from orchestrator.schema.experiment import Experiment
+
+exp_obj: Experiment = calculate_density._experiment
+print(exp_obj.identifier)  # e.g., 'calculate_density'
+print(exp_obj.requiredProperties)
+print(exp_obj.optionalProperties)
+print(exp_obj.targetProperties)
+```
+
+When you call the decorated function, its arguments are
+automatically validated against the required and optional inputs
+specified in the decorator, including domain constraints.
+If you call it with missing, extra, or out-of-domain arguments,
+the function will raise a `ValueError` describing what was invalid and why.
+For example:
+
+```python
+# Value outside domain - an error will be raised
+result = calculate_density(mass=0, volume=10)  
 ```
 
 ## Next Steps