Repository: caikit/caikit
Pull request #357 (2023-07-24 21:55:47)

:sparkles: Add training info DM
Supports: https://github.com/caikit/caikit/issues/354

### Changes
- Add TrainingInfo DM object
- Replace `get_status` with `get_info`
- Add error information for local trainer



Patch:
diff --git a/caikit/core/data_model/__init__.py b/caikit/core/data_model/__init__.py
index 71f5bce6..a6fd6eed 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 TrainingStatus
+from .training_status import TrainingInfo, TrainingStatus
diff --git a/caikit/core/data_model/training_status.py b/caikit/core/data_model/training_status.py
index a39e89b0..476309e8 100644
--- a/caikit/core/data_model/training_status.py
+++ b/caikit/core/data_model/training_status.py
@@ -17,14 +17,16 @@ 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 ..toolkit.wip_decorator import Action, WipCategory, work_in_progress
-from .dataobject import dataobject
+from .dataobject import DataObjectBase, dataobject
 from .package import PACKAGE_COMMON
 
 
-@work_in_progress(action=Action.WARNING, category=WipCategory.BETA)
 @dataobject(PACKAGE_COMMON)
 class TrainingStatus(Enum):
     QUEUED = 1
@@ -40,3 +42,12 @@ 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/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index 3c92c387..e1fe34a8 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -29,7 +29,7 @@ import aconfig
 import alog
 
 # Local
-from ..data_model import TrainingStatus
+from ..data_model import TrainingInfo, TrainingStatus
 from ..modules import ModuleBase
 from ..toolkit.destroyable_process import DestroyableProcess
 from ..toolkit.destroyable_thread import DestroyableThread
@@ -41,9 +41,6 @@ log = alog.use_channel("LOC-TRNR")
 error = error_handler.get(log)
 
 
-OOM_EXIT_CODE = 137
-
-
 # 🌶️🌶️🌶️
 # Fix for python3.9, 3.10 and 3.11 issue where forked processes always exit with exitcode 1
 # when it's created inside a ThreadPoolExecutor: https://github.com/python/cpython/issues/88110
@@ -105,7 +102,6 @@ class LocalModelTrainer(ModelTrainerBase):
                     return_result=True,
                     args=args,
                     kwargs={
-                        "return_completion_time": True,
                         **kwargs,
                     },
                 )
@@ -132,29 +128,32 @@ class LocalModelTrainer(ModelTrainerBase):
 
         ## Interface ##
 
-        def get_status(self) -> TrainingStatus:
+        def get_info(self) -> TrainingInfo:
             """Every model future must be able to poll the status of the
             training job
             """
+
             # The worker was canceled while doing work. It may still be in the
             # process of terminating and thus still alive.
             if self._worker.canceled:
-                return TrainingStatus.CANCELED
+                return TrainingInfo(status=TrainingStatus.CANCELED)
 
             # If the worker is currently alive it's doing work
             if self._worker.is_alive():
-                return TrainingStatus.RUNNING
+                return TrainingInfo(status=TrainingStatus.RUNNING)
 
             # The worker threw outside of a cancellation process
             if self._worker.threw:
-                return TrainingStatus.ERRORED
+                return TrainingInfo(
+                    status=TrainingStatus.ERRORED, errors=[str(self._worker.error)]
+                )
 
             # The worker completed its work without being canceled or raising
             if self._worker.ran:
-                return TrainingStatus.COMPLETED
+                return TrainingInfo(status=TrainingStatus.COMPLETED)
 
             # If it's not alive and not done, it hasn't started yet
-            return TrainingStatus.QUEUED
+            return TrainingInfo(status=TrainingStatus.QUEUED)
 
         def cancel(self):
             """Terminate the given training"""
@@ -170,17 +169,7 @@ class LocalModelTrainer(ModelTrainerBase):
             log.debug2("Waiting for %s", self.id)
             self._worker.join()
             log.debug2("Done waiting for %s", self.id)
-
-            # If running a subprocess, handle abnormal exit codes
-            if self._use_subprocess:
-                self._completion_time = self._worker.get_or_throw()
-                if self._worker.exitcode and self._worker.exitcode != os.EX_OK:
-                    if self._worker.exitcode == OOM_EXIT_CODE:
-                        raise MemoryError("Training process died with OOM error!")
-                    if not self._worker.canceled:
-                        raise RuntimeError(
-                            f"Training process died with exit code {self._worker.exitcode}"
-                        )
+            self._completion_time = self._completion_time or datetime.now()
 
         def load(self) -> ModuleBase:
             """Wait for the training to complete, then return the resulting
@@ -211,7 +200,7 @@ class LocalModelTrainer(ModelTrainerBase):
 
         ## Impl ##
 
-        def _train_and_save(self, *args, return_completion_time=False, **kwargs):
+        def _train_and_save(self, *args, **kwargs):
             """Function that will run in the worker thread"""
             with alog.ContextTimer(log.debug, "Training %s finished in: ", self.id):
                 trained_model = self._module_class.train(*args, **kwargs)
@@ -219,10 +208,7 @@ class LocalModelTrainer(ModelTrainerBase):
                 log.debug("Saving training %s to %s", self.id, self.save_path)
                 with alog.ContextTimer(log.debug, "Training %s saved in: ", self.id):
                     trained_model.save(self.save_path)
-            self._completion_time = datetime.now()
             log.debug2("Completion time for %s: %s", self.id, self._completion_time)
-            if return_completion_time:
-                return self._completion_time
             return trained_model
 
     ## Interface ##
diff --git a/caikit/core/model_management/model_trainer_base.py b/caikit/core/model_management/model_trainer_base.py
index 3047c887..311e786d 100644
--- a/caikit/core/model_management/model_trainer_base.py
+++ b/caikit/core/model_management/model_trainer_base.py
@@ -31,7 +31,7 @@ import abc
 import os
 
 # Local
-from ..data_model import TrainingStatus
+from ..data_model import TrainingInfo
 from ..modules import ModuleBase
 from ..toolkit.factory import FactoryConstructible
 from ..toolkit.reversible_hasher import ReversibleHasher
@@ -78,7 +78,7 @@ class ModelTrainerBase(FactoryConstructible):
             return self._save_path
 
         @abc.abstractmethod
-        def get_status(self) -> TrainingStatus:
+        def get_info(self) -> TrainingInfo:
             """Every model future must be able to poll the status of the
             training job
             """
diff --git a/caikit/core/toolkit/destroyable.py b/caikit/core/toolkit/destroyable.py
index 1140e576..b44b912d 100644
--- a/caikit/core/toolkit/destroyable.py
+++ b/caikit/core/toolkit/destroyable.py
@@ -52,3 +52,8 @@ class Destroyable(abc.ABC):
     @abc.abstractmethod
     def destroy(self):
         """Cancel any in-progress work"""
+
+    @property
+    @abc.abstractmethod
+    def error(self) -> bool:
+        """Return the error information to user if one occurred"""
diff --git a/caikit/core/toolkit/destroyable_process.py b/caikit/core/toolkit/destroyable_process.py
index 26db9a79..cadd9693 100644
--- a/caikit/core/toolkit/destroyable_process.py
+++ b/caikit/core/toolkit/destroyable_process.py
@@ -25,6 +25,7 @@ NOTE: The "fork" start method is used for two reasons:
 from functools import partial
 from typing import Any, Callable, Optional, Tuple
 import multiprocessing
+import os
 
 # First Party
 import alog
@@ -36,6 +37,8 @@ log = alog.use_channel("DESTROY-PROC")
 
 FORK_CTX = multiprocessing.get_context("fork")
 
+OOM_EXIT_CODE = 137
+
 
 class DestroyableProcess(
     FORK_CTX.Process, Destroyable
@@ -175,9 +178,18 @@ class DestroyableProcess(
     @property
     def error(self) -> Optional[Exception]:
         self._update_result()
+
         if isinstance(self.__result, Exception):
             return self.__result
 
+        if self.exitcode and self.exitcode != os.EX_OK:
+            if self.exitcode == OOM_EXIT_CODE:
+                return MemoryError("Training process died with OOM error!")
+            if not self.canceled:
+                return RuntimeError(
+                    f"Training process died with exit code {self.exitcode}"
+                )
+
     @property
     def completion_event(self) -> multiprocessing.Event:
         return self._completion_event
diff --git a/caikit/core/toolkit/destroyable_thread.py b/caikit/core/toolkit/destroyable_thread.py
index 34ac4ddb..9c692e11 100644
--- a/caikit/core/toolkit/destroyable_thread.py
+++ b/caikit/core/toolkit/destroyable_thread.py
@@ -221,6 +221,11 @@ class DestroyableThread(threading.Thread, Destroyable):
                 async_exception_result,
             )
 
+    @property
+    def error(self) -> Optional[Exception]:
+        if isinstance(self.__runnable_exception, Exception):
+            return self.__runnable_exception
+
     def __get_id(self):
         # Returns the thread if if the thread is running
         for thread in threading.enumerate():
diff --git a/caikit/runtime/servicers/global_train_servicer.py b/caikit/runtime/servicers/global_train_servicer.py
index c057865e..5c61784c 100644
--- a/caikit/runtime/servicers/global_train_servicer.py
+++ b/caikit/runtime/servicers/global_train_servicer.py
@@ -226,7 +226,7 @@ class GlobalTrainServicer:
                 # Create a callback to register termination of training
                 def rpc_termination_callback():
                     """Cancel the model future if it has not yet completed"""
-                    if not model_future.get_status().is_terminal:
+                    if not model_future.get_info().status.is_terminal:
                         log.warning(
                             "<RUN36361257W>", "Canceling training %s", model_future.id
                         )
@@ -250,6 +250,9 @@ class GlobalTrainServicer:
                 log.debug, "Training %s complete in: ", model_future.id
             ):
                 model_future.wait()
+                training_info = model_future.get_info()
+                if training_info.errors:
+                    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 54ec4e48..d61139f8 100644
--- a/caikit/runtime/servicers/training_management_servicer.py
+++ b/caikit/runtime/servicers/training_management_servicer.py
@@ -47,7 +47,7 @@ class TrainingManagementServicerImpl:
 
             return TrainingInfoResponse(
                 training_id=training_info.training_id,
-                status=model_future.get_status(),
+                status=model_future.get_info().status,
             ).to_proto()
         except ValueError as err:
             raise CaikitRuntimeException(


Test patch:
diff --git a/tests/core/helpers.py b/tests/core/helpers.py
index 01c88b48..ff780e98 100644
--- a/tests/core/helpers.py
+++ b/tests/core/helpers.py
@@ -21,7 +21,7 @@ import pytest
 
 # Local
 from caikit.core import MODEL_MANAGER
-from caikit.core.data_model import TrainingStatus
+from caikit.core.data_model import TrainingInfo, TrainingStatus
 from caikit.core.model_management import (
     ModelInitializerBase,
     ModelTrainerBase,
@@ -112,12 +112,12 @@ class TestTrainer(ModelTrainerBase):
             self._canceled = False
             self._completed = False
 
-        def get_status(self):
+        def get_info(self):
             if self._completed:
-                return TrainingStatus.COMPLETED
+                return TrainingInfo(status=TrainingStatus.COMPLETED)
             if self._canceled:
-                return TrainingStatus.CANCELED
-            return self._parent.canned_status
+                return TrainingInfo(status=TrainingStatus.CANCELED)
+            return TrainingInfo(status=self._parent.canned_status)
 
         def cancel(self):
             self._canceled = True
diff --git a/tests/core/model_management/test_local_model_trainer.py b/tests/core/model_management/test_local_model_trainer.py
index 31ea55d2..11466a51 100644
--- a/tests/core/model_management/test_local_model_trainer.py
+++ b/tests/core/model_management/test_local_model_trainer.py
@@ -30,10 +30,8 @@ import aconfig
 
 # Local
 from caikit.core.data_model import DataStream, TrainingStatus
-from caikit.core.model_management.local_model_trainer import (
-    OOM_EXIT_CODE,
-    LocalModelTrainer,
-)
+from caikit.core.model_management.local_model_trainer import LocalModelTrainer
+from caikit.core.toolkit.destroyable_process import OOM_EXIT_CODE
 from sample_lib.modules import SampleModule
 
 ## Helpers #####################################################################
@@ -62,7 +60,7 @@ def get_event(cfg: dict):
 ## Tests #######################################################################
 
 
-def test_train_and_get_status(trainer_type_cfg):
+def test_train_and_get_info(trainer_type_cfg):
     """Test that running a training can fetch status correctly"""
     trainer = local_trainer(**trainer_type_cfg)
 
@@ -75,14 +73,14 @@ def test_train_and_get_status(trainer_type_cfg):
         DataStream.from_iterable([]),
         wait_event=wait_event,
     )
-    assert model_future.get_status() == TrainingStatus.RUNNING
-    assert not model_future.get_status().is_terminal
+    assert model_future.get_info().status == TrainingStatus.RUNNING
+    assert not model_future.get_info().status.is_terminal
 
     # Let the training proceed and wait for it to complete
     wait_event.set()
     model_future.wait()
-    assert model_future.get_status() == TrainingStatus.COMPLETED
-    assert model_future.get_status().is_terminal
+    assert model_future.get_info().status == TrainingStatus.COMPLETED
+    assert model_future.get_info().status.is_terminal
 
     # Re-fetch the future by ID
     fetched_future = trainer.get_model_future(model_future.id)
@@ -137,13 +135,13 @@ def test_cancel_clean_termination(trainer_type_cfg):
         training_data=DataStream.from_iterable([]),
         sleep_time=1000,
     )
-    assert model_future.get_status() == TrainingStatus.RUNNING
-    assert not model_future.get_status().is_terminal
+    assert model_future.get_info().status == TrainingStatus.RUNNING
+    assert not model_future.get_info().status.is_terminal
 
     # Cancel the future
     model_future.cancel()
-    assert model_future.get_status() == TrainingStatus.CANCELED
-    assert model_future.get_status().is_terminal
+    assert model_future.get_info().status == TrainingStatus.CANCELED
+    assert model_future.get_info().status.is_terminal
     model_future.wait()
 
 
@@ -159,14 +157,14 @@ def test_cancel_without_waiting(trainer_type_cfg):
         sleep_time=0.5,
         sleep_increment=0.5,
     )
-    assert model_future.get_status() == TrainingStatus.RUNNING
-    assert not model_future.get_status().is_terminal
+    assert model_future.get_info().status == TrainingStatus.RUNNING
+    assert not model_future.get_info().status.is_terminal
 
     # Cancel the future and make sure it reports canceled, even though the
     # function is still sleeping
     model_future.cancel()
-    assert model_future.get_status() == TrainingStatus.CANCELED
-    assert model_future.get_status().is_terminal
+    assert model_future.get_info().status == TrainingStatus.CANCELED
+    assert model_future.get_info().status.is_terminal
 
 
 def test_no_retention_time(trainer_type_cfg):
@@ -204,3 +202,23 @@ def test_retention_duration_parse(test_params):
     """Make sure the regex for the duration can parse all expected durations"""
     trainer = local_trainer(retention_duration=test_params[0])
     assert trainer._retention_duration == test_params[1]
+
+
+def test_get_into_return_error(trainer_type_cfg):
+    """Test that failed training returns error properly"""
+    trainer = local_trainer(**trainer_type_cfg)
+
+    model_future = trainer.train(
+        SampleModule,
+        DataStream.from_iterable([]),
+        batch_size=SampleModule.POISON_PILL_BATCH_SIZE,
+    )
+    # assert model_future.get_info().status == TrainingStatus.RUNNING
+    # assert not model_future.get_info().status.is_terminal
+
+    # Let the training proceed and wait for it to complete
+    model_future.wait()
+    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!"
diff --git a/tests/core/test_model_manager.py b/tests/core/test_model_manager.py
index a8a7e441..cc078241 100644
--- a/tests/core/test_model_manager.py
+++ b/tests/core/test_model_manager.py
@@ -604,12 +604,12 @@ def test_train_with_wait(reset_globals):
         unfinished_train_future = caikit.train(
             SampleModule, DataStream.from_iterable([])
         )
-        assert unfinished_train_future.get_status() == TrainingStatus.RUNNING
+        assert unfinished_train_future.get_info().status == TrainingStatus.RUNNING
         # Call with wait and make sure it is COMPLETED
         finished_train_future = caikit.train(
             SampleModule, DataStream.from_iterable([]), wait=True
         )
-        assert finished_train_future.get_status() == TrainingStatus.COMPLETED
+        assert finished_train_future.get_info().status == TrainingStatus.COMPLETED
 
 
 def test_train_with_save_path(reset_globals):
@@ -620,9 +620,9 @@ def test_train_with_save_path(reset_globals):
             train_future = caikit.train(
                 SampleModule, DataStream.from_iterable([]), save_path=save_path
             )
-            assert train_future.get_status() == TrainingStatus.RUNNING
+            assert train_future.get_info().status == TrainingStatus.RUNNING
             assert not os.path.exists(save_path)
             assert train_future.save_path == save_path
             train_future.wait()
-            assert train_future.get_status() == TrainingStatus.COMPLETED
+            assert train_future.get_info().status == TrainingStatus.COMPLETED
             assert os.path.exists(save_path)
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index 4a3449bd..1a8858d5 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -330,13 +330,15 @@ def test_global_train_Edge_Case_Widget_should_raise_when_error_surfaces_from_mod
         training_data=training_data,
     )
 
-    with pytest.raises(ValueError) as context:
-        training_response = sample_train_servicer.Train(
-            train_request, Fixtures.build_context("foo")
-        )
-        MODEL_MANAGER.get_model_future(training_response.training_id).load()
+    training_response = sample_train_servicer.Train(
+        train_request, Fixtures.build_context("foo")
+    )
+    MODEL_MANAGER.get_model_future(training_response.training_id).wait()
 
-    assert f"Batch size of 999 is not allowed!" in str(context.value)
+    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]
 
 
 def test_global_train_returns_exit_code_with_oom(
@@ -356,13 +358,15 @@ def test_global_train_returns_exit_code_with_oom(
 
     # Enable sub-processing for test
     with set_use_subprocess(True):
-        with pytest.raises(MemoryError) as context:
-            training_response = sample_train_servicer.Train(
-                train_request, Fixtures.build_context("foo")
-            )
-            MODEL_MANAGER.get_model_future(training_response.training_id).wait()
+        training_response = sample_train_servicer.Train(
+            train_request, Fixtures.build_context("foo")
+        )
+        MODEL_MANAGER.get_model_future(training_response.training_id).wait()
 
-    assert f"Training process died with OOM error!" in str(context.value)
+        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]
 
 
 #####################################################################
diff --git a/tests/runtime/servicers/test_model_train_servicer_impl.py b/tests/runtime/servicers/test_model_train_servicer_impl.py
index 5556069e..75edacc5 100644
--- a/tests/runtime/servicers/test_model_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_model_train_servicer_impl.py
@@ -164,6 +164,25 @@ def test_model_train_incorrect_train_request_raises(
         assert e.value.status_code == grpc.StatusCode.INTERNAL
 
 
+def test_model_train_validation_error_raises(sample_model_train_servicer, output_dir):
+    """Test that if there is validation error we are able to raise it"""
+
+    training_id = str(uuid.uuid4())
+    model_train_request = process_pb2.ProcessRequest(
+        trainingID=training_id,
+        customTrainingID=str(uuid.uuid4()),
+        request_dict={
+            "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
+            "training_params": '{"model_name": "abc", "training_data": [1], "batch_size": 999}',
+        },
+        training_input_dir="training_input_dir",
+        training_output_dir=os.path.join(output_dir, training_id),
+    )
+    with pytest.raises(CaikitRuntimeException):
+        context = Fixtures.build_context("foo")
+        sample_model_train_servicer.Run(model_train_request, context)
+
+
 #####################################################################
 # Normal tests
 def test_model_train_sample_widget(sample_model_train_servicer, output_dir):
@@ -199,7 +218,7 @@ def test_model_train_sample_widget(sample_model_train_servicer, output_dir):
 
     # Make sure that the request completed synchronously
     model_future = MODEL_MANAGER.get_model_future(training_response.trainingID)
-    assert model_future.get_status() == TrainingStatus.COMPLETED
+    assert model_future.get_info().status == TrainingStatus.COMPLETED
 
     # Make sure that the return object looks right
     assert isinstance(training_response, process_pb2.ProcessResponse)
diff --git a/tests/runtime/servicers/test_training_management_servicer.py b/tests/runtime/servicers/test_training_management_servicer.py
index 27109501..49555606 100644
--- a/tests/runtime/servicers/test_training_management_servicer.py
+++ b/tests/runtime/servicers/test_training_management_servicer.py
@@ -23,8 +23,8 @@ import pytest
 
 # Local
 from caikit.core import MODEL_MANAGER
-from caikit.core.data_model import DataStream
-from caikit.interfaces.runtime.data_model import TrainingInfoRequest, TrainingStatus
+from caikit.core.data_model import DataStream, TrainingStatus
+from caikit.interfaces.runtime.data_model import TrainingInfoRequest
 from caikit.runtime.servicers.training_management_servicer import (
     TrainingManagementServicerImpl,
 )