Add TGISBackend.register_model_connection() method

caikit_tgis_backend/tgis_backend.py

@@ -11,12 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""This module implements a TGIS backend configuration
-"""
+"""This module implements a TGIS backend configuration"""
 
 # Standard
 from threading import Lock
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 # Third Party
 import grpc
@@ -35,6 +34,7 @@
 log = alog.use_channel("TGISBKND")
 error = error_handler.get(log)
 
+
 # pylint: disable=too-many-instance-attributes
 class TGISBackend(BackendBase):
     """Caikit backend with a connection to the TGIS server. If no connection
@@ -67,15 +67,17 @@ def __init__(self, config: Optional[dict] = None):
         self._mutex = Lock()
         self._local_tgis = None
         self._managed_tgis = None
-        self._model_connections = {}
+        self._model_connections: Dict[str, TGISConnection] = {}
         self._test_connections = self.config.get("test_connections", False)
         self._connect_timeout = self.config.get("connect_timeout", None)
 
         # Parse the config to see if we're managing a connection to a remote
         # TGIS instance or running a local copy
         connection_cfg = self.config.get("connection") or {}
         error.type_check("<TGB20235229E>", dict, connection=connection_cfg)
-        self._remote_models_cfg = self.config.get("remote_models") or {}
+        self._remote_models_cfg: Dict[str, dict] = (
+            self.config.get("remote_models") or {}
+        )
         error.type_check("<TGB20235338E>", dict, connection=self._remote_models_cfg)
         local_cfg = self.config.get("local") or {}
         error.type_check("<TGB20235225E>", dict, local=local_cfg)
@@ -195,6 +197,33 @@ def get_connection(
                 self._model_connections.setdefault(model_id, model_conn)
         return model_conn
 
+    def register_model_connection(
+        self, model_id: str, conn_cfg: Dict[str, Any]
+    ) -> None:
+        # TODO: Is this method necessary? .get_client() -> .get_connection() which will create
+        # the connection in self._model_connections if it doesn't exist by default.
+        model_conn = TGISConnection.from_config(model_id, conn_cfg)
+        error.value_check("<TGB81270235E>", model_conn is not None)
+
+        if self._test_connections:
+            try:
+                model_conn.test_connection()
+            except grpc.RpcError as err:
+                log.warning(
+                    "<TGB10601575W>",
+                    "Unable to connect to model %s: %s",
+                    model_id,
+                    err,
+                    exc_info=True,
+                )
+                model_conn = None
+        if model_conn is not None:
+            # NOTE: setdefault used here to avoid the need to hold the mutex
+            #   when running the connection test. It's possible that two
+            #   threads would stimulate the creation of the connection
+            #   concurrently, so just keep whichever dict update lands first
+            self._model_connections.setdefault(model_id, model_conn)
+
     def get_client(self, model_id: str) -> generation_pb2_grpc.GenerationServiceStub:
         model_conn = self.get_connection(model_id)
         if model_conn is None and self.local_tgis:

caikit_tgis_backend/tgis_connection.py

@@ -81,6 +81,8 @@ class TGISConnection:
 
     HOSTNAME_KEY = "hostname"
     HOSTNAME_TEMPLATE_MODEL_ID = "model_id"
+    NAMESPACE_KEY = "namespace"
+    HOSTNAME_TEMPLATE_NAMESPACE = "namespace"
     CA_CERT_FILE_KEY = "ca_cert_file"
     CLIENT_CERT_FILE_KEY = "client_cert_file"
     CLIENT_KEY_FILE_KEY = "client_key_file"
@@ -94,12 +96,22 @@ def from_config(cls, model_id: str, config: dict) -> Optional["TGISConnection"]:
         """Create an instance from a connection template and a model_id"""
         hostname = config.get(cls.HOSTNAME_KEY)
         if hostname:
-            hostname = hostname.format(
-                **{
+            assert isinstance(hostname, str)
+            # TODO: Namespace needs to be populated here as well into the hostname template
+            namespace = config.get(cls.NAMESPACE_KEY)
+
+            hostname = hostname.format_map(
+                {
                     cls.HOSTNAME_TEMPLATE_MODEL_ID: model_id,
+                    cls.HOSTNAME_TEMPLATE_NAMESPACE: namespace,
                 }
             )
-            log.debug("Resolved hostname [%s] for model %s", hostname, model_id)
+            log.debug(
+                "Resolved hostname [%s] for model %s in namespace %s",
+                hostname,
+                model_id,
+                namespace,
+            )
 
             tls_hostname_override = config.get(cls.TLS_HN_OVERRIDE_KEY)
 

tests/test_tgis_connection.py

@@ -14,6 +14,7 @@
 """
 Unit tests for the TGISConnection class
 """
+
 # Standard
 from contextlib import contextmanager
 from pathlib import Path
@@ -27,7 +28,7 @@
 
 # Local
 from caikit_tgis_backend.tgis_connection import TGISConnection
-from tests.tgis_mock import tgis_mock_insecure
+from tests.tgis_mock import tgis_mock_insecure  # noqa
 
 
 @contextmanager
@@ -78,6 +79,30 @@ def test_happy_path_template():
     )
 
 
+def test_happy_path_template_with_namespace():
+    template_model_id = "{{{}}}".format(TGISConnection.HOSTNAME_TEMPLATE_MODEL_ID)
+    template_namespace = "{{{}}}".format(TGISConnection.HOSTNAME_TEMPLATE_NAMESPACE)
+    template = (
+        f"{template_model_id}-predictor.'{template_namespace}'.svc.cluster.local:8033"
+    )
+
+    model_id = "some/model"
+    namespace = "byom4"
+    conn = TGISConnection.from_config(
+        model_id,
+        {
+            TGISConnection.HOSTNAME_KEY: template,
+            TGISConnection.NAMESPACE_KEY: namespace,
+        },
+    )
+    assert conn.hostname == template.format_map(
+        {
+            TGISConnection.HOSTNAME_TEMPLATE_MODEL_ID: model_id,
+            TGISConnection.HOSTNAME_TEMPLATE_NAMESPACE: namespace,
+        }
+    )
+
+
 def test_happy_path_tls(temp_ca_cert):
     conn = TGISConnection.from_config(
         "",