Repository: caikit/caikit
Pull request #733 (2024-06-05 17:12:59)

Add .complete file, env var for termination log
Adds .complete file to indicate to client who are watching the filesystem that we are done writing all of the files to disk. Also makes the termination log path configurable by environment variable.


Patch:
diff --git a/caikit/runtime/train.py b/caikit/runtime/train.py
index fe26990c..1f2f95c2 100644
--- a/caikit/runtime/train.py
+++ b/caikit/runtime/train.py
@@ -18,6 +18,7 @@ job using caikit.core.train
 """
 
 # Standard
+from pathlib import Path
 from typing import Type
 import argparse
 import importlib
@@ -65,7 +66,9 @@ class TrainArgumentParser(argparse.ArgumentParser):
         raise ArgumentParserError(f"{self.prog}: error: {message}")
 
 
-def write_termination_log(text: str, log_file="/dev/termination-log"):
+def write_termination_log(text: str, log_file: str, enabled: bool):
+    if not enabled:
+        return
     try:
         with open(log_file, "a") as handle:
             handle.write(text)
@@ -77,10 +80,38 @@ def write_termination_log(text: str, log_file="/dev/termination-log"):
         )
 
 
+# Final tasks before exiting the container
+def exit_complete(
+    exit_code: int,
+    save_path: str,
+    message: str,
+    termination_log_file: str,
+    enable_termination_log: bool,
+):
+    if exit_code != 0:
+        write_termination_log(message, termination_log_file, enable_termination_log)
+
+    if save_path:
+        try:
+            complete_path = os.path.join(save_path, ".complete")
+            log.info("Creating completion file at: %s", complete_path)
+            Path(complete_path).touch()
+        except Exception as e:
+            log.warning("Unable to write completion file due to execption: %s", e)
+
+    exit(exit_code)
+
+
 def main() -> int:
     """Main entrypoint for running training jobs"""
     parser = TrainArgumentParser(description=__doc__)
 
+    # Set default values for termination log incase parsing the arguments fail later on
+    enable_termination_log = os.environ.get("ENABLE_TERMINATION_LOG", True)
+    termination_log_file = os.environ.get(
+        "TERMINATION_LOG_FILE", "/dev/termination-log"
+    )
+
     # Required Args
     parser.add_argument(
         "--training-kwargs",
@@ -130,15 +161,30 @@ def main() -> int:
     parser.add_argument(
         "--termination-log-file",
         "-f",
-        action="store_true",
-        default="/dev/termination-log",
+        default=termination_log_file,
         help="Location of where to write a termination error message",
     )
+    parser.add_argument(
+        "--enable-termination-log",
+        "-e",
+        default=enable_termination_log,
+        help="Whether to enable writing to termination log when training fails",
+    )
 
     try:
         args = parser.parse_args()
         config_logging()
 
+        # Modify termination log variables if parsed
+        # Previously we grabbed the values from env variables (if present)
+        # Here, we allow overriding it with the parser values
+        # If the parser throws an exception parsing any of the args, the values
+        # captured in previous sections will be used.
+        if args.enable_termination_log:
+            enable_termination_log = args.enable_termination_log
+        if args.termination_log_file:
+            termination_log_file = args.termination_log_file
+
         # Initialize top-level kwargs
         train_kwargs = {
             "save_path": args.save_path,
@@ -147,7 +193,6 @@ def main() -> int:
         }
         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(
@@ -158,8 +203,14 @@ def main() -> int:
             },
             exc_info=True,
         )
-        write_termination_log(message)
-        exit(USER_ERROR_EXIT_CODE)
+        # We couldn't parse args, so cannot not pass save_path in
+        exit_complete(
+            USER_ERROR_EXIT_CODE,
+            None,
+            message,
+            termination_log_file,
+            enable_termination_log,
+        )
 
     # Import libraries to register modules
     try:
@@ -176,8 +227,13 @@ def main() -> int:
             },
             exc_info=True,
         )
-        write_termination_log(message)
-        exit(USER_ERROR_EXIT_CODE)
+        exit_complete(
+            USER_ERROR_EXIT_CODE,
+            args.save_path,
+            message,
+            termination_log_file,
+            enable_termination_log,
+        )
 
     # Try to import the root library of the provided module. It's ok if this
     # fails since the module may be a UID
@@ -212,8 +268,13 @@ def main() -> int:
             },
             exc_info=True,
         )
-        write_termination_log(message)
-        exit(USER_ERROR_EXIT_CODE)
+        exit_complete(
+            USER_ERROR_EXIT_CODE,
+            args.save_path,
+            message,
+            termination_log_file,
+            enable_termination_log,
+        )
 
     # Read training kwargs
     try:
@@ -262,8 +323,13 @@ def main() -> int:
             },
             exc_info=True,
         )
-        write_termination_log(message)
-        exit(USER_ERROR_EXIT_CODE)
+        exit_complete(
+            USER_ERROR_EXIT_CODE,
+            args.save_path,
+            message,
+            termination_log_file,
+            enable_termination_log,
+        )
     except ValueError as e:
         message = f"Invalid value for one or more input parameters: {e}"
         log.warning(
@@ -284,8 +350,13 @@ def main() -> int:
             },
             exc_info=True,
         )
-        write_termination_log(message)
-        exit(USER_ERROR_EXIT_CODE)
+        exit_complete(
+            USER_ERROR_EXIT_CODE,
+            args.save_path,
+            message,
+            termination_log_file,
+            enable_termination_log,
+        )
 
     try:
         # Run the training
@@ -295,16 +366,32 @@ def main() -> int:
             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
+                log.info(
+                    {
+                        "log_code": "<COR74526958I>",
+                        "message": "Training finished successfully",
+                    }
+                )
+                exit_complete(0, args.save_path, None, None, None)
             else:
-                log.error("Training finished unsuccessfully")
+                log.warning(
+                    {
+                        "log_code:": "<COR72523958E>",
+                        "message": "Training finished unsuccessfully",
+                    }
+                )
                 for err in info.errors or []:
                     log.error(err)
-                write_termination_log("Training finished unsuccessfully")
-                exit(INTERNAL_ERROR_EXIT_CODE)
+                exit_complete(
+                    INTERNAL_ERROR_EXIT_CODE,
+                    args.save_path,
+                    "Training finished unsuccessfully",
+                    termination_log_file,
+                    enable_termination_log,
+                )
     except MemoryError:
         message = "OOM error during training"
         log.warning(
@@ -315,8 +402,13 @@ def main() -> int:
             },
             exc_info=True,
         )
-        write_termination_log(message)
-        exit(INTERNAL_ERROR_EXIT_CODE)
+        exit_complete(
+            INTERNAL_ERROR_EXIT_CODE,
+            args.save_path,
+            message,
+            termination_log_file,
+            enable_termination_log,
+        )
     except Exception:
         message = "Unhandled exception during training"
         log.warning(
@@ -327,8 +419,13 @@ def main() -> int:
             },
             exc_info=True,
         )
-        write_termination_log(message)
-        exit(INTERNAL_ERROR_EXIT_CODE)
+        exit_complete(
+            INTERNAL_ERROR_EXIT_CODE,
+            args.save_path,
+            message,
+            termination_log_file,
+            enable_termination_log,
+        )
 
 
 if __name__ == "__main__":


Test patch:
diff --git a/tests/runtime/test_train.py b/tests/runtime/test_train.py
index c6db67c9..3cd29ed5 100644
--- a/tests/runtime/test_train.py
+++ b/tests/runtime/test_train.py
@@ -4,6 +4,7 @@ Unit tests for the train script entrypoint
 
 # Standard
 from contextlib import contextmanager
+from pathlib import Path
 from unittest import mock
 import copy
 import json
@@ -64,10 +65,14 @@ def test_train_sample_module(workdir):
         "--training-kwargs",
         json.dumps(SAMPLE_TRAIN_KWARGS),
     ):
-        assert main() == 0
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == 0
         model_dir = os.path.join(workdir, model_name)
         assert os.path.isdir(model_dir)
         assert os.path.isfile(os.path.join(model_dir, "config.yml"))
+        assert os.path.isfile(os.path.join(workdir, ".complete"))
 
 
 def test_train_from_file(workdir):
@@ -86,15 +91,20 @@ def test_train_from_file(workdir):
         "--training-kwargs",
         train_kwargs_file,
     ):
-        assert main() == 0
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == 0
         model_dir = os.path.join(workdir, model_name)
         assert os.path.isdir(model_dir)
         assert os.path.isfile(os.path.join(model_dir, "config.yml"))
+        assert os.path.isfile(os.path.join(workdir, ".complete"))
 
 
 def test_train_module_uid(workdir):
     """Test referencing the module by its UID"""
     model_name = "my-model"
+    log_path = os.path.join(workdir, "termination-log")
     with sys_argv(
         "--module",
         SampleModule.MODULE_ID,
@@ -104,11 +114,19 @@ def test_train_module_uid(workdir):
         workdir,
         "--training-kwargs",
         json.dumps(SAMPLE_TRAIN_KWARGS),
+        "--termination-log-file",
+        log_path,
     ):
-        assert main() == 0
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == 0
         model_dir = os.path.join(workdir, model_name)
         assert os.path.isdir(model_dir)
         assert os.path.isfile(os.path.join(model_dir, "config.yml"))
+        assert os.path.isfile(os.path.join(workdir, ".complete"))
+        # Ensure termination log doesn't exist, which indicates error
+        assert not os.path.isfile(log_path)
 
 
 def test_train_save_with_id(workdir):
@@ -125,7 +143,10 @@ def test_train_save_with_id(workdir):
         json.dumps(SAMPLE_TRAIN_KWARGS),
         "--save-with-id",
     ):
-        assert main() == 0
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            main()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == 0
         flat_model_dir = os.path.join(workdir, model_name)
         assert not os.path.isdir(flat_model_dir)
         dirs = list(
@@ -136,6 +157,7 @@ def test_train_save_with_id(workdir):
         )
         assert len(dirs) == 1
         assert os.path.isfile(os.path.join(dirs[0], model_name, "config.yml"))
+        assert os.path.isfile(os.path.join(dirs[0], workdir, ".complete"))
 
 
 def test_train_non_default_trainer(workdir):
@@ -172,10 +194,14 @@ def test_train_non_default_trainer(workdir):
             "--trainer",
             other_trainer,
         ):
-            assert main() == 0
+            with pytest.raises(SystemExit) as pytest_wrapped_e:
+                main()
+            assert pytest_wrapped_e.type == SystemExit
+            assert pytest_wrapped_e.value.code == 0
             model_dir = os.path.join(workdir, model_name)
             assert os.path.isdir(model_dir)
             assert os.path.isfile(os.path.join(model_dir, "config.yml"))
+            assert os.path.isfile(os.path.join(workdir, ".complete"))
 
 
 def test_train_import_library(workdir, reset_module_registry):
@@ -194,7 +220,10 @@ def test_train_import_library(workdir, reset_module_registry):
             "--library",
             "sample_lib",
         ):
-            assert main() == 0
+            with pytest.raises(SystemExit) as pytest_wrapped_e:
+                main()
+            assert pytest_wrapped_e.type == SystemExit
+            assert pytest_wrapped_e.value.code == 0
             model_dir = os.path.join(workdir, model_name)
             assert os.path.isdir(model_dir)
             assert os.path.isfile(os.path.join(model_dir, "config.yml"))
@@ -205,7 +234,7 @@ def test_train_import_library(workdir, reset_module_registry):
             ]
 
 
-def test_invalid_json():
+def test_invalid_json(workdir):
     """Make sure that an exception is raised for invalid json"""
     model_name = "my-model"
     with sys_argv(
@@ -216,15 +245,17 @@ def test_invalid_json():
         "--training-kwargs",
         "{invalid json",
     ):
+        log_path = os.path.join(workdir, "termination-log")
         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():
+def test_failed_training(workdir):
     """Make sure that a non-zero exit code is returned if training fails"""
     model_name = "my-model"
+    log_path = os.path.join(workdir, "termination-log")
     training_kwargs = copy.deepcopy(SAMPLE_TRAIN_KWARGS)
     training_kwargs["batch_size"] = SampleModule.POISON_PILL_BATCH_SIZE
     with sys_argv(
@@ -234,11 +265,14 @@ def test_failed_training():
         model_name,
         "--training-kwargs",
         json.dumps(training_kwargs),
+        "--termination-log-file",
+        log_path,
     ):
         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
+        assert os.path.isfile(log_path)
 
 
 def test_bad_module():