fix(sdk): Fix ensembler ensemble method (caraml-dev#413)

deadlycoconuts · web-flow · commit 9a53d9fecc1a · 2025-07-17T14:02:13.000+08:00
* Fix turing ensembler sdk

* Add sdk tests

* Fix broken tests
diff --git a/sdk/tests/__init__.py b/sdk/tests/__init__.py
@@ -1,5 +1,6 @@
 from datetime import date, datetime
 from dateutil.tz import tzutc
+from typing import Optional, Union, Any, Dict
 from turing import generated as client
 import turing.ensembler
 
@@ -17,9 +18,8 @@ def utc_date(date_str: str):
     return datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S.%fZ").replace(tzinfo=tzutc())
 
 
-class MyTestEnsembler(turing.ensembler.PyFunc):
+class MyTestEnsemblerJob(turing.ensembler.PyFunc):
     import pandas
-    from typing import Any, Optional
 
     def __init__(self, default: float):
         self._default = default
@@ -29,11 +29,26 @@ def initialize(self, artifacts: dict):
 
     def ensemble(
         self,
-        input: pandas.Series,
-        predictions: pandas.Series,
-        treatment_config: Optional[dict],
+        input: Union[pandas.Series, Dict[str, Any]],
+        predictions: Union[pandas.Series, Dict[str, Any]],
+        **kwargs,
     ) -> Any:
         if input["treatment"] in predictions:
             return predictions[input["treatment"]]
         else:
             return self._default
+
+
+class MyTestEnsemblerService(turing.ensembler.PyFunc):
+    def __init__(self, default: float):
+        self._default = default
+
+    def initialize(self, artifacts: dict):
+        pass
+
+    def ensemble(
+        self,
+        enricher_response,
+        **kwargs,
+    ) -> Any:
+        return enricher_response
diff --git a/sdk/tests/ensembler_test.py b/sdk/tests/ensembler_test.py
@@ -16,9 +16,10 @@
 with open(os.path.join(data_dir, "list_jobs_0000.json")) as f:
     list_jobs_0000 = f.read()
 
-def test_predict():
+
+def test_ensemble_job_predict():
     default_value = random.random()
-    ensembler = tests.MyTestEnsembler(default_value)
+    ensembler = tests.MyTestEnsemblerJob(default_value)
 
     model_input = pandas.DataFrame(
         data={
@@ -43,6 +44,98 @@ def test_predict():
 
     assert_series_equal(expected, result)
 
+
+def test_ensemble_service_predict_with_enricher_response():
+    default_value = random.random()
+    ensembler = tests.MyTestEnsemblerService(default_value)
+
+    model_input = {
+        "headers": {
+            "Host": "example.com",
+            "Referer": "https://example.com/",
+        },
+        "body": {
+          # original request payload unmodified
+            "request": {
+                "dummy": "dummy"
+            },
+            "response": {
+                "route_responses": [
+                    {
+                        "route": "control",
+                        "data": {
+                            "abc": 1,
+                        },
+                    },
+                    {
+                        "route": "xgboost-ordinal",
+                        "data": {
+                            "def": 2,
+                        }
+                    },
+                ],
+                "experiment": {
+                    "configuration": {
+                        "output": "xyz",
+                    },
+                },
+                "enricher_response": {
+                    "output": "meow"
+                }
+            },
+        },
+    }
+
+    expected = {"output": "meow"}
+    result = ensembler.predict(context=None, model_input=model_input)
+
+    assert expected == result
+
+
+def test_ensemble_service_predict_without_enricher_response():
+    default_value = random.random()
+    ensembler = tests.MyTestEnsemblerService(default_value)
+
+    model_input = {
+        "headers": {
+            "Host": "example.com",
+            "Referer": "https://example.com/",
+        },
+        "body": {
+          # original request payload unmodified
+            "request": {
+                "dummy": "dummy"
+            },
+            "response": {
+                "route_responses": [
+                    {
+                        "route": "control",
+                        "data": {
+                            "abc": 1,
+                        },
+                    },
+                    {
+                        "route": "xgboost-ordinal",
+                        "data": {
+                            "def": 2,
+                        }
+                    },
+                ],
+                "experiment": {
+                    "configuration": {
+                        "output": "xyz",
+                    },
+                },
+            },
+        },
+    }
+
+    expected = None
+    result = ensembler.predict(context=None, model_input=model_input)
+
+    assert expected == result
+
+
 @pytest.mark.parametrize("num_ensemblers", [6])
 def test_list_ensemblers(
     turing_api, project, generic_ensemblers, use_google_oauth, active_project_magic_mock
@@ -73,6 +166,7 @@ def test_list_ensemblers(
         for actual, expected in zip(actual, generic_ensemblers):
             assert actual == turing.PyFuncEnsembler.from_open_api(expected)
 
+
 @patch("google.cloud.storage.Client")
 @patch("requests.Session.request")
 @patch("urllib3.PoolManager.request")
@@ -103,15 +197,16 @@ def test_create_ensembler(
 
     actual = turing.PyFuncEnsembler.create(
         name=pyfunc_ensembler.name,
-        ensembler_instance=tests.MyTestEnsembler(0.01),
+        ensembler_instance=tests.MyTestEnsemblerJob(0.01),
         conda_env={
             "channels": ["defaults"],
             "dependencies": ["python=3.9.0", {"pip": ["test-lib==0.0.1"]}],
         },
     )
     
     assert actual == turing.PyFuncEnsembler.from_open_api(pyfunc_ensembler)
-    
+
+
 @pytest.mark.parametrize(("num_ensemblers", "ensembler_name"), [(3, "updated")])
 @patch("google.cloud.storage.Client")
 @patch("requests.Session.request")
@@ -177,7 +272,7 @@ def test_update_ensembler(
 
     actual.update(
         name=pyfunc_ensembler.name,
-        ensembler_instance=tests.MyTestEnsembler(0.06),
+        ensembler_instance=tests.MyTestEnsemblerJob(0.06),
         conda_env={
             "channels": ["defaults"],
             "dependencies": ["python>=3.8.0", {"pip": ["test-lib==0.0.1"]}],
@@ -190,7 +285,8 @@ def test_update_ensembler(
     )
     
     assert actual == turing.PyFuncEnsembler.from_open_api(pyfunc_ensembler)
-    
+
+
 @pytest.mark.parametrize(("num_ensemblers", "ensembler_name"), [(3, "updated")])
 def test_update_ensembler_existing_router_version(
     turing_api,
@@ -235,7 +331,7 @@ def test_update_ensembler_existing_router_version(
         with pytest.raises(ValueError) as error:
             actual.update(
                 name=pyfunc_ensembler.name,
-                ensembler_instance=tests.MyTestEnsembler(0.06),
+                ensembler_instance=tests.MyTestEnsemblerJob(0.06),
                 conda_env={
                     "channels": ["defaults"],
                     "dependencies": ["python>=3.8.0", {"pip": ["test-lib==0.0.1"]}],
@@ -251,7 +347,8 @@ def test_update_ensembler_existing_router_version(
         expected_error_message = "There is pending router version using this ensembler. Please wait for the router version to be deployed or undeploy it, before updating the ensembler."
         actual_error_message = str(error.value)
         assert expected_error_message == actual_error_message
-        
+
+
 @pytest.mark.parametrize(("num_ensemblers", "ensembler_name"), [(3, "updated")])
 def test_update_ensembler_existing_job(
     turing_api, project, generic_ensemblers, pyfunc_ensembler, use_google_oauth, active_project_magic_mock
@@ -297,7 +394,7 @@ def test_update_ensembler_existing_job(
         with pytest.raises(ValueError) as error:
             actual.update(
                 name=pyfunc_ensembler.name,
-                ensembler_instance=tests.MyTestEnsembler(0.06),
+                ensembler_instance=tests.MyTestEnsemblerJob(0.06),
                 conda_env={
                     "channels": ["defaults"],
                     "dependencies": ["python>=3.8.0", {"pip": ["test-lib==0.0.1"]}],
@@ -313,7 +410,8 @@ def test_update_ensembler_existing_job(
         expected_error_message = "There is pending ensembling job using this ensembler. Please wait for the ensembling job to be completed or terminate it, before updating the ensembler."
         actual_error_message = str(error.value)
         assert expected_error_message == actual_error_message
-        
+
+
 @patch("google.cloud.storage.Client")
 @patch("requests.Session.request")
 @patch("urllib3.PoolManager.request")
diff --git a/sdk/turing/ensembler.py b/sdk/turing/ensembler.py
@@ -113,12 +113,17 @@ def _ensemble_request(self, model_input: Dict[str, Any]) -> Any:
             # Deletes route from the dictionary as it is a duplicate of the key
             del routes_to_response[prediction["route"]]["route"]
 
+        # This is for older Turing routers which do not pass the enricher_response to the ensembler
+        enricher_response = None
+        if "enricher_response" in request_body["response"]:
+            enricher_response = request_body["response"]["enricher_response"]
+
         try:
             return self.ensemble(
                 input=request_body["request"],
                 predictions=routes_to_response,
                 treatment_config=request_body["response"]["experiment"],
-                enricher_response=request_body["response"]["enricher_response"],
+                enricher_response=enricher_response,
                 headers=model_input["headers"],
             )
         except TypeError as e:
