Repository: caikit/caikit
Pull request #378 (2023-07-27 21:41:13)

Add output_path to train apis
**What this PR does / why we need it**:

For #369 

This PR adds `output_path: S3Path` to all train apis

This PR _does not_ nest the rest of the training parameters. That turned out to be a bit of a bear so we're following up in a separate PR.

**Special notes for your reviewer**:

We went with just an `S3Path` for now and not also a `str` because we already have pretty good on-disk-path support via the `runtime.training.output_dir` configuration. That allows a deployment-wide base path to be set and then the `model_name` parameter provides the directory separation within that base path on a per-training-request basis. We think some more design thought is needed if we want to start overriding disk paths directly from the training API.

This change updates the base trainer implementation to fail if an S3 path is passed to the logic that builds disk paths- we expect trainers that support S3 paths to do their own path handling.

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/.gitignore b/.gitignore
index b6a7ee55..d3ce3652 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@
 __pycache__
 .coverage
 .coverage.*
+durations/*
 coverage*.xml
 dist
 htmlcov
diff --git a/caikit/core/model_management/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index 8c0c3964..1fc2ace4 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -18,7 +18,7 @@ The LocalModelTrainer uses a local thread to launch and manage each training job
 # Standard
 from concurrent.futures.thread import _threads_queues
 from datetime import datetime, timedelta
-from typing import Optional, Type
+from typing import Optional, Type, Union
 import os
 import re
 import threading
@@ -29,6 +29,7 @@ import aconfig
 import alog
 
 # Local
+from ...interfaces.common.data_model.stream_sources import S3Path
 from ..data_model import TrainingStatus
 from ..modules import ModuleBase
 from ..toolkit.destroyable_process import DestroyableProcess
@@ -66,7 +67,7 @@ class LocalModelTrainer(ModelTrainerBase):
             trainer_name: str,
             module_class: Type[ModuleBase],
             *args,
-            save_path: Optional[str],
+            save_path: Optional[Union[str, S3Path]],
             save_with_id: bool,
             external_training_id: Optional[str],
             use_subprocess: bool,
diff --git a/caikit/core/model_management/model_trainer_base.py b/caikit/core/model_management/model_trainer_base.py
index 78a8e80c..f08238d4 100644
--- a/caikit/core/model_management/model_trainer_base.py
+++ b/caikit/core/model_management/model_trainer_base.py
@@ -26,12 +26,13 @@ model_management:
 """
 
 # Standard
-from typing import List, Optional, Type
+from typing import List, Optional, Type, Union
 import abc
 import dataclasses
 import os
 
 # Local
+from ...interfaces.common.data_model.stream_sources import S3Path
 from ..data_model import TrainingStatus
 from ..modules import ModuleBase
 from ..toolkit.factory import FactoryConstructible
@@ -59,8 +60,11 @@ class ModelTrainerBase(FactoryConstructible):
             trainer_name: str,
             training_id: str,
             save_with_id: bool,
-            save_path: Optional[str],
+            save_path: Optional[Union[str, S3Path]],
         ):
+            # Trainers should deal with an S3 ref first and not pass it along here
+            if save_path and isinstance(save_path, S3Path):
+                raise ValueError("S3 output path not supported by this runtime")
             self._id = self.__class__.ID_DELIMITER.join(
                 [ReversibleHasher.hash(trainer_name), training_id]
             )
@@ -135,7 +139,7 @@ class ModelTrainerBase(FactoryConstructible):
         self,
         module_class: Type[ModuleBase],
         *args,
-        save_path: Optional[str] = None,
+        save_path: Optional[Union[str, S3Path]] = None,
         save_with_id: bool = False,
         **kwargs,
     ) -> "ModelFutureBase":
diff --git a/caikit/core/model_manager.py b/caikit/core/model_manager.py
index 11abe25f..12992cb9 100644
--- a/caikit/core/model_manager.py
+++ b/caikit/core/model_manager.py
@@ -29,6 +29,7 @@ import zipfile
 import alog
 
 # Local
+from ..interfaces.common.data_model.stream_sources import S3Path
 from .model_management import (
     ModelFinderBase,
     ModelInitializerBase,
@@ -95,7 +96,7 @@ class ModelManager:
         module: Union[Type[ModuleBase], str],
         *args,
         trainer: Union[str, ModelTrainerBase] = "default",
-        save_path: Optional[str] = None,
+        save_path: Optional[Union[str, S3Path]] = None,
         save_with_id: bool = False,
         wait: bool = False,
         **kwargs,
@@ -119,8 +120,9 @@ 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[str]): Path on disk where the model should be
-                saved (may be relative to a remote trainer's filesystem)
+            save_path (Optional[Union[str, S3Path]]): 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
             wait (bool): Wait for training to complete before returning
@@ -144,7 +146,7 @@ class ModelManager:
         error.subclass_check("<COR05418775E>", module, ModuleBase)
 
         # Get the trainer to use
-        trainer = self._get_trainer(trainer)
+        trainer: ModelTrainerBase = self._get_trainer(trainer)
 
         # Start the training
         with alog.ContextTimer(log.debug, "Started training in: "):
diff --git a/caikit/runtime/service_generation/rpcs.py b/caikit/runtime/service_generation/rpcs.py
index f55f8dbb..f600edf3 100644
--- a/caikit/runtime/service_generation/rpcs.py
+++ b/caikit/runtime/service_generation/rpcs.py
@@ -30,6 +30,7 @@ from py_to_proto.dataclass_to_proto import (  # NOTE: Imported from here for com
 import alog
 
 # Local
+from ...interfaces.common.data_model.stream_sources import S3Path
 from . import protoable, type_helpers
 from .compatibility_checker import ApiFieldNames
 from .data_stream_source import make_data_stream_source
@@ -169,7 +170,10 @@ class ModuleClassTrainRPC(CaikitRPCBase):
         # Change return type for async training interface
         return_type = TrainingJob
 
-        new_params = {"model_name": str}
+        # Start with extra metaparameters
+        # - model_name: user-provided custom ID for the model to train
+        # - output_path: pointer to some storage where the model will be saved
+        new_params = {"model_name": str, "output_path": S3Path}
         for name, typ in signature.parameters.items():
             if type_helpers.has_data_stream(typ):
                 # Assume this is training data
diff --git a/caikit/runtime/servicers/global_train_servicer.py b/caikit/runtime/servicers/global_train_servicer.py
index b561f29f..893b5afa 100644
--- a/caikit/runtime/servicers/global_train_servicer.py
+++ b/caikit/runtime/servicers/global_train_servicer.py
@@ -28,7 +28,7 @@ import alog
 # Local
 from caikit import get_config
 from caikit.core import MODEL_MANAGER, ModuleBase
-from caikit.core.data_model import DataBase
+from caikit.interfaces.common.data_model.stream_sources import S3Path
 from caikit.interfaces.runtime.data_model import TrainingJob
 from caikit.runtime.model_management.model_manager import ModelManager
 from caikit.runtime.service_factory import ServicePackage
@@ -86,7 +86,7 @@ class GlobalTrainServicer:
 
     def Train(
         self,
-        request,
+        request: ProtoMessageType,
         context: ServicerContext,
         *_,
         **__,
@@ -150,7 +150,7 @@ class GlobalTrainServicer:
             log.warning(log_dict)
             raise CaikitRuntimeException(
                 StatusCode.INVALID_ARGUMENT,
-                f"Exception raised during inference. This may be a problem with your input: {e}",
+                f"Exception raised during training. This may be a problem with your input: {e}",
             ) from e
 
         except Exception as e:
@@ -166,7 +166,7 @@ class GlobalTrainServicer:
 
     def run_training_job(
         self,
-        request: Union[ProtoMessageType, DataBase],
+        request: ProtoMessageType,
         module: Type[ModuleBase],
         training_output_dir: str,
         *,
@@ -178,8 +178,7 @@ class GlobalTrainServicer:
         then returns the thread id
 
         Args:
-            request (Union[ProtoMessageType, DataBase]): The message that
-                stimulated this request
+            request (ProtoMessageType): The message that stimulated this request
             module (Type[ModuleBase]): The module class to train
             training_output_dir (str): The base directory where trained models
                 should be saved
@@ -193,9 +192,20 @@ class GlobalTrainServicer:
             training_job (TrainingJob): The job handle for the training with the
                 job's ID and the model's name
         """
+        request_data_model = caikit.core.data_model.DataBase.get_class_for_proto(
+            request
+        ).from_proto(request)
 
         # Figure out where this model will be saved
-        model_path = self._get_model_path(training_output_dir, request.model_name)
+        model_path: Union[str, S3Path]
+        if request_data_model.output_path:
+            # 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
+            )
 
         # Build the full set of kwargs for the train call
         kwargs.update(
diff --git a/caikit/runtime/servicers/model_train_servicer.py b/caikit/runtime/servicers/model_train_servicer.py
index 52f9c731..79e6ec85 100644
--- a/caikit/runtime/servicers/model_train_servicer.py
+++ b/caikit/runtime/servicers/model_train_servicer.py
@@ -144,7 +144,7 @@ class ModelTrainServicerImpl(process_pb2_grpc.ProcessServicer):
             )
             raise CaikitRuntimeException(
                 grpc.StatusCode.INVALID_ARGUMENT,
-                f"Exception raised during inference. This may be a problem with your input: {e}",
+                f"Exception raised during training. This may be a problem with your input: {e}",
             ) from e
 
         except Exception as e:
diff --git a/scripts/check_deps.sh b/scripts/check_deps.sh
index ccbda887..2033700e 100755
--- a/scripts/check_deps.sh
+++ b/scripts/check_deps.sh
@@ -18,9 +18,15 @@ then
     exit 1
 fi
 
-if < deps.txt grep -q ".*caikit_interfaces.*\->.*caikit_core.*"
+if < deps.txt grep -q ".*caikit_interfaces.*\->.*caikit_core.module*"
 then
-    echo "Fail: The core is importing the interfaces!"
+    echo "Fail: The core module definitions are importing the interfaces!"
+    exit 1
+fi
+
+if < deps.txt grep -q ".*caikit_interfaces.*\->.*caikit_core.data_model.*"
+then
+    echo "Fail: The core data model is importing the interfaces!"
     exit 1
 fi
 
diff --git a/tox.ini b/tox.ini
index 5818d748..99abccba 100644
--- a/tox.ini
+++ b/tox.ini
@@ -12,7 +12,7 @@ passenv =
     LOG_FORMATTER
     LOG_THREAD_ID
     LOG_CHANNEL_WIDTH
-commands = pytest --cov=caikit --cov-report=html:coverage-{env_name} --cov-report=xml:coverage-{env_name}.xml {posargs:tests -m "not examples"}
+commands = pytest --cov=caikit --cov-report=html:coverage-{env_name} --cov-report=xml:coverage-{env_name}.xml --html=durations/{env_name}.html {posargs:tests -m "not examples"}
 
 ; Unclear: We probably want to test wheel packaging
 ; But! tox will fail when this is set and _any_ interpreter is missing


Test patch:
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index cf9ee9cd..0fa6fb5e 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -20,10 +20,12 @@ import threading
 import time
 
 # Third Party
+import grpc
 import pytest
 
 # Local
 from caikit.core import MODEL_MANAGER
+from caikit.interfaces.common.data_model.stream_sources import S3Path
 from caikit.runtime.servicers.global_train_servicer import GlobalTrainServicer
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from sample_lib.data_model.sample import (
@@ -369,6 +371,29 @@ def test_global_train_returns_exit_code_with_oom(
         assert f"Training process died with OOM error!" in str(future_info.errors[0])
 
 
+def test_local_trainer_rejects_s3_output_paths(
+    sample_train_service, sample_train_servicer
+):
+    """Test that if an output_path references S3 but the trainer does not know how to handle it, an error is raised"""
+    stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
+    training_data = stream_type(
+        jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
+    ).to_proto()
+    train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
+        model_name=random_test_id(),
+        output_path=S3Path(path="foo").to_proto(),
+        batch_size=42,
+        training_data=training_data,
+        oom_exit=True,
+    )
+
+    with pytest.raises(
+        CaikitRuntimeException, match=".*S3 output path not supported by this runtime"
+    ) as ctx:
+        sample_train_servicer.Train(train_request, Fixtures.build_context("foo"))
+    assert ctx.value.status_code == grpc.StatusCode.INVALID_ARGUMENT
+
+
 #####################################################################