Repository: caikit/caikit
Pull request #481 (2023-09-19 21:43:49)

Model train save with custom
**What this PR does / why we need it**:

This PR modifies the `model_train_servicer` to use the inbound `ProcessRequest`'s `customTrainingID` if given. It is needed so that the fully-external `caikit` ID can be used in the save path.

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [X] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/runtime/servicers/global_train_servicer.py b/caikit/runtime/servicers/global_train_servicer.py
index 923e2301..b8b74300 100644
--- a/caikit/runtime/servicers/global_train_servicer.py
+++ b/caikit/runtime/servicers/global_train_servicer.py
@@ -54,9 +54,6 @@ class GlobalTrainServicer:
     def __init__(self, training_service: ServicePackage):
         self._training_service = training_service
         self._model_manager = ModelManager.get_instance()
-        caikit_config = get_config()
-        self.training_output_dir = caikit_config.runtime.training.output_dir
-        self.save_with_id = caikit_config.runtime.training.save_with_id
 
         # TODO: think about if we really want to do this here:
         self.cdm = get_data_model()
@@ -69,7 +66,7 @@ class GlobalTrainServicer:
         # Or grab the `libraries` off of the `training_service` instead of config here?
         # Duplicate code in global_train_servicer
         # pylint: disable=duplicate-code
-        self.library = clean_lib_names(caikit_config.runtime.library)[0]
+        self.library = clean_lib_names(get_config().runtime.library)[0]
         try:
             lib_version = version(self.library)
         except Exception:  # pylint: disable=broad-exception-caught
@@ -83,6 +80,14 @@ class GlobalTrainServicer:
         )
         super()
 
+    @property
+    def training_output_dir(self) -> str:
+        return get_config().runtime.training.output_dir
+
+    @property
+    def save_with_id(self) -> str:
+        return get_config().runtime.training.save_with_id
+
     def Train(
         self,
         request: ProtoMessageType,
diff --git a/caikit/runtime/servicers/model_train_servicer.py b/caikit/runtime/servicers/model_train_servicer.py
index e79b9833..697fc870 100644
--- a/caikit/runtime/servicers/model_train_servicer.py
+++ b/caikit/runtime/servicers/model_train_servicer.py
@@ -125,12 +125,12 @@ class ModelTrainServicerImpl(process_pb2_grpc.ProcessServicer):
                 #   external_training_id override. If that is ever not the case
                 #   this will be brokenly passing that kwarg through to the
                 #   module's train function.
-                external_training_id=request.trainingID,
+                external_training_id=request.customTrainingID or request.trainingID,
             )
             log.debug("<RUN00837184D>", "training_response: %s", training_response)
             # return response
             process_response = process_pb2.ProcessResponse(
-                trainingID=training_response.training_id,
+                trainingID=request.trainingID,
                 customTrainingID=request.customTrainingID,
             )
             return process_response


Test patch:
diff --git a/tests/runtime/servicers/test_model_train_servicer_impl.py b/tests/runtime/servicers/test_model_train_servicer_impl.py
index d91090b7..4077a31a 100644
--- a/tests/runtime/servicers/test_model_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_model_train_servicer_impl.py
@@ -186,10 +186,14 @@ def test_model_train_validation_error_raises(sample_model_train_servicer, output
 #####################################################################
 # Normal tests
 @pytest.mark.parametrize("output_in_config", (True, False))
+@pytest.mark.parametrize("save_with_id", (True, False))
+@pytest.mark.parametrize("use_custom_id", (True, False))
 def test_model_train_sample_widget(
     sample_model_train_servicer,
     output_dir,
     output_in_config,
+    save_with_id,
+    use_custom_id,
 ):
     """This test tests end-to-end training. It includes verifying that the model
     is saved in the right place. The place where the model is saved comes from:
@@ -200,6 +204,8 @@ def test_model_train_sample_widget(
     IN THAT ORDER!
     """
     training_id = str(uuid.uuid4())
+    custom_id = f"custom-{uuid.uuid4()}" if use_custom_id else None
+    model_name = "abc"
 
     # If getting the output path from config, use a bogus dir in the request to
     # ensure that the config is preferred
@@ -212,17 +218,24 @@ def test_model_train_sample_widget(
         req_output_dir = output_dir
 
     with temp_config(
-        {"runtime": {"training": {"output_dir": config_output_dir}}}, "merge"
+        {
+            "runtime": {
+                "training": {
+                    "output_dir": config_output_dir,
+                    "save_with_id": save_with_id,
+                },
+            },
+        },
+        "merge",
     ):
-        training_output_dir = os.path.join(output_dir, training_id)
         model_train_request = process_pb2.ProcessRequest(
             trainingID=training_id,
-            customTrainingID=str(uuid.uuid4()),
+            customTrainingID=custom_id,
             request_dict={
                 "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
                 "training_params": json.dumps(
                     {
-                        "model_name": "abc",
+                        "model_name": model_name,
                         "parameters": {
                             "training_data": {
                                 "jsondata": {
@@ -244,16 +257,23 @@ def test_model_train_sample_widget(
         training_response = sample_model_train_servicer.Run(
             model_train_request, context
         )
-        assert os.path.isdir(training_output_dir)
+        if save_with_id:
+            if use_custom_id:
+                model_save_path = os.path.join(output_dir, custom_id, model_name)
+            else:
+                model_save_path = os.path.join(output_dir, training_id, model_name)
+        else:
+            model_save_path = os.path.join(output_dir, model_name)
+        assert os.path.isdir(model_save_path)
 
     # Make sure that the request completed synchronously
-    model_future = MODEL_MANAGER.get_model_future(training_response.trainingID)
+    model_future = MODEL_MANAGER.get_model_future(custom_id or training_id)
     assert model_future.get_info().status == TrainingStatus.COMPLETED
 
     # Make sure that the return object looks right
     assert isinstance(training_response, process_pb2.ProcessResponse)
     assert training_response == process_pb2.ProcessResponse(
-        trainingID=model_train_request.trainingID,
+        trainingID=training_id,
         customTrainingID=model_train_request.customTrainingID,
     )
 
@@ -277,7 +297,6 @@ def test_files_from_training_input_dir_are_used(
 
     model_train_request = process_pb2.ProcessRequest(
         trainingID=training_id,
-        customTrainingID=str(uuid.uuid4()),
         request_dict={
             "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
             "training_params": json.dumps(
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 1c0a3de5..f6b13339 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -122,7 +122,6 @@ def test_model_train(runtime_grpc_server):
     model_name = "abc"
     model_train_request = process_pb2.ProcessRequest(
         trainingID=training_id,
-        customTrainingID=str(uuid.uuid4()),
         request_dict={
             "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
             "training_params": json.dumps(