Repository: caikit/caikit (python)

Issue:
Lazy imports of model trainer modules causing issues in module using asyncio.Queue()
## Describe the bug

When attempting to integrate the Ray backend module, an error is thrown when the module is dynamically imported during a `train()` call from the Global Train Servicer. The Ray library that is used is intolerant of being imported dynamically from inside a gRPC thread pool, as opposed to a main thread. 

## Platform

Please provide details about the environment you are using, including the following:

- Interpreter version:
- Library version:

## Sample Code

Please include a minimal sample of the code that will (if possible) reproduce the bug in isolation

## Expected behavior

A clear and concise description of what you expected to happen.

## Observed behavior

File \"/app/.local/lib/python3.9/site-packages/caikit/runtime/servicers/global_train_servicer.py\", line 124, in Train
    return self.run_training_job(
  File \"/app/.local/lib/python3.9/site-packages/caikit/runtime/servicers/global_train_servicer.py\", line 210, in run_training_job
    model_future = MODEL_MANAGER.train(**kwargs)
  File \"/app/.local/lib/python3.9/site-packages/caikit/core/model_manager.py\", line 134, in train
    trainer = self._get_trainer(trainer)
  File \"/app/.local/lib/python3.9/site-packages/caikit/core/model_manager.py\", line 547, in _get_trainer
    return self._get_component(
  File \"/app/.local/lib/python3.9/site-packages/caikit/core/model_manager.py\", line 542, in _get_component
    component_dict[component] = component_factory.construct(cfg, component)
  File \"/app/.local/lib/python3.9/site-packages/caikit/core/model_management/factories.py\", line 58, in construct
    imported_module = importlib.import_module(module_name)
  File \"/usr/lib64/python3.9/importlib/__init__.py\", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File \"<frozen importlib._bootstrap>\", line 1030, in _gcd_import
  File \"<frozen importlib._bootstrap>\", line 1007, in _find_and_load
  File \"<frozen importlib._bootstrap>\", line 986, in _find_and_load_unlocked
  File \"<frozen importlib._bootstrap>\", line 680, in _load_unlocked
  File \"<frozen importlib._bootstrap_external>\", line 850, in exec_module
  File \"<frozen importlib._bootstrap>\", line 228, in _call_with_frames_removed
  File \"/app/.local/lib/python3.9/site-packages/caikit_ray_backend/blocks/ray_train.py\", line 22, in <module>
    from ray.job_submission import JobStatus
  File \"/app/.local/lib/python3.9/site-packages/ray/job_submission/__init__.py\", line 3, in <module>
    from ray.dashboard.modules.job.sdk import JobSubmissionClient
  File \"/app/.local/lib/python3.9/site-packages/ray/dashboard/modules/job/sdk.py\", line 6, in <module>
    from ray.dashboard.utils import get_address_for_submission_client
  File \"/app/.local/lib/python3.9/site-packages/ray/dashboard/utils.py\", line 277, in <module>
    class NotifyQueue:
  File \"/app/.local/lib/python3.9/site-packages/ray/dashboard/utils.py\", line 280, in NotifyQueue
    _queue = asyncio.Queue()
  File \"/usr/lib64/python3.9/asyncio/queues.py\", line 36, in __init__
    self._loop = events.get_event_loop()
  File \"/usr/lib64/python3.9/asyncio/events.py\", line 642, in get_event_loop
    raise RuntimeError('There is no current event loop in thread %r.'
RuntimeError: There is no current event loop in thread 'ThreadPoolExecutor-0_0'.

## Additional context

Add any other context about the problem here.


Pull request:
Proactive mm init 355
## Description

Closes #355 

This PR moves the initialization of `caikit.core.MODEL_MANAGER` components to happen proactively when a server (`grpc` or `http`) is initialized. This means that each component instance will be instantiated in the parent thread of the server rather than in a worker thread.

Patch:
diff --git a/caikit/core/model_manager.py b/caikit/core/model_manager.py
index d8735ce8..11abe25f 100644
--- a/caikit/core/model_manager.py
+++ b/caikit/core/model_manager.py
@@ -75,6 +75,19 @@ class ModelManager:
         self._initializers = {}
         self.__singleton_lock = Lock()
 
+    def initialize_components(self):
+        """Proactively initialize all configured trainer/finder/initializer
+        component instances. This is a separate call to enable explicit config.
+        """
+        # Initialize all configured components
+        mm_config = get_config().model_management
+        for trainer in mm_config.trainers:
+            self._get_trainer(trainer)
+        for finder in mm_config.finders:
+            self._get_finder(finder)
+        for initializer in mm_config.initializers:
+            self._get_initializer(initializer)
+
     ## Public ##################################################################
 
     def train(
diff --git a/caikit/runtime/utils/import_util.py b/caikit/runtime/utils/import_util.py
index a5a83725..864a9e6b 100644
--- a/caikit/runtime/utils/import_util.py
+++ b/caikit/runtime/utils/import_util.py
@@ -30,6 +30,7 @@ import alog
 
 # Local
 from caikit import get_config
+from caikit.core import MODEL_MANAGER
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 import caikit.core
 
@@ -86,8 +87,7 @@ def get_data_model(config: aconfig.Config = None) -> UnifiedDataModel:
     Returns:
         (module): Handle to the module after dynamic wild import
     """
-    if not config:
-        config = get_config()
+    config = config or get_config()
     lib_names = clean_lib_names(config.runtime.library)
 
     # Add all caikit.interfaces.X modules
@@ -119,6 +119,14 @@ def get_data_model(config: aconfig.Config = None) -> UnifiedDataModel:
     for lib_name in base_lib_names:
         cdm = _get_cdm_from_lib(lib_name, cdm)
 
+    # Ensure that all model management components have been initialized
+    # TODO: This function has a _ton_ of side effects! We need to split it up
+    #   and isolate these side effects in appropriately named functions.
+    #   Specifically, this function is not only responsible for creating the
+    #   data model, but it performs the dynamic import of the domain library and
+    #   initializes the model management components.
+    MODEL_MANAGER.initialize_components()
+
     return cdm
 
 


Test patch:
diff --git a/tests/core/test_model_manager.py b/tests/core/test_model_manager.py
index a8a7e441..dc1381e5 100644
--- a/tests/core/test_model_manager.py
+++ b/tests/core/test_model_manager.py
@@ -569,6 +569,37 @@ def test_load_loader_finder_by_value(reset_globals):
         assert isinstance(model, SampleModule)
 
 
+def test_initialize_all_components(reset_globals):
+    """Make sure that calling initialize_components will proactively create all
+    component instances from config
+    """
+    with temp_config(
+        {
+            "model_management": {
+                "trainers": {
+                    "default": {"type": "LOCAL"},
+                    "foobar": {"type": "LOCAL"},
+                },
+                "finders": {
+                    "default": {"type": "LOCAL"},
+                    "foobar": {"type": "LOCAL"},
+                },
+                "initializers": {
+                    "default": {"type": "LOCAL"},
+                    "foobar": {"type": "LOCAL"},
+                },
+            }
+        }
+    ):
+        assert not MODEL_MANAGER._trainers
+        assert not MODEL_MANAGER._finders
+        assert not MODEL_MANAGER._initializers
+        MODEL_MANAGER.initialize_components()
+        assert set(MODEL_MANAGER._trainers.keys()) == {"default", "foobar"}
+        assert set(MODEL_MANAGER._finders.keys()) == {"default", "foobar"}
+        assert set(MODEL_MANAGER._initializers.keys()) == {"default", "foobar"}
+
+
 def test_train_by_module_class(reset_globals):
     """Make sure training can be accessed through the central train function
     with the class directly
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 66ea83a6..0b0f4778 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -64,6 +64,7 @@ from sample_lib.data_model import (
     SampleTrainingType,
 )
 from tests.conftest import random_test_id
+from tests.core.helpers import *
 from tests.fixtures import Fixtures
 from tests.runtime.conftest import register_trained_model, runtime_grpc_test_server
 import caikit.interfaces.common
@@ -183,6 +184,25 @@ def test_model_train(runtime_grpc_server):
     assert result.learning_rate == 0.0015
 
 
+def test_components_preinitialized(
+    reset_globals, open_port, sample_inference_service, sample_train_service
+):
+    """Test that all model management components get pre-initialized when the
+    server is instantiated
+    """
+    assert not MODEL_MANAGER._trainers
+    assert not MODEL_MANAGER._finders
+    assert not MODEL_MANAGER._initializers
+    with runtime_grpc_test_server(
+        open_port,
+        inference_service=sample_inference_service,
+        training_service=sample_train_service,
+    ):
+        assert MODEL_MANAGER._trainers
+        assert MODEL_MANAGER._finders
+        assert MODEL_MANAGER._initializers
+
+
 def test_predict_sample_module_ok_response(
     sample_task_model_id, runtime_grpc_server, sample_inference_service
 ):
diff --git a/tests/runtime/utils/test_import_util.py b/tests/runtime/utils/test_import_util.py
index 27f98b89..6a58d36d 100644
--- a/tests/runtime/utils/test_import_util.py
+++ b/tests/runtime/utils/test_import_util.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 # Standard
+from unittest import mock
 import inspect
 import sys
 
@@ -20,6 +21,7 @@ import sys
 import pytest
 
 # Local
+from caikit.core import MODEL_MANAGER
 from caikit.core.data_model.base import DataBase
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from caikit.runtime.utils.import_util import (
@@ -140,3 +142,28 @@ def test_multiple_caikit_libraries():
                 and issubclass(attr_val, DataBase)
             ]
             assert all(attrs_match)
+
+
+def test_get_data_model_init_components_after_import():
+    """Make sure that the model_manager's init_components is called _after_ the
+    dynamic import of the runtime library
+    """
+    with temp_config({"runtime": {"library": "sample_lib"}}):
+        call_mock = mock.MagicMock()
+        with mock.patch.object(
+            MODEL_MANAGER,
+            "initialize_components",
+            call_mock,
+        ):
+            with mock.patch(
+                "caikit.runtime.utils.import_util._get_cdm_from_lib", call_mock
+            ):
+                get_data_model()
+                # Make sure the mock was called for both _get_cdm_from_lib and
+                # initialize_components
+                assert len(call_mock.call_args_list) > 1
+
+                # Make sure the last call is the only one with no args
+                # indicating that it's the call to initialize_components
+                assert not call_mock.call_args_list[-1].args
+                assert all(call.args for call in call_mock.call_args_list[:-1])