Repository: caikit/caikit
Pull request #709 (2024-05-03 19:51:34)

Add exception catching for everything in train.py
Because train.py is being launched via the command line, we need to convert all thrown exceptions to exit codes. When running in a k8s container environment, the error message must also be written to the termination log.


Patch:
diff --git a/caikit/runtime/train.py b/caikit/runtime/train.py
index 97d5efe8..fe26990c 100644
--- a/caikit/runtime/train.py
+++ b/caikit/runtime/train.py
@@ -24,6 +24,7 @@ import importlib
 import json
 import os
 import sys
+import traceback
 
 # Third Party
 from google.protobuf import json_format
@@ -44,10 +45,41 @@ from .utils.servicer_util import build_caikit_library_request_dict
 log = alog.use_channel("TRAIN")
 error = error_handler.get(log)
 
+# The USER_ERROR_EXIT_CODE will be thrown when the process must exit
+# as result of a user input error. User-related errors should be
+# >= 1 and <=127 due to how some kubernetes operators interpret them.
+USER_ERROR_EXIT_CODE = 1
+# The INTERNAL_ERROR_EXIT_CODE will be thrown when training
+# abnormally terminates, and it is not clearly fault of the user.
+# System-level errors should be >= 128 and <= 254
+INTERNAL_ERROR_EXIT_CODE = 203
+
+
+class ArgumentParserError(Exception):
+    """Custom exception class for ArgumentParser errors."""
+
+
+class TrainArgumentParser(argparse.ArgumentParser):
+    def error(self, message):
+        """Error handler that raises an exception instead of exiting."""
+        raise ArgumentParserError(f"{self.prog}: error: {message}")
+
+
+def write_termination_log(text: str, log_file="/dev/termination-log"):
+    try:
+        with open(log_file, "a") as handle:
+            handle.write(text)
+    except Exception as e:
+        log.warning(
+            "<COR96300323W>",
+            "Unable to write termination log due to error %s",
+            e,
+        )
+
 
 def main() -> int:
     """Main entrypoint for running training jobs"""
-    parser = argparse.ArgumentParser(description=__doc__)
+    parser = TrainArgumentParser(description=__doc__)
 
     # Required Args
     parser.add_argument(
@@ -95,44 +127,93 @@ def main() -> int:
         default=False,
         help="Include the training ID in the save path",
     )
+    parser.add_argument(
+        "--termination-log-file",
+        "-f",
+        action="store_true",
+        default="/dev/termination-log",
+        help="Location of where to write a termination error message",
+    )
 
-    args = parser.parse_args()
-    config_logging()
+    try:
+        args = parser.parse_args()
+        config_logging()
 
-    # Initialize top-level kwargs
-    train_kwargs = {
-        "save_path": args.save_path,
-        "save_with_id": args.save_with_id,
-        "model_name": args.model_name,
-    }
-    if args.trainer is not None:
-        train_kwargs["trainer"] = args.trainer
+        # Initialize top-level kwargs
+        train_kwargs = {
+            "save_path": args.save_path,
+            "save_with_id": args.save_with_id,
+            "model_name": args.model_name,
+        }
+        if args.trainer is not None:
+            train_kwargs["trainer"] = args.trainer
+
+    except Exception as e:
+        message = f"Exception raised during training. This may be a problem with your input: {e}"
+        log.warning(
+            {
+                "log_code": "<COR39662029E>",
+                "message": message,
+                "stack_trace": traceback.format_exc(),
+            },
+            exc_info=True,
+        )
+        write_termination_log(message)
+        exit(USER_ERROR_EXIT_CODE)
 
     # Import libraries to register modules
-    for library in args.library or []:
-        log.info("<COR88091092I>", "Importing library %s", library)
-        importlib.import_module(library)
+    try:
+        for library in args.library or []:
+            log.info("<COR88091092I>", "Importing library %s", library)
+            importlib.import_module(library)
+    except Exception:
+        message = "Unable to import module {}".format(library)
+        log.warning(
+            {
+                "log_code": "<COR17776539E>",
+                "message": message,
+                "stack_trace": traceback.format_exc(),
+            },
+            exc_info=True,
+        )
+        write_termination_log(message)
+        exit(USER_ERROR_EXIT_CODE)
 
     # Try to import the root library of the provided module. It's ok if this
     # fails since the module may be a UID
-    mod_root_lib = args.module.split(".")[0]
     try:
+        mod_root_lib = args.module.split(".")[0]
         importlib.import_module(mod_root_lib)
-    except ImportError:
+    except (ImportError, ValueError):
         log.debug("Unable to import module root lib: %s", mod_root_lib)
 
     # Figure out the module to train
-    mod_reg = module_registry()
-    mod_pkg_to_mod = {
-        f"{mod.__module__}.{mod.__name__}": mod for mod in mod_reg.values()
-    }
-    module: Type[ModuleBase] = mod_reg.get(args.module, mod_pkg_to_mod.get(args.module))
-    error.value_check(
-        "<COR03876205E>",
-        module is not None,
-        "Unable to find module {} to train",
-        args.module,
-    )
+    try:
+        mod_reg = module_registry()
+        mod_pkg_to_mod = {
+            f"{mod.__module__}.{mod.__name__}": mod for mod in mod_reg.values()
+        }
+        module: Type[ModuleBase] = mod_reg.get(
+            args.module, mod_pkg_to_mod.get(args.module)
+        )
+        error.value_check(
+            "<COR03876205E>",
+            module is not None,
+            "Unable to find module {} to train",
+            args.module,
+        )
+    except Exception:
+        message = "Unable to find module {} to train".format(args.module)
+        log.warning(
+            {
+                "log_code": "<COR17476539E>",
+                "message": message,
+                "stack_trace": traceback.format_exc,
+            },
+            exc_info=True,
+        )
+        write_termination_log(message)
+        exit(USER_ERROR_EXIT_CODE)
 
     # Read training kwargs
     try:
@@ -141,57 +222,113 @@ def main() -> int:
                 training_kwargs = json.load(handle)
         else:
             training_kwargs = json.loads(args.training_kwargs)
+
+        # Convert datatypes to match the training API
+        training_service = ServicePackageFactory.get_service_package(
+            ServicePackageFactory.ServiceType.TRAINING,
+        )
+        train_rpcs = [
+            rpc
+            for rpc in training_service.caikit_rpcs.values()
+            if rpc.module_list == [module]
+        ]
+        error.value_check(
+            "<COR11978965E>",
+            len(train_rpcs) == 1,
+            "Unable to find a unique train signature",
+        )
+        package_name = get_service_package_name(
+            ServicePackageFactory.ServiceType.TRAINING
+        )
+        train_rpc_req = (
+            train_rpcs[0].create_request_data_model(package_name).get_proto_class()
+        )
+        request_proto = json_format.Parse(
+            json.dumps({"parameters": training_kwargs}),
+            train_rpc_req(),
+        )
+        req_kwargs = build_caikit_library_request_dict(
+            request_proto.parameters, module.TRAIN_SIGNATURE
+        )
+        train_kwargs.update(req_kwargs)
+        log.debug3("All train kwargs: %s", train_kwargs)
     except json.decoder.JSONDecodeError:
-        log.error(
-            "<COR65834760E>",
-            "--training-kwargs must be valid json or point to a valid json file",
+        message = "training-kwargs must be valid json or point to a valid json file"
+        log.warning(
+            {
+                "log_code": "<COR65834760E>",
+                "message": message,
+                "stack_trace": traceback.format_exc(),
+            },
+            exc_info=True,
+        )
+        write_termination_log(message)
+        exit(USER_ERROR_EXIT_CODE)
+    except ValueError as e:
+        message = f"Invalid value for one or more input parameters: {e}"
+        log.warning(
+            {
+                "log_code": "<COR65474760E>",
+                "message": message,
+                "stack_trace": traceback.format_exc(),
+            },
+            exc_info=True,
         )
-        raise
+    except Exception:
+        message = "Exception encountered when attempting to parse input parameters"
+        log.warning(
+            {
+                "log_code": "<COR17776549E>",
+                "message": message,
+                "stack_trace": traceback.format_exc(),
+            },
+            exc_info=True,
+        )
+        write_termination_log(message)
+        exit(USER_ERROR_EXIT_CODE)
 
-    # Convert datatypes to match the training API
-    training_service = ServicePackageFactory.get_service_package(
-        ServicePackageFactory.ServiceType.TRAINING,
-    )
-    train_rpcs = [
-        rpc
-        for rpc in training_service.caikit_rpcs.values()
-        if rpc.module_list == [module]
-    ]
-    error.value_check(
-        "<COR11978965E>",
-        len(train_rpcs) == 1,
-        "Unable to find a unique train signature",
-    )
-    package_name = get_service_package_name(ServicePackageFactory.ServiceType.TRAINING)
-    train_rpc_req = (
-        train_rpcs[0].create_request_data_model(package_name).get_proto_class()
-    )
-    request_proto = json_format.Parse(
-        json.dumps({"parameters": training_kwargs}),
-        train_rpc_req(),
-    )
-    req_kwargs = build_caikit_library_request_dict(
-        request_proto.parameters, module.TRAIN_SIGNATURE
-    )
-    train_kwargs.update(req_kwargs)
-    log.debug3("All train kwargs: %s", train_kwargs)
-
-    # Run the training
-    with alog.ContextTimer(
-        log.info,
-        "Finished training %s in: ",
-        args.model_name,
-    ):
-        future = train(module, wait=True, **train_kwargs)
-        info = future.get_info()
-        if info.status == TrainingStatus.COMPLETED:
-            log.info("Training finished successfully")
-            return 0
-        else:
-            log.error("Training finished unsuccessfully")
-            for err in info.errors or []:
-                log.error(err)
-            return 1
+    try:
+        # Run the training
+        with alog.ContextTimer(
+            log.info,
+            "Finished training %s in: ",
+            args.model_name,
+        ):
+            future = train(module, wait=True, **train_kwargs)
+            info = future.get_info()
+            if info.status == TrainingStatus.COMPLETED:
+                log.info("Training finished successfully")
+                return 0
+            else:
+                log.error("Training finished unsuccessfully")
+                for err in info.errors or []:
+                    log.error(err)
+                write_termination_log("Training finished unsuccessfully")
+                exit(INTERNAL_ERROR_EXIT_CODE)
+    except MemoryError:
+        message = "OOM error during training"
+        log.warning(
+            {
+                "log_code": "<COR04280062E>",
+                "message": message,
+                "stack_trace": traceback.format_exc(),
+            },
+            exc_info=True,
+        )
+        write_termination_log(message)
+        exit(INTERNAL_ERROR_EXIT_CODE)
+    except Exception:
+        message = "Unhandled exception during training"
+        log.warning(
+            {
+                "log_code": "<COR04280062E>",
+                "message": message,
+                "stack_trace": traceback.format_exc(),
+            },
+            exc_info=True,
+        )
+        write_termination_log(message)
+        exit(INTERNAL_ERROR_EXIT_CODE)
 
 
 if __name__ == "__main__":


Test patch:
diff --git a/tests/runtime/test_train.py b/tests/runtime/test_train.py
index ee18aaaf..c6db67c9 100644
--- a/tests/runtime/test_train.py
+++ b/tests/runtime/test_train.py
@@ -16,6 +16,7 @@ import pytest
 
 # Local
 from caikit.core.registries import module_registry
+from caikit.runtime import train
 from caikit.runtime.train import main
 from sample_lib.modules import SampleModule
 from tests.conftest import reset_module_registry, temp_config
@@ -215,8 +216,10 @@ def test_invalid_json():
         "--training-kwargs",
         "{invalid json",
     ):
-        with pytest.raises(json.decoder.JSONDecodeError):
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
             main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == train.USER_ERROR_EXIT_CODE
 
 
 def test_failed_training():
@@ -232,4 +235,94 @@ def test_failed_training():
         "--training-kwargs",
         json.dumps(training_kwargs),
     ):
-        assert main() == 1
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == train.INTERNAL_ERROR_EXIT_CODE
+
+
+def test_bad_module():
+    """Make sure that a non-zero exit code is returned if an invalid module is provided"""
+    model_name = "my-model"
+    training_kwargs = copy.deepcopy(SAMPLE_TRAIN_KWARGS)
+    with sys_argv(
+        "--module",
+        "this.is.a.bad.module",
+        "--model-name",
+        model_name,
+        "--training-kwargs",
+        json.dumps(training_kwargs),
+    ):
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == train.USER_ERROR_EXIT_CODE
+
+
+def test_no_module_provided():
+    """Make sure that a non-zero exit code is returned if an invalid module is provided"""
+    model_name = "my-model"
+    training_kwargs = copy.deepcopy(SAMPLE_TRAIN_KWARGS)
+    with sys_argv(
+        "--model-name",
+        model_name,
+        "--training-kwargs",
+        json.dumps(training_kwargs),
+    ):
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == 1
+
+
+def test_blank_kwargs():
+    """Make sure that a non-zero exit code is returned if kwargs are blank"""
+    model_name = "my-model"
+    with sys_argv(
+        "--model-name",
+        SAMPLE_MODULE,
+        "--model-name",
+        model_name,
+    ):
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == train.USER_ERROR_EXIT_CODE
+
+
+def test_empty_module_name():
+    """Test handling of empty module parameter"""
+    model_name = "my-model"
+    with sys_argv(
+        "--module",
+        "",
+        "--model-name",
+        model_name,
+        "--training-kwargs",
+        json.dumps(SAMPLE_TRAIN_KWARGS),
+    ):
+        with pytest.raises(SystemExit) as e:
+            main()
+        assert e.value.code == train.USER_ERROR_EXIT_CODE
+
+
+def test_non_existent_save_path():
+    """Test with a non-existent save path"""
+    # We cannot verify save path ahead of time, so if it is unable
+    # to be written to, the training will fail with a system error
+    model_name = "my-model"
+    non_existent_path = "/path/that/does/not/exist"
+    with sys_argv(
+        "--module",
+        SAMPLE_MODULE,
+        "--model-name",
+        model_name,
+        "--save-path",
+        non_existent_path,
+        "--training-kwargs",
+        json.dumps(SAMPLE_TRAIN_KWARGS),
+    ):
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == train.INTERNAL_ERROR_EXIT_CODE