Repository: caikit/caikit
Pull request #373 (2023-07-26 22:18:08)

:recycle: Move training info out of interfaces
<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:
1.  Move `TrainingInfo` down to `caikit.core` as part of the ModelTrainer / ModelFuture api
2. Make trainers return exception instead of returning error strings

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 4cf06366..3fd5d499 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -119,6 +119,9 @@ runtime:
         enabled: true
         port: 8086
 
+    # Use arbitrary available ports if the specified ones aren't free
+    find_available_ports: false
+
     # Whether or not to abort work on client cancellation
     use_abortable_threads: true
 
diff --git a/caikit/core/data_model/__init__.py b/caikit/core/data_model/__init__.py
index a6fd6eed..71f5bce6 100644
--- a/caikit/core/data_model/__init__.py
+++ b/caikit/core/data_model/__init__.py
@@ -29,4 +29,4 @@ from .enums import *
 from .producer import PACKAGE_COMMON, ProducerId
 from .streams import data_stream
 from .streams.data_stream import *
-from .training_status import TrainingInfo, TrainingStatus
+from .training_status import TrainingStatus
diff --git a/caikit/core/data_model/training_status.py b/caikit/core/data_model/training_status.py
index 476309e8..24eb1243 100644
--- a/caikit/core/data_model/training_status.py
+++ b/caikit/core/data_model/training_status.py
@@ -17,13 +17,9 @@ Common data model enum used for reporting training status
 
 # Standard
 from enum import Enum
-from typing import List
-
-# First Party
-from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
 
 # Local
-from .dataobject import DataObjectBase, dataobject
+from .dataobject import dataobject
 from .package import PACKAGE_COMMON
 
 
@@ -42,12 +38,3 @@ class TrainingStatus(Enum):
             self.__class__.CANCELED,
             self.__class__.ERRORED,
         ]
-
-
-@dataobject(PACKAGE_COMMON)
-class TrainingInfo(DataObjectBase):
-    errors: Annotated[List[str], FieldNumber(1)]
-    # TODO: Add elements to conveying other useful information
-    # regarding training status, such as iterations progressed
-    # evaluation so far, etc.
-    status: Annotated[TrainingStatus, FieldNumber(2)]
diff --git a/caikit/core/model_management/__init__.py b/caikit/core/model_management/__init__.py
index abd77e75..bcb2b0a5 100644
--- a/caikit/core/model_management/__init__.py
+++ b/caikit/core/model_management/__init__.py
@@ -27,4 +27,4 @@ from .factories import (
 )
 from .model_finder_base import ModelFinderBase
 from .model_initializer_base import ModelInitializerBase
-from .model_trainer_base import ModelTrainerBase
+from .model_trainer_base import ModelTrainerBase, TrainingInfo
diff --git a/caikit/core/model_management/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index e1fe34a8..8c0c3964 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -29,12 +29,12 @@ import aconfig
 import alog
 
 # Local
-from ..data_model import TrainingInfo, TrainingStatus
+from ..data_model import TrainingStatus
 from ..modules import ModuleBase
 from ..toolkit.destroyable_process import DestroyableProcess
 from ..toolkit.destroyable_thread import DestroyableThread
 from ..toolkit.errors import error_handler
-from .model_trainer_base import ModelTrainerBase
+from .model_trainer_base import ModelTrainerBase, TrainingInfo
 import caikit
 
 log = alog.use_channel("LOC-TRNR")
@@ -145,7 +145,7 @@ class LocalModelTrainer(ModelTrainerBase):
             # The worker threw outside of a cancellation process
             if self._worker.threw:
                 return TrainingInfo(
-                    status=TrainingStatus.ERRORED, errors=[str(self._worker.error)]
+                    status=TrainingStatus.ERRORED, errors=[self._worker.error]
                 )
 
             # The worker completed its work without being canceled or raising
diff --git a/caikit/core/model_management/model_trainer_base.py b/caikit/core/model_management/model_trainer_base.py
index 311e786d..78a8e80c 100644
--- a/caikit/core/model_management/model_trainer_base.py
+++ b/caikit/core/model_management/model_trainer_base.py
@@ -26,17 +26,24 @@ model_management:
 """
 
 # Standard
-from typing import Optional, Type
+from typing import List, Optional, Type
 import abc
+import dataclasses
 import os
 
 # Local
-from ..data_model import TrainingInfo
+from ..data_model import TrainingStatus
 from ..modules import ModuleBase
 from ..toolkit.factory import FactoryConstructible
 from ..toolkit.reversible_hasher import ReversibleHasher
 
 
+@dataclasses.dataclass
+class TrainingInfo:
+    status: TrainingStatus
+    errors: Optional[List[Exception]] = None
+
+
 class ModelTrainerBase(FactoryConstructible):
     __doc__ = __doc__
 
diff --git a/caikit/interfaces/runtime/data_model/__init__.py b/caikit/interfaces/runtime/data_model/__init__.py
index 657de3d8..db2bf67d 100644
--- a/caikit/interfaces/runtime/data_model/__init__.py
+++ b/caikit/interfaces/runtime/data_model/__init__.py
@@ -17,7 +17,7 @@ from . import training_management
 from .training_management import (
     ModelPointer,
     TrainingInfoRequest,
-    TrainingInfoResponse,
     TrainingJob,
     TrainingStatus,
+    TrainingStatusResponse,
 )
diff --git a/caikit/interfaces/runtime/data_model/training_management.py b/caikit/interfaces/runtime/data_model/training_management.py
index 2f1fefb9..1c8683b5 100644
--- a/caikit/interfaces/runtime/data_model/training_management.py
+++ b/caikit/interfaces/runtime/data_model/training_management.py
@@ -12,7 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Standard
+from datetime import datetime
+from typing import List
+
 # First Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
 import alog
 
 # Local
@@ -45,9 +50,9 @@ class ModelPointer(DataObjectBase):
 
 @work_in_progress(action=Action.WARNING, category=WipCategory.BETA)
 @dataobject(RUNTIME_PACKAGE)
-class TrainingInfoResponse(DataObjectBase):
-    training_id: str
-    status: TrainingStatus
-    submission_timestamp: str
-    completion_timestamp: str
-    error_code: str
+class TrainingStatusResponse(DataObjectBase):
+    training_id: Annotated[str, FieldNumber(1)]
+    state: Annotated[TrainingStatus, FieldNumber(2)]
+    submission_timestamp: Annotated[datetime, FieldNumber(3)]
+    completion_timestamp: Annotated[datetime, FieldNumber(4)]
+    reasons: Annotated[List[str], FieldNumber(5)]
diff --git a/caikit/runtime/server_base.py b/caikit/runtime/server_base.py
index 2baf594d..3fce363f 100644
--- a/caikit/runtime/server_base.py
+++ b/caikit/runtime/server_base.py
@@ -16,6 +16,7 @@
 # Standard
 from typing import Optional
 import abc
+import socket
 
 # First Party
 import aconfig
@@ -33,12 +34,49 @@ class RuntimeServerBase(abc.ABC):
 
     def __init__(self, base_port: int, tls_config_override: Optional[aconfig.Config]):
         self.config = get_config()
-        self.port = base_port
+        self.port = (
+            self._find_port(base_port)
+            if self.config.runtime.find_available_ports
+            else base_port
+        )
+        if self.port != base_port:
+            log.warning(
+                "Port %d was in use, had to find another! %d", base_port, self.port
+            )
         self.tls_config = (
             tls_config_override if tls_config_override else self.config.runtime.tls
         )
         log.debug4("Full caikit config: %s", get_config())
 
+    @classmethod
+    def _find_port(cls, start=8888, end=None, host="127.0.0.1"):
+        """Function to find an available port in a given range
+        Args:
+            start: int
+                Starting number for port search (inclusive)
+                Default: 8888
+            end: Optional[int]
+                End number for port search (exclusive)
+                Default: start + 1000
+            host: str
+                Host name or ip address to search on
+                Default: localhost
+        Returns:
+            int
+                Available port
+        """
+        end = start + 1000 if end is None else end
+        if start < end:
+            with socket.socket() as soc:
+                # soc.connect_ex returns 0 if connection is successful and thus
+                # indicating port is available
+                if soc.connect_ex((host, start)) == 0:
+                    # port is in use, thus connection to it is successful
+                    return cls._find_port(start + 1, end, host)
+
+                # port is open
+                return start
+
     def _shut_down_model_manager(self):
         """Shared utility for shutting down the model manager"""
         ModelManager.get_instance().shut_down()
diff --git a/caikit/runtime/service_factory.py b/caikit/runtime/service_factory.py
index ce05b440..71133342 100644
--- a/caikit/runtime/service_factory.py
+++ b/caikit/runtime/service_factory.py
@@ -33,7 +33,7 @@ from caikit import get_config
 from caikit.core import ModuleBase
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
-    TrainingInfoResponse,
+    TrainingStatusResponse,
 )
 from caikit.runtime import service_generation
 from caikit.runtime.service_generation.rpcs import CaikitRPCBase, snake_to_upper_camel
@@ -49,7 +49,7 @@ TRAINING_MANAGEMENT_SERVICE_SPEC = {
             {
                 "name": "GetTrainingStatus",
                 "input_type": TrainingInfoRequest.get_proto_class().DESCRIPTOR.full_name,
-                "output_type": TrainingInfoResponse.get_proto_class().DESCRIPTOR.full_name,
+                "output_type": TrainingStatusResponse.get_proto_class().DESCRIPTOR.full_name,
             }
         ]
     }
diff --git a/caikit/runtime/servicers/global_train_servicer.py b/caikit/runtime/servicers/global_train_servicer.py
index f964b570..b561f29f 100644
--- a/caikit/runtime/servicers/global_train_servicer.py
+++ b/caikit/runtime/servicers/global_train_servicer.py
@@ -29,7 +29,6 @@ import alog
 from caikit import get_config
 from caikit.core import MODEL_MANAGER, ModuleBase
 from caikit.core.data_model import DataBase
-from caikit.core.model_management.local_model_trainer import LocalModelTrainer
 from caikit.interfaces.runtime.data_model import TrainingJob
 from caikit.runtime.model_management.model_manager import ModelManager
 from caikit.runtime.service_factory import ServicePackage
@@ -258,11 +257,9 @@ class GlobalTrainServicer:
                         "Training %s failed with error: %s. "
                         "Re-raising exception for synchronous response",
                         model_future.id,
-                        training_info.errors[0],
+                        str(training_info.errors[0]),
                     )
-                    if isinstance(model_future, LocalModelTrainer.LocalModelFuture):
-                        raise model_future._worker.error
-                    raise RuntimeError(training_info.errors[0])
+                    raise training_info.errors[0]
 
         # return TrainingJob object
         return TrainingJob(
diff --git a/caikit/runtime/servicers/training_management_servicer.py b/caikit/runtime/servicers/training_management_servicer.py
index d61139f8..14216cfe 100644
--- a/caikit/runtime/servicers/training_management_servicer.py
+++ b/caikit/runtime/servicers/training_management_servicer.py
@@ -25,7 +25,7 @@ import alog
 from caikit.core import MODEL_MANAGER
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
-    TrainingInfoResponse,
+    TrainingStatusResponse,
 )
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 
@@ -45,9 +45,14 @@ class TrainingManagementServicerImpl:
                 training_id=training_info.training_id
             )
 
-            return TrainingInfoResponse(
+            reasons = []
+            if model_future.get_info().errors:
+                reasons = [str(error) for error in model_future.get_info().errors]
+
+            return TrainingStatusResponse(
                 training_id=training_info.training_id,
-                status=model_future.get_info().status,
+                state=model_future.get_info().status,
+                reasons=reasons,
             ).to_proto()
         except ValueError as err:
             raise CaikitRuntimeException(


Test patch:
diff --git a/tests/core/helpers.py b/tests/core/helpers.py
index ff780e98..d1864694 100644
--- a/tests/core/helpers.py
+++ b/tests/core/helpers.py
@@ -21,10 +21,11 @@ import pytest
 
 # Local
 from caikit.core import MODEL_MANAGER
-from caikit.core.data_model import TrainingInfo, TrainingStatus
+from caikit.core.data_model import TrainingStatus
 from caikit.core.model_management import (
     ModelInitializerBase,
     ModelTrainerBase,
+    TrainingInfo,
     model_initializer_factory,
     model_trainer_factory,
 )
diff --git a/tests/core/model_management/test_local_model_trainer.py b/tests/core/model_management/test_local_model_trainer.py
index cd9404b3..101cb355 100644
--- a/tests/core/model_management/test_local_model_trainer.py
+++ b/tests/core/model_management/test_local_model_trainer.py
@@ -220,4 +220,5 @@ def test_get_into_return_error(trainer_type_cfg):
     assert model_future.get_info().status == TrainingStatus.ERRORED
     assert model_future.get_info().status.is_terminal
     assert isinstance(model_future.get_info().errors, list)
-    assert model_future.get_info().errors[0] == "Batch size of 999 is not allowed!"
+    assert isinstance(model_future.get_info().errors[0], ValueError)
+    assert str(model_future.get_info().errors[0]) == "Batch size of 999 is not allowed!"
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index 1a8858d5..cf9ee9cd 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -338,7 +338,7 @@ def test_global_train_Edge_Case_Widget_should_raise_when_error_surfaces_from_mod
     future_info = MODEL_MANAGER.get_model_future(
         training_response.training_id
     ).get_info()
-    assert f"Batch size of 999 is not allowed!" in future_info.errors[0]
+    assert f"Batch size of 999 is not allowed!" in str(future_info.errors[0])
 
 
 def test_global_train_returns_exit_code_with_oom(
@@ -366,7 +366,7 @@ def test_global_train_returns_exit_code_with_oom(
         future_info = MODEL_MANAGER.get_model_future(
             training_response.training_id
         ).get_info()
-        assert f"Training process died with OOM error!" in future_info.errors[0]
+        assert f"Training process died with OOM error!" in str(future_info.errors[0])
 
 
 #####################################################################
diff --git a/tests/runtime/servicers/test_training_management_servicer.py b/tests/runtime/servicers/test_training_management_servicer.py
index 49555606..a0a941b5 100644
--- a/tests/runtime/servicers/test_training_management_servicer.py
+++ b/tests/runtime/servicers/test_training_management_servicer.py
@@ -64,7 +64,7 @@ def test_training_runs(training_management_servicer, training_pool):
     # send a request, check it's running
     request = TrainingInfoRequest(training_id=model_future.id).to_proto()
     response = training_management_servicer.GetTrainingStatus(request, context=None)
-    assert response.status == TrainingStatus.RUNNING.value
+    assert response.state == TrainingStatus.RUNNING.value
 
     event.set()
     model_future.wait()
@@ -72,7 +72,7 @@ def test_training_runs(training_management_servicer, training_pool):
     # Ensure it's now done
     request = TrainingInfoRequest(training_id=model_future.id).to_proto()
     response = training_management_servicer.GetTrainingStatus(request, context=None)
-    assert response.status == TrainingStatus.COMPLETED.value
+    assert response.state == TrainingStatus.COMPLETED.value
 
 
 def test_training_complete_status(training_management_servicer, training_pool):
@@ -90,8 +90,8 @@ def test_training_complete_status(training_management_servicer, training_pool):
     request = TrainingInfoRequest(training_id=model_future.id).to_proto()
     response = training_management_servicer.GetTrainingStatus(request, context=None)
 
-    assert re.match("<class '.*TrainingInfoResponse'>", str(type(response)))
-    assert response.status == TrainingStatus.COMPLETED.value
+    assert re.match("<class '.*TrainingStatusResponse'>", str(type(response)))
+    assert response.state == TrainingStatus.COMPLETED.value
 
 
 def test_training_status_incorrect_id(training_management_servicer):
@@ -119,4 +119,4 @@ def test_training_fails(training_management_servicer, training_pool):
     request = TrainingInfoRequest(training_id=model_future.id).to_proto()
     response = training_management_servicer.GetTrainingStatus(request, context=None)
 
-    assert response.status == TrainingStatus.ERRORED.value
+    assert response.state == TrainingStatus.ERRORED.value
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index f732aa48..84b878f2 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -45,9 +45,9 @@ from caikit.core import MODEL_MANAGER
 from caikit.core.data_model.base import DataBase
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
-    TrainingInfoResponse,
     TrainingJob,
     TrainingStatus,
+    TrainingStatusResponse,
 )
 from caikit.runtime.grpc_server import RuntimeGRPCServer
 from caikit.runtime.model_management.model_manager import ModelManager
@@ -101,14 +101,14 @@ def is_good_train_response(
         training_info_request = TrainingInfoRequest(
             training_id=actual_response.training_id
         )
-        training_management_response: TrainingInfoResponse = (
-            TrainingInfoResponse.from_proto(
+        training_management_response: TrainingStatusResponse = (
+            TrainingStatusResponse.from_proto(
                 training_management_stub.GetTrainingStatus(
                     training_info_request.to_proto()
                 )
             )
         )
-        status = training_management_response.status
+        status = training_management_response.state
         assert status != TrainingStatus.ERRORED.value
         i += 1
         assert i < 100, "Waited too long for training to complete"
@@ -173,10 +173,10 @@ def test_model_train(runtime_grpc_server):
     )
 
     # MT training is sync, so it should be COMPLETE immediately
-    response: TrainingInfoResponse = TrainingInfoResponse.from_proto(
+    response: TrainingStatusResponse = TrainingStatusResponse.from_proto(
         training_management_stub.GetTrainingStatus(training_info_request)
     )
-    assert response.status == TrainingStatus.COMPLETED.value
+    assert response.state == TrainingStatus.COMPLETED.value
 
     # Make sure we wait for training to finish
     result = MODEL_MANAGER.get_model_future(response.training_id).load()
diff --git a/tests/runtime/test_http_server.py b/tests/runtime/test_http_server.py
index 2f5e3bef..9731ee92 100644
--- a/tests/runtime/test_http_server.py
+++ b/tests/runtime/test_http_server.py
@@ -34,7 +34,7 @@ from caikit.core import DataObjectBase, dataobject
 from caikit.interfaces.nlp.data_model import GeneratedTextStreamResult, GeneratedToken
 from caikit.runtime import http_server
 from tests.conftest import temp_config
-from tests.runtime.conftest import ModuleSubproc
+from tests.runtime.conftest import ModuleSubproc, open_port
 
 ## Helpers #####################################################################
 
@@ -66,7 +66,7 @@ class TLSConfig:
 
 @contextmanager
 def generate_tls_configs(
-    port: int, tls: bool = False, mtls: bool = False, **http_config_overrides
+    tls: bool = False, mtls: bool = False, **http_config_overrides
 ):
     """Helper to generate tls configs"""
     with tempfile.TemporaryDirectory() as workdir:
@@ -110,6 +110,7 @@ def generate_tls_configs(
                 config_overrides["use_in_test"]["client_key"] = client_keyfile
 
             config_overrides["runtime"] = {"tls": tls_config}
+        port = http_server.RuntimeServerBase._find_port()
         config_overrides.setdefault("runtime", {})["http"] = {
             "server_shutdown_grace_period_seconds": 0.01,  # this is so the server is killed after 0.1 if no test is running
             "port": port,
@@ -123,8 +124,8 @@ def generate_tls_configs(
 ## Insecure and TLS Tests #######################################################################
 
 
-def test_insecure_server(open_port):
-    with generate_tls_configs(open_port):
+def test_insecure_server():
+    with generate_tls_configs():
         insecure_http_server = http_server.RuntimeHTTPServer()
         # start a non-blocking http server
         with insecure_http_server:
@@ -132,9 +133,9 @@ def test_insecure_server(open_port):
             resp.raise_for_status()
 
 
-def test_basic_tls_server(open_port):
+def test_basic_tls_server():
     with generate_tls_configs(
-        open_port, tls=True, mtls=False, http_config_overrides={}
+        tls=True, mtls=False, http_config_overrides={}
     ) as config_overrides:
         http_server_with_tls = http_server.RuntimeHTTPServer(
             tls_config_override=config_overrides["runtime"]["tls"]
@@ -148,9 +149,9 @@ def test_basic_tls_server(open_port):
             resp.raise_for_status()
 
 
-def test_basic_tls_server_with_wrong_cert(open_port):
+def test_basic_tls_server_with_wrong_cert():
     with generate_tls_configs(
-        open_port, tls=True, mtls=False, http_config_overrides={}
+        tls=True, mtls=False, http_config_overrides={}
     ) as config_overrides:
         http_server_with_tls = http_server.RuntimeHTTPServer(
             tls_config_override=config_overrides["runtime"]["tls"]
@@ -164,9 +165,9 @@ def test_basic_tls_server_with_wrong_cert(open_port):
                 )
 
 
-def test_mutual_tls_server(open_port):
+def test_mutual_tls_server():
     with generate_tls_configs(
-        open_port, tls=True, mtls=True, http_config_overrides={}
+        tls=True, mtls=True, http_config_overrides={}
     ) as config_overrides:
         http_server_with_mtls = http_server.RuntimeHTTPServer(
             tls_config_override=config_overrides["runtime"]["tls"]
@@ -184,9 +185,9 @@ def test_mutual_tls_server(open_port):
             resp.raise_for_status()
 
 
-def test_mutual_tls_server_with_wrong_cert(open_port):
+def test_mutual_tls_server_with_wrong_cert():
     with generate_tls_configs(
-        open_port, tls=True, mtls=True, http_config_overrides={}
+        tls=True, mtls=True, http_config_overrides={}
     ) as config_overrides:
         http_server_with_mtls = http_server.RuntimeHTTPServer(
             tls_config_override=config_overrides["runtime"]["tls"]
diff --git a/tests/runtime/utils/test_servicer_util.py b/tests/runtime/utils/test_servicer_util.py
index b987c251..264ce6dd 100644
--- a/tests/runtime/utils/test_servicer_util.py
+++ b/tests/runtime/utils/test_servicer_util.py
@@ -45,12 +45,12 @@ import sample_lib
 def test_servicer_util_validate_caikit_library_class_exists_returns_caikit_class():
     """Test that validate_caikit_library_class_exists returns a caikit class object"""
     validated_class = validate_caikit_library_class_exists(
-        get_data_model(), "TrainingInfoResponse"
+        get_data_model(), "TrainingStatusResponse"
     )
 
     assert issubclass(
         validated_class,
-        caikit.interfaces.runtime.data_model.training_management.TrainingInfoResponse,
+        caikit.interfaces.runtime.data_model.training_management.TrainingStatusResponse,
     )
 
     """