Repository: caikit/caikit
Pull request #382 (2023-07-28 20:39:35)

:sparkles: Add model name to Trainer API
<!--  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**:

For #381 

This PR adds `model_name` as an optional top-level arg to the model trainer's `.train` interface.

This allows model trainer implementations to construct their own save paths based on
- the base path given
- the training id
- the model name

The abstract base class is updated so that the save path is constructed as
```
${save_path}/${training_id}/${model_name}
```

This should not be a breaking change since the parameter is optional with a default of none

(except maybe how the .train function may have keyword collisions with actual training kwargs but... /shrug)

**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/core/model_management/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index 1fc2ace4..988751bf 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -69,6 +69,7 @@ class LocalModelTrainer(ModelTrainerBase):
             *args,
             save_path: Optional[Union[str, S3Path]],
             save_with_id: bool,
+            model_name: Optional[str],
             external_training_id: Optional[str],
             use_subprocess: bool,
             **kwargs,
@@ -78,15 +79,16 @@ class LocalModelTrainer(ModelTrainerBase):
                 training_id=str(uuid.uuid4()),
                 save_with_id=save_with_id,
                 save_path=save_path,
+                model_name=model_name,
             )
-
-            # If an external id is given, use it explicitly
+            # 🌶️🌶️🌶️ For the external training id override, we don't want to include the
+            # reversible hash bit on the id. Therefore, we have to re-do the id and save
+            # path stuff done in the super() init here instead
             if external_training_id is not None:
+                # No extra hash bit
                 self._id = external_training_id
                 self._save_path = self._save_path_with_id(
-                    save_path,
-                    save_with_id,
-                    external_training_id,
+                    save_path, save_with_id, external_training_id, model_name
                 )
 
             self._module_class = module_class
@@ -257,6 +259,7 @@ class LocalModelTrainer(ModelTrainerBase):
         save_path: Optional[str] = None,
         save_with_id: bool = False,
         external_training_id: Optional[str] = None,
+        model_name: Optional[str] = None,
         **kwargs,
     ) -> "LocalModelFuture":
         """Start training the given module and return a future to the trained
@@ -274,6 +277,7 @@ class LocalModelTrainer(ModelTrainerBase):
             save_with_id=save_with_id,
             external_training_id=external_training_id,
             use_subprocess=self._use_subprocess,
+            model_name=model_name,
             **kwargs,
         )
 
diff --git a/caikit/core/model_management/model_trainer_base.py b/caikit/core/model_management/model_trainer_base.py
index f08238d4..ee3992ae 100644
--- a/caikit/core/model_management/model_trainer_base.py
+++ b/caikit/core/model_management/model_trainer_base.py
@@ -61,6 +61,7 @@ class ModelTrainerBase(FactoryConstructible):
             training_id: str,
             save_with_id: bool,
             save_path: Optional[Union[str, S3Path]],
+            model_name: Optional[str] = None,
         ):
             # Trainers should deal with an S3 ref first and not pass it along here
             if save_path and isinstance(save_path, S3Path):
@@ -72,6 +73,7 @@ class ModelTrainerBase(FactoryConstructible):
                 save_path,
                 save_with_id,
                 self._id,
+                model_name,
             )
 
         @property
@@ -118,21 +120,25 @@ class ModelTrainerBase(FactoryConstructible):
             save_path: Optional[str],
             save_with_id: bool,
             training_id: str,
+            model_name: Optional[str],
         ) -> Optional[str]:
             """If asked to save_with_id, child classes should use this shared
             utility to construct the final save path
             """
             if save_path is None:
                 return save_path
-            if not save_with_id or training_id in save_path:
-                return save_path
 
-            # If told to save with the ID in the path, inject it right before the
-            # final portion of the path which is assumed to be the model ID.
-            path_parts = os.path.split(save_path)
-            return os.path.join(
-                *list(path_parts[:-1] + (training_id,) + path_parts[-1:])
-            )
+            final_path_parts = [save_path]
+            # If told to save with the ID in the path, inject it before the
+            # model name.
+            if save_with_id and training_id not in save_path:
+                # (Don't inject training id if its already in the path)
+                final_path_parts.append(training_id)
+
+            if model_name and model_name not in save_path:
+                final_path_parts.append(model_name)
+
+            return os.path.join(*final_path_parts)
 
     @abc.abstractmethod
     def train(
@@ -141,6 +147,7 @@ class ModelTrainerBase(FactoryConstructible):
         *args,
         save_path: Optional[Union[str, S3Path]] = None,
         save_with_id: bool = False,
+        model_name: Optional[str] = None,
         **kwargs,
     ) -> "ModelFutureBase":
         """Start training the given module and return a future to the trained
diff --git a/caikit/core/model_manager.py b/caikit/core/model_manager.py
index 12992cb9..7cd34ff0 100644
--- a/caikit/core/model_manager.py
+++ b/caikit/core/model_manager.py
@@ -98,6 +98,7 @@ class ModelManager:
         trainer: Union[str, ModelTrainerBase] = "default",
         save_path: Optional[Union[str, S3Path]] = None,
         save_with_id: bool = False,
+        model_name: Optional[str] = None,
         wait: bool = False,
         **kwargs,
     ) -> ModelTrainerBase.ModelFutureBase:
@@ -120,11 +121,13 @@ class ModelManager:
             trainer (Union[str, ModelTrainerBase]): The trainer to use. If given
                 as a string, this is a key in the global config at
                 model_management.trainers.
-            save_path (Optional[Union[str, S3Path]]): Path where the model should be
+            save_path (Optional[Union[str, S3Path]]): Base path where the model should be
                 saved (may be relative to a remote trainer's filesystem, or link to S3
                 storage)
             save_with_id (bool): Inject the training ID into the save path for
                 the output model
+            model_name (Optional[str]): Name of model that will be appended
+                to the end of the save_path
             wait (bool): Wait for training to complete before returning
             **kwargs: Additional keyword arguments to pass through to the
                 modules's train function
@@ -155,6 +158,7 @@ class ModelManager:
                 *args,
                 save_path=save_path,
                 save_with_id=save_with_id,
+                model_name=model_name,
                 **kwargs,
             )
             log.debug(
diff --git a/caikit/runtime/servicers/global_train_servicer.py b/caikit/runtime/servicers/global_train_servicer.py
index 893b5afa..175700f1 100644
--- a/caikit/runtime/servicers/global_train_servicer.py
+++ b/caikit/runtime/servicers/global_train_servicer.py
@@ -14,7 +14,6 @@
 # Standard
 from importlib.metadata import version
 from typing import Optional, Type, Union
-import os
 import traceback
 
 # Third Party
@@ -202,10 +201,10 @@ class GlobalTrainServicer:
             # If we got an S3 storage link, just pass that along to the trainer
             model_path: S3Path = request_data_model.output_path
         else:
-            # Otherwise, append the model name to the specified output directory
-            model_path: str = self._get_model_path(
-                training_output_dir, request_data_model.model_name
-            )
+            # Otherwise, use either:
+            # 1. The provided `training_output_dir` here, or
+            # 2. The configured `runtime.training.output_dir`
+            model_path: str = training_output_dir or self.training_output_dir
 
         # Build the full set of kwargs for the train call
         kwargs.update(
@@ -213,6 +212,7 @@ class GlobalTrainServicer:
                 "module": module,
                 "save_path": model_path,
                 "save_with_id": self.save_with_id,
+                "model_name": request_data_model.model_name,
                 **build_caikit_library_request_dict(request, module.TRAIN_SIGNATURE),
             }
         )
@@ -277,19 +277,6 @@ class GlobalTrainServicer:
             training_id=model_future.id,
         )
 
-    def _get_model_path(
-        self,
-        training_output_dir: Optional[str],
-        model_name: str,
-    ) -> str:
-        """Get the right output path for a given model"""
-        base_dir = (
-            training_output_dir
-            if training_output_dir is not None
-            else self.training_output_dir
-        )
-        return os.path.join(base_dir, model_name)
-
     def _load_trained_model(self, model_name: str, model_path: str):
         log.debug("Autoloading trained model %s", model_name)
         self._model_manager.load_model(


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 101cb355..dfd9ce0d 100644
--- a/tests/core/model_management/test_local_model_trainer.py
+++ b/tests/core/model_management/test_local_model_trainer.py
@@ -124,6 +124,44 @@ def test_save_with_id(trainer_type_cfg, save_path):
     assert os.path.exists(model_future.save_path)
 
 
+def test_save_with_id_and_model_name(trainer_type_cfg, save_path):
+    """Test that saving with the training id and model name
+    correctly injects the ID and name in the save path
+    """
+    trainer = local_trainer(**trainer_type_cfg)
+    model_future = trainer.train(
+        SampleModule,
+        training_data=DataStream.from_iterable([]),
+        save_path=save_path,
+        save_with_id=True,
+        model_name="abc",
+    )
+    model_future.wait()
+    assert model_future.save_path != save_path
+    assert model_future.id in model_future.save_path
+    assert "abc" in model_future.save_path
+    assert os.path.exists(model_future.save_path)
+
+
+def test_save_with_model_name(trainer_type_cfg, save_path):
+    """Test that saving with the model name correctly
+    injects the model name in the save path
+    """
+    trainer = local_trainer(**trainer_type_cfg)
+    model_future = trainer.train(
+        SampleModule,
+        training_data=DataStream.from_iterable([]),
+        save_path=save_path,
+        save_with_id=False,
+        model_name="abc",
+    )
+    model_future.wait()
+    assert model_future.save_path != save_path
+    assert model_future.id not in model_future.save_path
+    assert "abc" in model_future.save_path
+    assert os.path.exists(model_future.save_path)
+
+
 def test_cancel_clean_termination(trainer_type_cfg):
     """Test that cancelling an in-progress training successfully destroys the
     training when the training is run in a way that can be shut down cleanly
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index 0fa6fb5e..2801436b 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -16,6 +16,7 @@ from contextlib import contextmanager
 from tempfile import TemporaryDirectory
 from unittest.mock import patch
 import multiprocessing
+import os
 import threading
 import time
 
@@ -114,7 +115,13 @@ def test_global_train_sample_task(
         training_response.training_id,
     )
 
-    result = MODEL_MANAGER.get_model_future(training_response.training_id).load()
+    model_future = MODEL_MANAGER.get_model_future(training_response.training_id)
+    expected_save_path_bits = os.path.join(
+        training_response.training_id, training_response.model_name
+    )
+    assert expected_save_path_bits in model_future.save_path
+
+    result = model_future.load()
     assert result.batch_size == 42
     assert (
         result.MODULE_CLASS
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 473183ab..03ec0f58 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -116,6 +116,7 @@ def test_model_train(runtime_grpc_server):
         runtime_grpc_server.make_local_channel()
     )
     training_id = str(uuid.uuid4())
+    model_name = "abc"
     model_train_request = process_pb2.ProcessRequest(
         trainingID=training_id,
         customTrainingID=str(uuid.uuid4()),
@@ -123,7 +124,7 @@ def test_model_train(runtime_grpc_server):
             "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
             "training_params": json.dumps(
                 {
-                    "model_name": "abc",
+                    "model_name": model_name,
                     "training_data": {
                         "jsondata": {
                             "data": [
@@ -170,7 +171,12 @@ def test_model_train(runtime_grpc_server):
     assert response.state == TrainingStatus.COMPLETED.value
 
     # Make sure we wait for training to finish
-    result = MODEL_MANAGER.get_model_future(response.training_id).load()
+    model_future = MODEL_MANAGER.get_model_future(response.training_id)
+    result = model_future.load()
+
+    # Make sure we put the path bits in the right order: base/training_id/model_name
+    training_id_and_model_name_path_bit = os.path.join(training_id, model_name)
+    assert training_id_and_model_name_path_bit in model_future.save_path
 
     assert (
         result.MODULE_CLASS