Merge pull request caikit#123 from gabe-l-hart/TgisAutoFinder

Tgis auto finder
Signed-off-by: gkumbhat <[email protected]>

caikit_nlp/__init__.py

@@ -24,7 +24,7 @@
 
 # Local
 # Import subpackages
-from . import config, data_model
+from . import config, data_model, model_management
 from .config import *
 from .data_model import *
 from .modules import *

caikit_nlp/model_management/__init__.py

@@ -0,0 +1,16 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+# Local
+from .tgis_auto_finder import TGISAutoFinder

caikit_nlp/model_management/tgis_auto_finder.py

@@ -0,0 +1,142 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""
+The TGISAutoFinder implements the ModelFinder interface to provide automatic
+discovery of text-generation models that can be auto-configured to run against
+a remote TGIS model.
+"""
+# Standard
+from typing import Optional
+
+# First Party
+from caikit.core import MODEL_MANAGER, error_handler
+from caikit.core.model_management import ModelFinderBase, model_finder_factory
+from caikit.core.modules import ModuleConfig
+from caikit_tgis_backend import TGISBackend
+import aconfig
+import alog
+
+# Local
+from ..modules.text_generation import TextGenerationTGIS
+
+log = alog.use_channel("TGIS_FND")
+error = error_handler.get(log)
+
+
+class TGISAutoFinder(ModelFinderBase):
+    __doc__ = __doc__
+
+    name = "TGIS-AUTO"
+
+    # Constants for the keys of the config blob
+    _LOCAL_INITIALIZER_NAME_KEY = "local_initializer_name"
+    _TGIS_BACKEND_PRIORITY_KEY = "tgis_backend_priority"
+
+    def __init__(self, config: aconfig.Config, instance_name: str = ""):
+        """Initialize from the model finder factory config
+
+        Config schema:
+
+        local_initializer_name:
+            type: string
+            default: "default"
+            description: The name within the initializers config for the LOCAL
+                initializer that will hold the tgis backend to use
+
+        tgis_backend_priority:
+            type: integer
+            description: Index within the backend_priority list for the TGIS
+                backend to use. If not set, the first TGIS backend found will be
+                used.
+
+        Args:
+            config (aconfig.Config): The configuration blob from caikit's
+                model_management factory construction
+            instance_name (str): The name of this finder instance
+        """
+        local_initializer_name = config.get(self._LOCAL_INITIALIZER_NAME_KEY, "default")
+        tgis_backend_priority = config.get(self._TGIS_BACKEND_PRIORITY_KEY)
+        error.type_check(
+            "<NLP97312902E>", str, local_initializer_name=local_initializer_name
+        )
+        error.type_check(
+            "<NLP97312903E>",
+            int,
+            tgis_backend_priority=tgis_backend_priority,
+            allow_none=True,
+        )
+
+        # Extract the TGIS backend instance
+        local_initializer = MODEL_MANAGER.get_initializer(local_initializer_name)
+        backends = local_initializer.backends
+        if tgis_backend_priority is not None:
+            error.value_check(
+                "<NLP87928813E>",
+                0 <= tgis_backend_priority < len(backends),
+                "Invalid {}: {}",
+                self._TGIS_BACKEND_PRIORITY_KEY,
+                tgis_backend_priority,
+            )
+            self._tgis_backend = backends[tgis_backend_priority]
+            error.value_check(
+                "<NLP77150201E>",
+                self._tgis_backend.backend_type == TGISBackend.backend_type,
+                "Index {} is not a TGIS backend",
+                tgis_backend_priority,
+            )
+        else:
+            tgis_backend = None
+            for backend in backends:
+                if backend.backend_type == TGISBackend.backend_type:
+                    tgis_backend = backend
+                    break
+            error.value_check(
+                "<NLP96294266E>",
+                tgis_backend is not None,
+                "No TGIS backend found!",
+            )
+            self._tgis_backend = tgis_backend
+
+    def find_model(
+        self,
+        model_path: str,
+        **kwargs,
+    ) -> Optional[ModuleConfig]:
+        """Find the model if"""
+
+        # Get a connection to this model in tgis
+        log.debug2("Attempting to setup TGIS client for %s", model_path)
+        if self._tgis_backend.get_connection(model_id=model_path) is None:
+            log.debug2("TGIS cannot connect to model %s", model_path)
+            return None
+
+        # If connection is ok, set up the module config to point to the remote
+        # TGIS text generation module
+        cfg = ModuleConfig(
+            {
+                "module_id": TextGenerationTGIS.MODULE_ID,
+                "module_class": TextGenerationTGIS.MODULE_CLASS,
+                "name": TextGenerationTGIS.MODULE_NAME,
+                "version": TextGenerationTGIS.MODULE_VERSION,
+                "model_name": model_path,
+            }
+        )
+        # Set a special indicator in the module config to use the backend that
+        # this finder found. This will override the backend found by the local
+        # initializer.
+        cfg.tgis_backend = self._tgis_backend
+        return cfg
+
+
+model_finder_factory.register(TGISAutoFinder)

caikit_nlp/modules/text_generation/text_generation_tgis.py

@@ -148,6 +148,7 @@ def load(cls, model_path: str, load_backend: BackendBase) -> "TextGeneration":
         error.type_check("<NLP03521359E>", TGISBackend, load_backend=load_backend)
 
         config = ModuleConfig.load(model_path)
+        tgis_backend = config.tgis_backend or load_backend
         artifacts_path = config.artifact_path
         if artifacts_path:
             model_name = os.path.join(model_path, artifacts_path)
@@ -163,7 +164,7 @@ def load(cls, model_path: str, load_backend: BackendBase) -> "TextGeneration":
             sep_token=config.sep_token,
             eos_token=config.eos_token,
             pad_token=config.pad_token,
-            tgis_backend=load_backend,
+            tgis_backend=tgis_backend,
         )
 
     def save(self, model_path: str):

pyproject.toml

@@ -14,8 +14,8 @@ classifiers=[
     "License :: OSI Approved :: Apache Software License"
 ]
 dependencies = [
-    "caikit[runtime-grpc,runtime-http]>=0.13.0,<0.15.0",
-    "caikit-tgis-backend>=0.1.14,<0.2.0",
+    "caikit[runtime-grpc,runtime-http]>=0.15.0,<0.16.0",
+    "caikit-tgis-backend>=0.1.15,<0.2.0",
 
     # TODO: loosen dependencies
     "accelerate>=0.18.0",

runtime_config.yaml

@@ -9,6 +9,13 @@ runtime:
       size: 0 # Set to batch size for batching
 
 model_management:
+  finders:
+    default:
+      type: LOCAL
+    remote_tgis:
+      type: TGIS-AUTO
+      config:
+        test_connection: true
   initializers:
     default:
       type: LOCAL

tests/conftest.py

@@ -0,0 +1,29 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 sets up global test configs when pytest starts
+"""
+
+# Standard
+import os
+
+# First Party
+import alog
+
+# Configure logging from the environment
+alog.configure(
+    default_level=os.environ.get("LOG_LEVEL", "off"),
+    filters=os.environ.get("LOG_FILTERS", "urllib3:off"),
+    thread_id=os.environ.get("LOG_THREAD_ID", "") == "true",
+)

tests/fixtures/__init__.py

@@ -242,6 +242,7 @@ def __init__(
     def get_client(self, base_model_name):
         self._model_connections[base_model_name] = TGISConnection(
             hostname="foo.bar",
+            model_id=base_model_name,
             prompt_dir=self._temp_dir,
         )
         return StubTGISClient(base_model_name)