Repository: caikit/caikit (python)

Issue:
Support multiple remote servers in `RemoteModelFinder`
## Is your feature request related to a problem? Please describe.

The `RemoteModelFinder` allows a single `caikit.runtime` to proxy requests to a single remote `caikit.runtime` server for the collection of models loaded there. This feature request is to extend that functionality to support multiple remote servers behind a single instance of `RemoteModelFinder`.

## Describe the solution you'd like

* Allow the `connection.hostname` key in the config to include a template that can be filled in with the `model_name`/`model_id`
    * This allows the set of connections to change dynamically without rolling over the server to reload config
* Add support for a `remote_connections` config to pre-populate connections to multiple known servers.

## Describe alternatives you've considered

The desired feature can be accomplished with the current implementation by having one configured `RemoteModelFinder` for each remote server, however this set of servers would be immutable after boot, so the server acting as the proxy would need to be rolled over to update the mapping.


Pull request:
Multi connection remote model finder 680
**What this PR does / why we need it**:

Closes #680 

This PR adds support to the `caikit.runtime.client.RemoteModelFinder` so that it can find models across multiple remote `caikit.runtime` servers. There are two new ways to enable multi-connection finding:

1. The `config.connection.hostname` string can contain a single `{}` template that will be replaced with `model_id`.
2. The `config.remote_connections` entry can be a list of connection blobs, each one with the same set of arguments as `config.connection` (`hostname`, `port`, `tls`, `options`).

Here's an example of a full configuration with `RemoteModelFinder`, `LocalFinder`, and [TGISAutoFinder](https://github.com/caikit/caikit-nlp/blob/main/caikit_nlp/model_management/tgis_auto_finder.py):

```yaml
model_management:
  initializers:
    default:
      type: LOCAL
      config:
        backend_priority:
          - type: LOCAL
          - type: TGIS
            config:
              connection:
                hostname: tgis.server:8080
  finders:
    default:
      type: LOCAL
      config:
        load_path: null
    remote:
      type: REMOTE
      config:
        connection:
          hostname: foo.bar.{}
          port: 1234
        remote_connection:
          - hostname: bar.baz
            port: 54321
          - hostname: biz.buz
            port: 12121
    tgis:
      type: TGIS-AUTO
      config:
        local_initializer_name: default
        tgis_backend_priority: 1
```

**Special notes for your reviewer**:

PR still needs more unit tests!

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility


Patch:
diff --git a/caikit/runtime/client/remote_model_finder.py b/caikit/runtime/client/remote_model_finder.py
index 461c10bb..4e749b27 100644
--- a/caikit/runtime/client/remote_model_finder.py
+++ b/caikit/runtime/client/remote_model_finder.py
@@ -35,9 +35,10 @@ model_management:
 
 """
 # Standard
+from dataclasses import dataclass
 from datetime import datetime, timedelta
 from threading import Lock
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 # Third Party
 from requests import RequestException
@@ -72,6 +73,12 @@ error = error_handler.get(log)
 ### Finder Definitions
 
 
+@dataclass
+class ModuleConnectionInfo:
+    conn: ConnectionInfo
+    module_id: str
+
+
 class RemoteModelFinder(ModelFinderBase):
     __doc__ = __doc__
 
@@ -82,9 +89,22 @@ class RemoteModelFinder(ModelFinderBase):
 
         self._instance_name = instance_name
 
-        # Type/Value check connection parameters
-        self._connection = ConnectionInfo(**config.connection)
-        self._tls = self._connection.tls
+        # Initialize model_name -> connection map
+        self._connections: Dict[str, ConnectionInfo] = {}
+        self._connection_template: Optional[ConnectionInfo] = None
+
+        # If a remote_models key is found, it's a mapping from model name to
+        # connection info
+        for remote_conn in config.get("remote_connections", []):
+            conn = ConnectionInfo(**remote_conn)
+            self._connections[f"{conn.hostname}:{conn.port}"] = conn
+
+        # If a single "global" connection given, initialize with model_name None
+        if config.connection:
+            default_conn = ConnectionInfo(**config.connection)
+            if f"{default_conn.hostname}:{default_conn.port}" not in self._connections:
+                self._connection_template = default_conn
+                self._connections[default_conn.hostname] = default_conn
 
         # Type/Value check default parameters
         self._model_key = config.get("model_key", MODEL_MESH_MODEL_ID_KEY)
@@ -98,17 +118,31 @@ class RemoteModelFinder(ModelFinderBase):
             self._protocol,
         )
 
-        if self._protocol == "grpc" and self._tls.enabled:
-            error.value_check(
-                "<COR74451567E>",
-                not self._tls.insecure_verify,
-                "GRPC does not support insecure TLS connections."
-                "Please provide a valid CA certificate",
-            )
+        if self._protocol == "grpc":
+            for conn in self._connections.values():
+                error.value_check(
+                    "<COR74451567E>",
+                    not conn.tls.enabled or not conn.tls.insecure_verify,
+                    "GRPC does not support insecure TLS connections."
+                    "Please provide a valid CA certificate",
+                )
+
+        # Initialize the supported models using the model connection info
+        self._supported_models: Dict[str, ModuleConnectionInfo] = {}
+        supported_models = config.get("supported_models") or {}
+        error.value_check(
+            "<NLP77334255E>",
+            not supported_models or self._connection_template,
+            "Cannot provide 'supported_models' without 'connection'",
+        )
+        for model_name, module_id in supported_models.items():
+            if model_conn := self._render_conn_template(model_name):
+                self._supported_models[model_name] = ModuleConnectionInfo(
+                    model_conn, module_id
+                )
 
         # Type/Value check model parameters
         self._discover_models = config.get("discover_models", True)
-        self._supported_models = config.get("supported_models") or {}
         self._min_poll_time = config.get("min_poll_time", 30)
         error.type_check(
             "<COR74343245E>",
@@ -140,7 +174,7 @@ class RemoteModelFinder(ModelFinderBase):
         # If model_path is not detected and discover models is enabled attempt
         # rediscovery
         if model_path not in self._supported_models and self._discover_models:
-            self._safe_discover()
+            self._safe_discover(model_path)
 
         # If model_path is not one of the supported models then raise an error
         if model_path not in self._supported_models:
@@ -151,9 +185,10 @@ class RemoteModelFinder(ModelFinderBase):
             )
             return
 
+        module_conn_info = self._supported_models.get(model_path)
         return RemoteModuleConfig.load_from_module(
-            module_reference=self._supported_models.get(model_path),
-            connection_info=self._connection,
+            module_reference=module_conn_info.module_id,
+            connection_info=module_conn_info.conn,
             protocol=self._protocol,
             model_key=self._model_key,
             model_path=model_path,
@@ -161,7 +196,9 @@ class RemoteModelFinder(ModelFinderBase):
 
     ### Discovery Helper Functions
 
-    def _discover(self) -> Dict[str, str]:
+    def _discover(
+        self, model_name: Optional[str] = None
+    ) -> Dict[str, ModuleConnectionInfo]:
         """Helper method to discover models from a remote
         runtime. This is a separate function to help with subclassing
 
@@ -169,13 +206,19 @@ class RemoteModelFinder(ModelFinderBase):
             model_map: Dict[str, str]
                 The map of models to modules
         """
+        error.value_check(
+            "<COR45835854E>",
+            self._protocol in ["grpc", "http"],
+            "Invalid protocol: {}",
+            self._protocol,
+        )
         if self._protocol == "grpc":
-            return self._discover_grpc_models()
-        elif self._protocol == "http":
-            return self._discover_http_models()
-        error("<COR45835854E>", ValueError(f"Invalid protocol: {self._protocol}"))
+            return self._discover_grpc_models(model_name)
+        return self._discover_http_models(model_name)
 
-    def _safe_discover(self) -> Dict[str, str]:
+    def _safe_discover(
+        self, model_name: Optional[str] = None
+    ) -> Dict[str, ModuleConnectionInfo]:
         """Helper function that lazily discovers models in a
         thread safe manor. This function also ensures we don't overload
         the remote server with discovery requests
@@ -195,10 +238,13 @@ class RemoteModelFinder(ModelFinderBase):
 
             # Run discovery
             self._last_discovered_time = current_time
-            self._supported_models = self._discover()
+            self._supported_models = self._discover(model_name)
             return self._supported_models
 
-    def _discover_grpc_models(self) -> Dict[str, str]:
+    def _discover_grpc_models(
+        self,
+        model_name: Optional[str],
+    ) -> Dict[str, ModuleConnectionInfo]:
         """Helper function to get all the supported models and modules
         from a remote GRPC runtime
 
@@ -206,45 +252,52 @@ class RemoteModelFinder(ModelFinderBase):
            support_models: Dict[str, str
                Mapping of remote model names to module ids
         """
-
-        target = f"{self._connection.hostname}:{self._connection.port}"
-        options = tuple(self._connection._options.items())
-        with construct_grpc_channel(target, options, self._tls) as channel:
-            info_service_rpc = channel.unary_unary(
-                get_grpc_route_name(ServiceType.INFO, "GetModelsInfo"),
-                request_serializer=ModelInfoRequest.get_proto_class().SerializeToString,
-                response_deserializer=ModelInfoResponse.get_proto_class().FromString,
-            )
-            try:
-                model_info_proto = info_service_rpc(
-                    ModelInfoRequest().to_proto(), timeout=self._connection.timeout
-                )
-            except grpc.RpcError as exc:
-                log.warning(
-                    "Unable to discover modules from remote: %s. Error: %s",
-                    self._connection.hostname,
-                    str(exc),
-                )
-                return {}
-
-        model_info_response = ModelInfoResponse.from_proto(model_info_proto)
-
-        # Parse response into dictionary of name->id
         supported_modules = {}
-        for model_info in model_info_response.models:
-            model_name = model_info.name
-            module_id = model_info.module_id
-
-            log.debug(
-                "Discovered model %s with module_id %s from remote runtime",
-                model_name,
-                module_id,
-            )
-            supported_modules[model_name] = module_id
+        for conn in self._get_conn_candidates(model_name):
+            target = f"{conn.hostname}:{conn.port}"
+            options = [tuple(opt) for opt in conn.options.items()]
+            with construct_grpc_channel(target, options, conn.tls) as channel:
+                info_service_rpc = channel.unary_unary(
+                    get_grpc_route_name(ServiceType.INFO, "GetModelsInfo"),
+                    request_serializer=ModelInfoRequest.get_proto_class().SerializeToString,
+                    response_deserializer=ModelInfoResponse.get_proto_class().FromString,
+                )
+                try:
+                    model_info_proto = info_service_rpc(
+                        ModelInfoRequest().to_proto(), timeout=conn.timeout
+                    )
+
+                    model_info_response = ModelInfoResponse.from_proto(model_info_proto)
+
+                    # Parse response into dictionary of name->conn
+                    for model_info in model_info_response.models:
+                        model_name = model_info.name
+                        module_id = model_info.module_id
+
+                        log.debug(
+                            "Discovered model %s with module_id %s from remote runtime %s",
+                            model_name,
+                            module_id,
+                            target,
+                        )
+                        # NOTE: If multiple servers support the same model, the
+                        #   first to be checked will win
+                        supported_modules.setdefault(
+                            model_name, ModuleConnectionInfo(conn, module_id)
+                        )
+                except grpc.RpcError as exc:
+                    log.warning(
+                        "Unable to discover modules from remote: %s. Error: %s",
+                        target,
+                        str(exc),
+                    )
 
         return supported_modules
 
-    def _discover_http_models(self) -> Dict[str, str]:
+    def _discover_http_models(
+        self,
+        model_name: Optional[str],
+    ) -> Dict[str, ConnectionInfo]:
         """Helper function to get all the supported models and modules
         from a remote HTTP runtime
 
@@ -252,53 +305,75 @@ class RemoteModelFinder(ModelFinderBase):
             supported_models:Dict[str, str]
                 Mapping of remote model names to module_ids
         """
+        supported_modules = {}
+        for conn in self._get_conn_candidates(model_name):
 
-        # Configure HTTP target and Session object
-        http_scheme = "https://" if self._tls.enabled else "http://"
-        target = (
-            f"{http_scheme}{self._connection.hostname}:"
-            f"{self._connection.port}{MODELS_INFO_ENDPOINT}"
-        )
-        session = construct_requests_session(
-            self._connection.options, self._tls, self._connection.timeout
-        )
-
-        # Send HTTP Request
-        try:
-            resp = session.get(target)
-        except RequestException as exc:
-            log.warning(
-                "Unable to discover modules from remote: %s. Error: %s",
-                self._connection.hostname,
-                str(exc),
-            )
-            return {}
-
-        if resp.status_code != 200:
-            log.warning(
-                "Unable to discover modules from remote: %s. Error: %s",
-                target,
-                resp.reason,
+            # Configure HTTP target and Session object
+            http_scheme = "https://" if conn.tls.enabled else "http://"
+            target = (
+                f"{http_scheme}{conn.hostname}:" f"{conn.port}{MODELS_INFO_ENDPOINT}"
             )
-            return {}
+            session = construct_requests_session(conn.options, conn.tls, conn.timeout)
 
-        # Load the response as a json object
-        model_info = resp.json()
+            # Send HTTP Request
+            try:
+                resp = session.get(target)
+
+                if resp.status_code != 200:
+                    log.warning(
+                        "Unable to discover modules from remote: %s. Error: %s",
+                        target,
+                        resp.reason,
+                    )
+                else:
+
+                    # Load the response as a json object
+                    model_info = resp.json()
+
+                    # Parse response into dictionary of name->id
+                    for model_dict in model_info.get("models", []):
+                        model_name = model_dict.get("name")
+                        module_id = model_dict.get("module_id")
+
+                        log.debug(
+                            "Discovered model %s with module_id %s from remote runtime",
+                            model_name,
+                            module_id,
+                        )
+                        # NOTE: If multiple servers support the same model, the
+                        #   first to be checked will win
+                        supported_modules.setdefault(
+                            model_name, ModuleConnectionInfo(conn, module_id)
+                        )
+            except RequestException as exc:
+                log.warning(
+                    "Unable to discover modules from remote: %s. Error: %s",
+                    target,
+                    str(exc),
+                )
 
-        # Parse response into dictionary of name->id
-        supported_modules = {}
-        for model_dict in model_info.get("models", []):
-            model_name = model_dict.get("name")
-            module_id = model_dict.get("module_id")
+        return supported_modules
 
-            log.debug(
-                "Discovered model %s with module_id %s from remote runtime",
-                model_name,
-                module_id,
+    def _render_conn_template(self, model_name: str) -> Optional[ConnectionInfo]:
+        """Common utility to get the connection for a given model"""
+        if self._connection_template is not None:
+            conn_dict = self._connection_template.to_dict()
+            conn_dict["hostname"] = self._connection_template.hostname.format(
+                model_name
             )
-            supported_modules[model_name] = module_id
-
-        return supported_modules
+            return ConnectionInfo(**conn_dict)
+
+    def _get_conn_candidates(self, model_name: Optional[str]) -> List[ConnectionInfo]:
+        """Common utility to get all connections to try"""
+        candidate_conns = []
+        if (
+            model_name is not None
+            and self._connection_template is not None
+            and (model_conn := self._render_conn_template(model_name))
+        ):
+            candidate_conns.append(model_conn)
+        candidate_conns.extend(self._connections.values())
+        return candidate_conns
 
 
 # Register the remote finder once it has been constructed
diff --git a/pyproject.toml b/pyproject.toml
index cef7c09e..079e5c0e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -99,7 +99,7 @@ dev-docs = [
 ]
 
 dev-fmt = [
-    "ruff==0.3.2",
+    "ruff==0.3.1",
     "pre-commit>=3.0.4,<4.0",
     "pydeps>=1.12.12,<2",
 ]


Test patch:
diff --git a/tests/runtime/client/test_remote_model_finder.py b/tests/runtime/client/test_remote_model_finder.py
index 8873d621..73447ac1 100644
--- a/tests/runtime/client/test_remote_model_finder.py
+++ b/tests/runtime/client/test_remote_model_finder.py
@@ -17,23 +17,31 @@ Tests for the RemoteModelFinder
 
 # Standard
 from contextlib import contextmanager
+from typing import Optional
+from unittest.mock import MagicMock, patch
 
 # Third Party
+import grpc
 import pytest
 
 # First Party
 from aconfig import Config, ImmutableConfig
 
 # Local
+from caikit.interfaces.runtime.data_model import ModelInfo, ModelInfoResponse
 from caikit.runtime.client import RemoteModelFinder, RemoteModuleConfig
 from caikit.runtime.model_management.model_manager import ModelManager
 from sample_lib.modules.file_processing import BoundingBoxModule
+from sample_lib.modules.sample_task import SampleModule
 from tests.conftest import random_test_id
 from tests.fixtures import Fixtures
 from tests.runtime.conftest import multi_task_model_id  # noqa: F401
-from tests.runtime.conftest import open_port  # noqa: F401
 from tests.runtime.conftest import sample_task_model_id  # noqa: F401
-from tests.runtime.conftest import generate_tls_configs, runtime_test_server
+from tests.runtime.conftest import (  # noqa: F401
+    generate_tls_configs,
+    open_port,
+    runtime_test_server,
+)
 
 ## Test Helpers #######################################################################
 
@@ -75,6 +83,7 @@ def temp_finder(
     multi_finder_name="remote",
     multi_finder_cfg=None,
     connection_cfg=None,
+    remote_connections_cfg=None,
     min_poll_time=0,
     protocol="grpc",
 ):
@@ -88,15 +97,17 @@ def temp_finder(
 
     if connection_cfg:
         multi_finder_cfg["connection"] = connection_cfg
-
-    if "protocol" not in multi_finder_cfg:
-        multi_finder_cfg["protocol"] = protocol
-
-    if "connection" not in multi_finder_cfg:
+    elif connection_cfg is None:
         multi_finder_cfg["connection"] = {
             "hostname": "localhost",
         }
 
+    if remote_connections_cfg:
+        multi_finder_cfg["remote_connections"] = remote_connections_cfg
+
+    if "protocol" not in multi_finder_cfg:
+        multi_finder_cfg["protocol"] = protocol
+
     yield RemoteModelFinder(ImmutableConfig(multi_finder_cfg), multi_finder_name)
 
 
@@ -121,6 +132,27 @@ def test_remote_finder_static_model(sample_module_id):
         assert len(config.task_methods[0][1]) == 4
 
 
+def test_remote_finder_connection_template(sample_module_id):
+    """Test to ensure that the connection can be a template"""
+    hn_template = "foo.{}.svc"
+    with temp_finder(
+        connection_cfg={
+            "hostname": hn_template,
+            "port": 12345,
+        },
+        multi_finder_cfg={
+            "discover_models": False,
+            "supported_models": {
+                "sample1": sample_module_id,
+                "sample2": sample_module_id,
+            },
+        },
+    ) as finder:
+        for model_id in ["sample1", "sample2"]:
+            config = finder.find_model(model_id)
+            assert config.connection.hostname == hn_template.format(model_id)
+
+
 @pytest.mark.parametrize("protocol", ["grpc", "http"])
 def test_remote_finder_multi_task_model(multi_task_model_id, open_port, protocol):
     """Test to ensure model finder works for models with multiple tasks"""
@@ -138,7 +170,9 @@ def test_remote_finder_multi_task_model(multi_task_model_id, open_port, protocol
 
 
 @pytest.mark.parametrize("protocol", ["grpc", "http"])
-def test_remote_finder_discover_models(sample_task_model_id, open_port, protocol):
+def test_remote_finder_discover_single_conn_models(
+    sample_task_model_id, open_port, protocol
+):
     """Test to ensure discovering models works for http"""
     with runtime_test_server(open_port, protocol=protocol) as server, temp_finder(
         connection_cfg={
@@ -155,6 +189,107 @@ def test_remote_finder_discover_models(sample_task_model_id, open_port, protocol
         assert len(config.task_methods[0][1]) == 4
 
 
+@pytest.mark.parametrize("protocol", ["grpc", "http"])
+def test_remote_finder_discover_multi_conn_models(protocol):
+    """Test to ensure discovery works with multiple servers"""
+    hn_a = "foo.bar.com"
+    hn_b = "baz.biz.com"
+    port1 = 12345
+    port2 = 23456
+    mod_id_x = SampleModule.MODULE_ID
+    mod_id_y = BoundingBoxModule.MODULE_ID
+    model_id1 = random_test_id()
+    model_id2 = random_test_id()
+    model_id3 = random_test_id()
+    model_id4 = random_test_id()
+
+    class MockChannelSession:
+        def __init__(self, *_, target: Optional[str] = None, **__):
+            self.target = target
+
+        @staticmethod
+        def _get_resp(target: str):
+            return {
+                # hn A / port1 -> model1, model2
+                f"{hn_a}:{port1}": ModelInfoResponse(
+                    [
+                        ModelInfo(name=model_id1, module_id=mod_id_x),
+                        ModelInfo(name=model_id2, module_id=mod_id_x),
+                    ]
+                ),
+                # hn A / port2 -> model3
+                f"{hn_a}:{port2}": ModelInfoResponse(
+                    [
+                        ModelInfo(name=model_id3, module_id=mod_id_y),
+                    ]
+                ),
+                # hn B / port3 -> model4
+                f"{hn_b}:{port2}": ModelInfoResponse(
+                    [
+                        ModelInfo(name=model_id4, module_id=mod_id_y),
+                    ]
+                ),
+            }.get(target)
+
+        def get(self, target: str):
+            resp_mock = MagicMock()
+            resp = self._get_resp(target.split("/")[2])
+            if not resp:
+                resp_mock.status_code = 404
+            else:
+                resp_mock.status_code = 200
+                resp_mock.json = MagicMock(return_value=resp.to_dict())
+            return resp_mock
+
+        def unary_unary(self, *_, **__):
+            assert self.target
+            resp = self._get_resp(self.target)
+            if not resp:
+                return MagicMock(side_effect=grpc.RpcError)
+            return MagicMock(return_value=resp.to_proto())
+
+    @contextmanager
+    def mock_construct_grpc_channel(target, *_, **__):
+        yield MockChannelSession(target=target)
+
+    with patch(
+        "caikit.runtime.client.remote_model_finder.construct_grpc_channel",
+        new=mock_construct_grpc_channel,
+    ), patch(
+        "caikit.runtime.client.remote_model_finder.construct_requests_session",
+        new=MockChannelSession,
+    ):
+        with temp_finder(
+            remote_connections_cfg=[
+                {"hostname": hn_a, "port": port1},
+                {"hostname": hn_a, "port": port2},
+                {"hostname": hn_b, "port": port2},
+            ],
+            protocol=protocol,
+        ) as finder:
+            # hn A / port1 -> model1, model2
+            config1 = finder.find_model(model_id1)
+            assert config1
+            assert config1.connection.hostname == hn_a
+            assert config1.connection.port == port1
+            config2 = finder.find_model(model_id2)
+            assert config2
+            assert config2.connection.hostname == hn_a
+            assert config2.connection.port == port1
+            # hn A / port2 -> model3
+            config3 = finder.find_model(model_id3)
+            assert config3
+            assert config3.connection.hostname == hn_a
+            assert config3.connection.port == port2
+            # hn B / port3 -> model4
+            config4 = finder.find_model(model_id4)
+            assert config4
+            assert config4.connection.hostname == hn_b
+            assert config4.connection.port == port2
+            # Unknown model
+            assert finder.find_model("unknown") is None
+
+
 @pytest.mark.parametrize("protocol", ["grpc", "http"])
 def test_remote_finder_discover_mtls_models(sample_task_model_id, open_port, protocol):
     """Test to ensure discovering models works for https with MTLS and secure CA"""