Bugfix: Clear default route Id from custom ensembler configs (#200)

* Miscellaneous bug fixes

* Add unit tests for the SDK changes, address PR comments

* Bugfix: Regression in display of container configs for Pyfunc

* Add type annotation to class methods

Co-authored-by: Krithika Sundararajan <[email protected]>

Author: krithika369

sdk/tests/conftest.py

@@ -16,7 +16,12 @@
 from turing.router.config.resource_request import ResourceRequest
 from turing.router.config.log_config import LogConfig, ResultLoggerType
 from turing.router.config.enricher import Enricher
-from turing.router.config.router_ensembler_config import DockerRouterEnsemblerConfig
+from turing.router.config.router_ensembler_config import (
+    EnsemblerNopConfig,
+    EnsemblerStandardConfig,
+    DockerRouterEnsemblerConfig,
+    PyfuncRouterEnsemblerConfig
+)
 from turing.router.config.common.env_var import EnvVar
 from turing.router.config.experiment_config import ExperimentConfig
 from tests.fixtures.mlflow import mock_mlflow
@@ -80,6 +85,58 @@ def generic_ensemblers(project, num_ensemblers):
             updated_at=datetime.now() + timedelta(seconds=i + 10)
         ) for i in range(1, num_ensemblers + 1)]
 
+@pytest.fixture
+def nop_router_ensembler_config():
+    return EnsemblerNopConfig(final_response_route_id="test")
+
+@pytest.fixture
+def standard_router_ensembler_config():
+    return EnsemblerStandardConfig(
+        experiment_mappings=[
+            turing.generated.models.EnsemblerStandardConfigExperimentMappings(
+                experiment="experiment-1",
+                treatment="treatment-1",
+                route="route-1"
+            ),
+            turing.generated.models.EnsemblerStandardConfigExperimentMappings(
+                experiment="experiment-2",
+                treatment="treatment-2",
+                route="route-2"
+            )
+        ],
+        fallback_response_route_id="route-1"
+    )
+
+@pytest.fixture
+def docker_router_ensembler_config():
+    return DockerRouterEnsemblerConfig(
+        image="test.io/just-a-test/turing-ensembler:0.0.0-build.0",
+        resource_request=ResourceRequest(
+            min_replica=1,
+            max_replica=3,
+            cpu_request="500m",
+            memory_request="512Mi"
+        ),
+        endpoint=f"http://localhost:5000/ensembler_endpoint",
+        timeout="500ms",
+        port=5120,
+        env=[],
+    )
+
+@pytest.fixture
+def pyfunc_router_ensembler_config():
+    return PyfuncRouterEnsemblerConfig(
+        project_id=1,
+        ensembler_id=1,
+        resource_request=ResourceRequest(
+            min_replica=1,
+            max_replica=3,
+            cpu_request="500m",
+            memory_request="512Mi"
+        ),
+        timeout="500ms",
+        env=[],
+    )
 
 @pytest.fixture
 def bucket_name():
@@ -469,7 +526,7 @@ def generic_router_version(
 
 
 @pytest.fixture
-def generic_router_config():
+def generic_router_config(docker_router_ensembler_config):
     return RouterConfig(
         environment_name="id-dev",
         name="router-1",
@@ -530,19 +587,7 @@ def generic_router_config():
                 )
             ]
         ),
-        ensembler=DockerRouterEnsemblerConfig(
-            image="test.io/just-a-test/turing-ensembler:0.0.0-build.0",
-            resource_request=ResourceRequest(
-                min_replica=1,
-                max_replica=3,
-                cpu_request="500m",
-                memory_request="512Mi"
-            ),
-            endpoint=f"http://localhost:5000/ensembler_endpoint",
-            timeout="500ms",
-            port=5120,
-            env=[],
-        )
+        ensembler=docker_router_ensembler_config
     )
 
 

sdk/tests/router/config/router_config_test.py

@@ -1,6 +1,13 @@
 import pytest
 from turing.router.config.route import Route, DuplicateRouteException, InvalidRouteException
-
+from turing.router.config.resource_request import ResourceRequest
+from turing.router.config.router_ensembler_config import (
+    DockerRouterEnsemblerConfig,
+    NopRouterEnsemblerConfig,
+    StandardRouterEnsemblerConfig,
+    PyfuncRouterEnsemblerConfig,
+    RouterEnsemblerConfig
+)
 
 @pytest.mark.parametrize(
     "actual,new_routes,expected", [
@@ -29,6 +36,7 @@ def test_set_router_config_with_invalid_routes(actual, new_routes, expected, req
     ])
 def test_set_router_config_with_invalid_default_route(actual, invalid_route_id, expected, request):
     actual = request.getfixturevalue(actual)
+    actual.ensembler = None
     actual.default_route_id = invalid_route_id
     with pytest.raises(expected):
         actual.to_open_api()
@@ -43,3 +51,76 @@ def test_set_router_config_with_invalid_default_route(actual, invalid_route_id,
 def test_remove_router_config_default_route(actual, expected, request):
     actual = request.getfixturevalue(actual)
     assert "default_route_id" not in actual.to_open_api()
+
+@pytest.mark.parametrize(
+    "router,type,nop_config,standard_config,docker_config,pyfunc_config,expected_class", [
+        pytest.param(
+            "generic_router_config",
+            "nop", "nop_router_ensembler_config", None, None, None,
+            NopRouterEnsemblerConfig
+        ),
+        pytest.param(
+            "generic_router_config",
+            "standard", None, "standard_router_ensembler_config", None, None,
+            StandardRouterEnsemblerConfig
+        ),
+        pytest.param(
+            "generic_router_config",
+            "docker", None, None, "generic_ensembler_docker_config", None,
+            DockerRouterEnsemblerConfig
+        ),
+        pytest.param(
+            "generic_router_config",
+            "pyfunc",  None, None, None, "generic_ensembler_pyfunc_config",
+            PyfuncRouterEnsemblerConfig
+        )
+    ])
+def test_set_router_config_base_ensembler(
+    router, type,
+    nop_config, standard_config, docker_config, pyfunc_config,
+    expected_class, request):
+    actual = request.getfixturevalue(router)
+    ensembler = RouterEnsemblerConfig(type=type,
+        nop_config=None if nop_config is None else request.getfixturevalue(nop_config),
+        standard_config=None if standard_config is None else request.getfixturevalue(standard_config),
+        docker_config=None if docker_config is None else request.getfixturevalue(docker_config),
+        pyfunc_config=None if pyfunc_config is None else request.getfixturevalue(pyfunc_config)
+    )
+    actual.ensembler = ensembler
+    assert isinstance(actual.ensembler, expected_class)
+
+@pytest.mark.parametrize(
+    "router_config,default_route_id,ensembler,expected", [
+        pytest.param(
+            "generic_router_config",
+            "model-a",
+            StandardRouterEnsemblerConfig(experiment_mappings=[], fallback_response_route_id="model-b"),
+            "model-b",
+        ),
+        pytest.param(
+            "generic_router_config",
+            "model-a",
+            NopRouterEnsemblerConfig(final_response_route_id="model-b"),
+            "model-b",
+        ),
+        pytest.param(
+            "generic_router_config",
+            "model-a",
+            "docker_router_ensembler_config",
+            None,
+        ),
+        pytest.param(
+            "generic_router_config",
+            "model-a",
+            "pyfunc_router_ensembler_config",
+            None,
+        )
+    ])
+def test_default_route_id_by_ensembler_config(router_config, default_route_id, ensembler, expected, request):
+    router = request.getfixturevalue(router_config)
+    router.default_route_id = default_route_id
+    router.ensembler = request.getfixturevalue(ensembler) if isinstance(ensembler, str) else ensembler
+    if expected:
+        assert router.to_open_api().to_dict()["config"]["default_route_id"] == expected
+    else:
+        assert "default_route_id" not in router.to_open_api().to_dict()["config"]

sdk/tests/router/config/router_ensembler_config_test.py

@@ -6,7 +6,6 @@
 from turing.router.config.route import InvalidRouteException
 from turing.router.config.router_ensembler_config import (RouterEnsemblerConfig,
                                                           EnsemblerNopConfig,
-                                                          EnsemblerStandardConfig,
                                                           NopRouterEnsemblerConfig,
                                                           PyfuncRouterEnsemblerConfig,
                                                           DockerRouterEnsemblerConfig,
@@ -18,55 +17,26 @@
         pytest.param(
             1,
             "standard",
-            EnsemblerStandardConfig(
-                experiment_mappings=[
-                    turing.generated.models.EnsemblerStandardConfigExperimentMappings(
-                        experiment="experiment-1",
-                        treatment="treatment-1",
-                        route="route-1"
-                    ),
-                    turing.generated.models.EnsemblerStandardConfigExperimentMappings(
-                        experiment="experiment-2",
-                        treatment="treatment-2",
-                        route="route-2"
-                    )
-                ],
-                fallback_response_route_id="route-1"
-            ),
+            "standard_router_ensembler_config",
             None,
             "generic_standard_router_ensembler_config"
         ),
         pytest.param(
             1,
             "docker",
             None,
-            turing.generated.models.EnsemblerDockerConfig(
-                image="test.io/just-a-test/turing-ensembler:0.0.0-build.0",
-                resource_request=turing.generated.models.ResourceRequest(
-                    min_replica=1,
-                    max_replica=3,
-                    cpu_request='100m',
-                    memory_request='512Mi'
-                ),
-                endpoint=f"http://localhost:5000/ensembler_endpoint",
-                timeout="500ms",
-                port=5120,
-                env=[
-                    turing.generated.models.EnvVar(
-                        name="env_name",
-                        value="env_val")
-                ],
-                service_account="secret-name-for-google-service-account"
-            ),
+            "generic_ensembler_docker_config",
             "generic_docker_router_ensembler_config"
         )
     ])
 def test_create_router_ensembler_config(id, type, standard_config, docker_config, expected, request):
+    docker_config_data = None if docker_config is None else request.getfixturevalue(docker_config)
+    standard_config_data = None if standard_config is None else request.getfixturevalue(standard_config)
     actual = RouterEnsemblerConfig(
         id=id,
         type=type,
-        standard_config=standard_config,
-        docker_config=docker_config
+        standard_config=standard_config_data,
+        docker_config=docker_config_data
     ).to_open_api()
     assert actual == request.getfixturevalue(expected)
 
@@ -532,3 +502,88 @@ def test_create_nop_router_ensembler_config_with_invalid_route(
     router.ensembler = ensembler_config
     with pytest.raises(expected):
        router.to_open_api()
+
+@pytest.mark.parametrize(
+    "cls,config,expected", [
+        pytest.param(
+            NopRouterEnsemblerConfig,
+            "nop_router_ensembler_config",
+            {"type":"nop", "final_response_route_id": "test"}
+        ),
+        pytest.param(
+            StandardRouterEnsemblerConfig,
+            "standard_router_ensembler_config",
+            {
+                "type":"standard",
+                "experiment_mappings": [
+                    {
+                        "experiment": "experiment-1",
+                        "treatment": "treatment-1",
+                        "route": "route-1"
+                    },
+                    {
+                        "experiment": "experiment-2",
+                        "treatment": "treatment-2",
+                        "route": "route-2"
+                    }
+                ],
+                "fallback_response_route_id": "route-1"
+            }
+        ),
+        pytest.param(
+            DockerRouterEnsemblerConfig,
+            "generic_ensembler_docker_config",
+            {
+                "type": "docker",
+                "image": "test.io/just-a-test/turing-ensembler:0.0.0-build.0",
+                "resource_request": ResourceRequest(
+                    min_replica=1,
+                    max_replica=3,
+                    cpu_request='100m',
+                    memory_request='512Mi'
+                ),
+                "endpoint": "http://localhost:5000/ensembler_endpoint",
+                "timeout": "500ms",
+                "port": 5120,
+                "env": [EnvVar(name="env_name", value="env_val")],
+                "service_account": "secret-name-for-google-service-account"
+            }
+        ),
+        pytest.param(
+            PyfuncRouterEnsemblerConfig,
+            "generic_ensembler_pyfunc_config",
+            {
+                "type":"pyfunc",
+                "project_id": 77,
+                "ensembler_id": 11,
+                "resource_request": ResourceRequest(
+                    min_replica=1,
+                    max_replica=3,
+                    cpu_request='100m',
+                    memory_request='512Mi'
+                ),
+                "timeout": "500ms",
+                "env": [EnvVar(name="env_name", value="env_val")]
+            }
+        )
+    ])
+def test_create_ensembler_config_from_config(
+    cls, config, expected, request
+):
+    config_data = request.getfixturevalue(config)
+    assert cls.from_config(config_data).to_dict() == expected
+
+def test_set_nop_ensembler_config_with_default_route(request):
+    router = request.getfixturevalue("generic_router_config")
+    router.ensembler = None
+    router.default_route_id = "model-b"
+    assert router.ensembler.final_response_route_id == "model-b"
+
+def test_set_standard_ensembler_config_with_default_route(request):
+    router = request.getfixturevalue("generic_router_config")
+    router.ensembler = StandardRouterEnsemblerConfig(
+        experiment_mappings=[],
+        fallback_response_route_id=""
+    )
+    router.default_route_id = "model-b"
+    assert router.ensembler.fallback_response_route_id == "model-b"