Repository: caikit/caikit (python)

Issue:
Retrying local training with external ID raises assertion
## Describe the bug

The assertion [here](https://github.com/caikit/caikit/blob/main/caikit/core/model_management/local_model_trainer.py#L313) in `LocalModelTrainer` can be triggered during a valid execution scenario where a training has failed and is being retried using the same external ID.

## Sample Code

```py
import aconfig
import alog
from caikit.core.toolkit.logging import configure
from caikit.core.model_management.local_model_trainer import LocalModelTrainer
from caikit.core import ModuleBase, module

log = alog.use_channel("REPRO")
configure()

@module("fail-train-once", "Module that will fail training once", "0.0.0")
class FailTrainOnce(ModuleBase):
    _calls = 0
    def __init__(self, val: int):
        self.val = val
    @classmethod
    def train(cls, arg: int):
        cls._calls = cls._calls + 1
        if cls._calls == 1:
            raise RuntimeError("Yikes!")
        return cls(arg)


trainer = LocalModelTrainer(aconfig.Config({}), "")
training_id = "my-training"

# First try should fail
try:
    log.info("First Try")
    model_future = trainer.train(FailTrainOnce, 42, external_training_id=training_id)
    model_future.load()
    raise AssertionError("Shouldn't get here")
except RuntimeError:
    # Second time should succeed
    log.info("Second Try")
    model_future = trainer.train(FailTrainOnce, 42, external_training_id=training_id)
    model = model_future.load()

assert model
```

## Expected behavior

The final `assert model` should pass because the first attempt reached a terminal state, so the second attempt is a retry

## Observed behavior

The call to `trainer.train` raises the linked `AssertionError` because the external ID is being reused


Pull request:
UUID asssertion on terminal 656
**What this PR does / why we need it**:

Closes #656 

This PR handles the edge case where an external job manager attempts to retry a training job against a running instance of the `LocalModelTrainer` using the same ID as a previously failed run.

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [x] this PR has been tested for backwards compatibility


Patch:
diff --git a/caikit/core/model_management/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index b6758766..553461f5 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -293,6 +293,18 @@ class LocalModelTrainer(ModelTrainerBase):
             log.debug2("Subprocess wrapped models: %s", wrapped_models.keys())
             kwargs.update(wrapped_models)
 
+        # If there's an external ID, make sure it's not currently running before
+        # launching the job
+        if external_training_id and (
+            current_future := self._futures.get(external_training_id)
+        ):
+            error.value_check(
+                "<COR79850561E>",
+                current_future.get_info().status.is_terminal,
+                "Cannot restart training {} that is currently running",
+                external_training_id,
+            )
+
         # Create the new future
         model_future = self.LocalModelFuture(
             self._instance_name,
@@ -309,9 +321,13 @@ class LocalModelTrainer(ModelTrainerBase):
 
         # Lock the global futures dict and add it to the dict
         with self._futures_lock:
-            assert (
-                model_future.id not in self._futures
-            ), f"UUID collision for model future {model_future.id}"
+            if current_future := self._futures.get(model_future.id):
+                error.value_check(
+                    "<COR35431427E>",
+                    current_future.get_info().status.is_terminal,
+                    "UUID collision for model future {}",
+                    model_future.id,
+                )
             self._futures[model_future.id] = model_future
 
         # Return the future


Test patch:
diff --git a/tests/core/model_management/test_local_model_trainer.py b/tests/core/model_management/test_local_model_trainer.py
index c84c96d6..2997e8b2 100644
--- a/tests/core/model_management/test_local_model_trainer.py
+++ b/tests/core/model_management/test_local_model_trainer.py
@@ -29,6 +29,7 @@ import aconfig
 
 # Local
 from caikit.config import get_config
+from caikit.core import ModuleBase
 from caikit.core.data_model import DataStream, TrainingStatus
 from caikit.core.exceptions.caikit_core_exception import CaikitCoreException
 from caikit.core.model_management.local_model_trainer import LocalModelTrainer
@@ -62,6 +63,28 @@ def get_event(cfg: dict):
     return threading.Event()
 
 
+class FailTrainOnce(ModuleBase):
+    """Dummy module that will fail training the first time"""
+
+    _calls = 0
+
+    @classmethod
+    def train(cls):
+        cls._calls = cls._calls + 1
+        if cls._calls == 1:
+            raise RuntimeError("Yikes!")
+        return cls()
+
+
+class WaitTrain(ModuleBase):
+    """Dummy module that will block training on an event"""
+
+    @classmethod
+    def train(cls, wait_event: threading.Event):
+        wait_event.wait()
+        return cls()
+
+
 ## Tests #######################################################################
 
 
@@ -266,3 +289,39 @@ def test_get_into_return_error(trainer_type_cfg):
     assert isinstance(model_future.get_info().errors, list)
     assert isinstance(model_future.get_info().errors[0], ValueError)
     assert str(model_future.get_info().errors[0]) == "Batch size of 999 is not allowed!"
+
+
+def test_retry_duplicate_external_id():
+    """Test that a training can be retried safely reusing an external ID"""
+    trainer = local_trainer()
+    training_id = "my-training"
+
+    # First try should fail
+    try:
+        model_future = trainer.train(FailTrainOnce, external_training_id=training_id)
+        model_future.load()
+        raise AssertionError("Shouldn't get here")
+    except RuntimeError:
+        # Second time should succeed
+        model_future = trainer.train(FailTrainOnce, external_training_id=training_id)
+        assert model_future.load()
+
+
+def test_duplicate_external_id_cannot_restart_while_running():
+    """Make sure that if a training is actively running, it cannot be replaced
+    by a rerun
+    """
+    trainer = local_trainer()
+    training_id = "my-training"
+    wait_event = threading.Event()
+    model_future = trainer.train(
+        WaitTrain, wait_event, external_training_id=training_id
+    )
+    try:
+        with pytest.raises(ValueError, match="Cannot restart training.*"):
+            trainer.train(WaitTrain, wait_event, external_training_id=training_id)
+
+        assert trainer.get_model_future(training_id) is model_future
+    finally:
+        wait_event.set()
+        model_future.wait()