Repository: caikit/caikit
Pull request #520 (2023-10-11 20:39:37)

Pluggable data input
**What this PR does / why we need it**:

Addresses #517 

This PR refactors `DataStreamSource` to us a pluggable/extensible factory pattern to allow for external library extensions. Concretely, this adds the following:

* New config section `data_streams.source_plugins` to define the set of source plugins
    * The keys in this map are arbitrary labels that are used for logging, but nothing else. Do we like this?
* New base class (`DataStreamSourcePlugin`) that can be extended for new plugins
* New factory for stream source plugins

The ADR also contains details about plans to handle pluggable output with `OutputTarget`s, though those are not implemented in this PR

**If applicable**:
- [x] this PR contains documentation
- [x] this PR contains unit tests
- [x] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 28862e77..fdf72014 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -69,6 +69,22 @@ data_streams:
     # Base directory where stream source relative paths should be found
     file_source_base: null
 
+    # Config for data stream source plugins. The keys in the map are names for
+    # environment scoping and the values are factory blobs for an importable
+    # factory. Each blob requires a `type` and can optionally have `config` and
+    # `import_class`.
+    source_plugins:
+        inline:
+            type: JsonData
+        file:
+            type: FileData
+        file_list:
+            type: ListOfFiles
+        directory:
+            type: Directory
+        s3:
+            type: S3Files
+
 ### Runtime configurations
 runtime:
     # The runtime library (or libraries) whose models we want to serve using Caikit Runtime. This should
diff --git a/caikit/core/model_management/factories.py b/caikit/core/model_management/factories.py
index c0aaccdb..c39cc5db 100644
--- a/caikit/core/model_management/factories.py
+++ b/caikit/core/model_management/factories.py
@@ -14,78 +14,14 @@
 """
 Global factories for model management
 """
-# Standard
-from typing import Optional
-import importlib
-
-# First Party
-import alog
 
 # Local
-from ..exceptions import error_handler
-from ..toolkit.factory import Factory, FactoryConstructible
+from ..toolkit.factory import ImportableFactory
 from .local_model_finder import LocalModelFinder
 from .local_model_initializer import LocalModelInitializer
 from .local_model_trainer import LocalModelTrainer
 from .multi_model_finder import MultiModelFinder
 
-log = alog.use_channel("MMFCTRY")
-error = error_handler.get(log)
-
-
-class ImportableFactory(Factory):
-    """An ImportableFactory extends the base Factory to allow the construction
-    to specify an "import_class" field that will be used to import and register
-    the implementation class before attempting to initialize it.
-    """
-
-    IMPORT_CLASS_KEY = "import_class"
-
-    def construct(
-        self,
-        instance_config: dict,
-        instance_name: Optional[str] = None,
-    ):
-        # Look for an import_class and import and register it if found
-        import_class_val = instance_config.get(self.__class__.IMPORT_CLASS_KEY)
-        if import_class_val:
-            error.type_check(
-                "<COR85108801E>",
-                str,
-                **{self.__class__.IMPORT_CLASS_KEY: import_class_val},
-            )
-            module_name, class_name = import_class_val.rsplit(".", 1)
-            try:
-                imported_module = importlib.import_module(module_name)
-            except ImportError:
-                error(
-                    "<COR46837141E>",
-                    ValueError(
-                        "Invalid {}: Module cannot be imported [{}]".format(
-                            self.__class__.IMPORT_CLASS_KEY,
-                            module_name,
-                        )
-                    ),
-                )
-            try:
-                imported_class = getattr(imported_module, class_name)
-            except AttributeError:
-                error(
-                    "<COR46837142E>",
-                    ValueError(
-                        "Invalid {}: No such class [{}] on module [{}]".format(
-                            self.__class__.IMPORT_CLASS_KEY,
-                            class_name,
-                            module_name,
-                        )
-                    ),
-                )
-            error.subclass_check("<COR52306423E>", imported_class, FactoryConstructible)
-
-            self.register(imported_class)
-        return super().construct(instance_config, instance_name)
-
-
 # Model trainer factory. A trainer is responsible for performing the train
 # operation against a configured framework connection.
 model_trainer_factory = ImportableFactory("ModelTrainer")
diff --git a/caikit/core/toolkit/factory.py b/caikit/core/toolkit/factory.py
index e5b88cb7..c2c2dd59 100644
--- a/caikit/core/toolkit/factory.py
+++ b/caikit/core/toolkit/factory.py
@@ -19,6 +19,7 @@ base class for classes that can be constructed via caikit config
 # Standard
 from typing import Optional, Type
 import abc
+import importlib
 
 # First Party
 import aconfig
@@ -105,3 +106,56 @@ class Factory:
         )
         instance_name = instance_name or inst_cls.name
         return inst_cls(inst_cfg, instance_name)
+
+
+class ImportableFactory(Factory):
+    """An ImportableFactory extends the base Factory to allow the construction
+    to specify an "import_class" field that will be used to import and register
+    the implementation class before attempting to initialize it.
+    """
+
+    IMPORT_CLASS_KEY = "import_class"
+
+    def construct(
+        self,
+        instance_config: dict,
+        instance_name: Optional[str] = None,
+    ):
+        # Look for an import_class and import and register it if found
+        import_class_val = instance_config.get(self.__class__.IMPORT_CLASS_KEY)
+        if import_class_val:
+            error.type_check(
+                "<COR85108801E>",
+                str,
+                **{self.__class__.IMPORT_CLASS_KEY: import_class_val},
+            )
+            module_name, class_name = import_class_val.rsplit(".", 1)
+            try:
+                imported_module = importlib.import_module(module_name)
+            except ImportError:
+                error(
+                    "<COR46837141E>",
+                    ValueError(
+                        "Invalid {}: Module cannot be imported [{}]".format(
+                            self.__class__.IMPORT_CLASS_KEY,
+                            module_name,
+                        )
+                    ),
+                )
+            try:
+                imported_class = getattr(imported_module, class_name)
+            except AttributeError:
+                error(
+                    "<COR46837142E>",
+                    ValueError(
+                        "Invalid {}: No such class [{}] on module [{}]".format(
+                            self.__class__.IMPORT_CLASS_KEY,
+                            class_name,
+                            module_name,
+                        )
+                    ),
+                )
+            error.subclass_check("<COR52306423E>", imported_class, FactoryConstructible)
+
+            self.register(imported_class)
+        return super().construct(instance_config, instance_name)
diff --git a/caikit/runtime/service_generation/data_stream_source.py b/caikit/runtime/service_generation/data_stream_source.py
index d203f744..21a8aea5 100644
--- a/caikit/runtime/service_generation/data_stream_source.py
+++ b/caikit/runtime/service_generation/data_stream_source.py
@@ -13,8 +13,10 @@
 # limitations under the License.
 
 # Standard
+from functools import cached_property, partial
 from glob import glob
-from typing import Any, List, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Type, Union
+import abc
 import os
 import sys
 
@@ -22,14 +24,17 @@ import sys
 import grpc
 
 # First Party
-from py_to_proto.dataclass_to_proto import Annotated, OneofField
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber, OneofField
+import aconfig
 import alog
 
 # Local
+from caikit.config import get_config
 from caikit.core.data_model.base import DataBase
 from caikit.core.data_model.dataobject import _make_oneof_init, make_dataobject
 from caikit.core.data_model.streams.data_stream import DataStream
 from caikit.core.exceptions import error_handler
+from caikit.core.toolkit.factory import FactoryConstructible, ImportableFactory
 from caikit.interfaces.common.data_model.stream_sources import (
     Directory,
     File,
@@ -52,124 +57,76 @@ log = alog.use_channel("DSTRM-SRC")
 error = error_handler.get(log)
 
 
-class DataStreamSourceBase(DataStream):
-    """This base class acts as a sentinel so that dynamically generated data
-    stream source classes can be identified programmatically.
+## Plugin Bases ################################################################
+
+
+class DataStreamSourcePlugin(FactoryConstructible):
+    """A DataStreamSourcePlugin is a pluggable source that defines the shape of
+    the data object needed as well as the code for accessing the data from some
+    source type.
     """
 
-    def __init__(self):
-        super().__init__(self._generator)
+    def __init__(self, config: aconfig.Config, instance_name: str):
+        """Construct with the basic factory constructible interface and store the
+        args for use by the child
+        """
+        self._config = config
+        self._instance_name = instance_name
 
-    def _generator(self):
-        stream = self.to_data_stream()
-        return stream.generator_func(*stream.generator_args, **stream.generator_kwargs)
+    ## Abstract Interface ##
 
-    def __getstate__(self) -> bytes:
-        """A DataStreamSource is pickled by serializing its source
-        representation. This is particularly useful when sharing data streams
-        across subprocesses to run training in an isolated process.
+    @abc.abstractmethod
+    def get_stream_message_type(self, element_type: type) -> Type[DataBase]:
+        """Get the type of the dataobject class that will be used as the source
+        information
         """
-        return self.to_binary_buffer()
 
-    def __setstate__(self, pickle_bytes: bytes):
-        """Unpickling a DataStreamSource basically involves unpacking the
-        serialized source representation. The catch is that the oneof is
-        represented strangely in __dict__, so we need to explicitly set all
-        oneof members.
+    @abc.abstractmethod
+    def to_data_stream(
+        self, source_message: Type[DataBase], element_type: type
+    ) -> DataStream:
+        """Convert an instance of the source message type into a DataStream"""
+
+    @abc.abstractmethod
+    def get_field_number(self) -> int:
+        """Each plugin must define its field number which may be informed by
+        self._config
         """
-        new_inst = self.__class__.from_binary_buffer(pickle_bytes)
-        setattr(self, new_inst.which_oneof("data_stream"), new_inst.data_stream)
-        self.generator_func = self._generator
-        self.generator_args = tuple()
-        self.generator_kwargs = {}
 
-    # pylint: disable=too-many-return-statements
-    def to_data_stream(self) -> DataStream:
-        """Convert to the target data stream type based on the source type"""
+    ## Public Methods ##
 
-        # Determine which of the value types is set
-        set_field = None
-        for field_name in self.get_proto_class().DESCRIPTOR.fields_by_name:
-            if getattr(self, field_name) is not None:
-                assert (
-                    set_field is None
-                ), "Found DataStreamSource with multiple sources set"
-                set_field = field_name
+    def get_field_name(self, element_type: type) -> str:
+        """The name of the field that this plugin will use in the source oneof"""
+        return self.get_stream_message_type(element_type).__name__.lower()
 
-        # If no field is set, return an empty DataStream
-        if set_field is None:
-            log.debug3("Returning empty data stream")
-            return DataStream.from_iterable([])
+    ## Shared Impl ##
 
-        # If a S3 pointer is given, raise not implemented
-        if set_field == "s3files":
-            error(
-                "<COR80419785E>",
-                NotImplementedError(
-                    "S3Files are not implemented as stream sources in this runtime."
-                ),
-            )
+    @staticmethod
+    def _to_element_type(element_type: type, raw_element: Any) -> Any:
+        if issubclass(element_type, DataBase):
+            # To allow for extra fields (e.g. in training data) that may not
+            # be needed by the data objects, we ignore unknown fields
+            return element_type.from_json(raw_element, ignore_unknown_fields=True)
+        return raw_element
+
+    @staticmethod
+    def _to_element_partial(element_type: type) -> Callable:
+        return partial(DataStreamSourcePlugin._to_element_type, element_type)
 
-        # If jsondata, pull from the data elements directly
-        if set_field == "jsondata":
-            log.debug3("Pulling data stream from inline json")
-            return DataStream.from_iterable(self.jsondata.data)
-
-        # If jsonfile, attempt to read the file and pull in the data from there
-        if set_field == "file":
-            return self._create_data_stream_from_file(fname=self.file.filename)
-
-        # If list of files, attempt to read all files and combine datastreams from all
-        if set_field == "listoffiles":
-            # combined list of data streams that we will chain and send back
-            data_stream_list = []
-            for fname in self.listoffiles.files:
-                data_stream_list.append(self._create_data_stream_from_file(fname=fname))
-
-            return DataStream.chain(data_stream_list).flatten()
-
-        # If directory, attempt to read an element from each file in the dir
-        if set_field == "directory":
-            dirname = self.directory.dirname
-            full_dirname = self._get_resolved_source_path(dirname)
-            extension = self.directory.extension or "json"
-            if not dirname or not os.path.isdir(full_dirname):
-                raise CaikitRuntimeException(
-                    grpc.StatusCode.INVALID_ARGUMENT,
-                    f"Invalid {extension} directory source file: {full_dirname}",
-                )
-            files_with_ext = list(glob(os.path.join(full_dirname, "*." + extension)))
-            # make sure at least 1 file with the given extension exists
-            if len(files_with_ext) == 0:
-                raise CaikitRuntimeException(
-                    grpc.StatusCode.INVALID_ARGUMENT,
-                    f"directory {dirname} contains no source files with extension {extension}",
-                )
-            if extension == "json":
-                return DataStream.from_json_collection(full_dirname, extension).map(
-                    self._to_element_type
-                )
-            if extension == "csv":
-                return DataStream.from_csv_collection(full_dirname).map(
-                    self._to_element_type
-                )
-            if extension == "jsonl":
-                return DataStream.from_jsonl_collection(full_dirname).map(
-                    self._to_element_type
-                )
-            raise CaikitRuntimeException(
-                grpc.StatusCode.INVALID_ARGUMENT,
-                f"Extension not supported! {extension}",
-            )
+
+class FilePluginBase(DataStreamSourcePlugin):
+    """Intermediate base class for file-based plugins with helper utilities"""
 
     @classmethod
-    def _create_data_stream_from_file(cls, fname: str) -> DataStream:
+    def _create_data_stream_from_file(
+        cls, fname: str, element_type: type
+    ) -> DataStream:
         """Create a data stream object by deducing file extension
         and reading the file accordingly"""
 
         _, extension = os.path.splitext(fname)
         if not extension:
-            return cls._load_from_file_without_extension(fname)
+            return cls._load_from_file_without_extension(fname, element_type)
 
         full_fname = cls._get_resolved_source_path(fname)
         log.debug3("Pulling data stream from %s file [%s]", extension, full_fname)
@@ -179,22 +136,23 @@ class DataStreamSourceBase(DataStream):
                 grpc.StatusCode.INVALID_ARGUMENT,
                 f"Invalid {extension} data source file: {fname}",
             )
+        to_element_type = cls._to_element_partial(element_type)
         if extension == ".json":
-            stream = DataStream.from_json_array(full_fname).map(cls._to_element_type)
+            stream = DataStream.from_json_array(full_fname).map(to_element_type)
             # Iterate once to make sure this is a json array
             stream.peek()
             return stream
         if extension == ".csv":
-            return DataStream.from_header_csv(full_fname).map(cls._to_element_type)
+            return DataStream.from_header_csv(full_fname).map(to_element_type)
         if extension == ".jsonl":
-            return DataStream.from_jsonl(full_fname).map(cls._to_element_type)
+            return DataStream.from_jsonl(full_fname).map(to_element_type)
         raise CaikitRuntimeException(
             grpc.StatusCode.INVALID_ARGUMENT,
             f"Extension not supported! {extension}",
         )
 
     @classmethod
-    def _load_from_file_without_extension(cls, fname) -> DataStream:
+    def _load_from_file_without_extension(cls, fname, element_type: type) -> DataStream:
         """Similar to _create_data_stream_from_file, but we don't have a file extension to work
         with. Attempt to create a data stream using one of a few well-known formats.
         🌶🌶🌶️ on ordering here:
@@ -207,6 +165,7 @@ class DataStreamSourceBase(DataStream):
             confidently return a stream even if that's not the case.
         """
         full_fname = cls._get_resolved_source_path(fname)
+        to_element_type = cls._to_element_partial(element_type)
         log.debug3("Attempting to guess file type for file: %s", full_fname)
         for factory_method in (
             DataStream.from_json_array,
@@ -215,13 +174,14 @@ class DataStreamSourceBase(DataStream):
             DataStream.from_header_csv,
         ):
             try:
-                stream = factory_method(full_fname).map(cls._to_element_type)
-                # Iterate once and assume we have the correct file type if this works
+                stream = factory_method(full_fname).map(to_element_type)
+                # Iterate once and assume we have the correct file type if this
+                # works
                 stream.peek()
                 return stream
             except Exception as e:  # pylint: disable=broad-exception-caught
-                # Catch any exception: it's hard to know which all could be thrown by any of the
-                # formatters
+                # Catch any exception: it's hard to know which all could be
+                # thrown by any of the formatters
                 log.debug3(
                     "Failed to load file %s using data stream factory method %s: %s",
                     full_fname,
@@ -234,17 +194,6 @@ class DataStreamSourceBase(DataStream):
             f"Could not load input file with no extension: {full_fname}",
         )
 
-    @classmethod
-    def _to_element_type(cls, raw_element: Any) -> "ElementType":
-        """Stream adapter to adapt from the raw json object data representations
-        to the underlying data objects
-        """
-        if issubclass(cls.ELEMENT_TYPE, DataBase):
-            # To allow for extra fields (e.g. in training data) that may not
-            # be needed by the data objects, we ignore unknown fields
-            return cls.ELEMENT_TYPE.from_json(raw_element, ignore_unknown_fields=True)
-        return raw_element
-
     @staticmethod
     def _get_resolved_source_path(input_path: str) -> str:
         """Get a fully resolved path, including any shared prefix"""
@@ -255,53 +204,351 @@ class DataStreamSourceBase(DataStream):
         return os.path.join(source_pfx, input_path) if source_pfx else input_path
 
 
-def make_data_stream_source(data_element_type: Type) -> Type[DataBase]:
+## Source Plugins ##############################################################
+
+
+class FileDataStreamSourcePlugin(FilePluginBase):
+    """Plugin for a single file"""
+
+    name = "FileData"
+
+    def get_stream_message_type(self, *_, **__) -> Type[DataBase]:
+        return File
+
+    def to_data_stream(self, source_message: File, element_type: type) -> DataStream:
+        return self._create_data_stream_from_file(
+            fname=source_message.filename, element_type=element_type
+        )
+
+    def get_field_number(self) -> int:
+        return 2
+
+
+class ListOfFilesDataStreamSourcePlugin(FilePluginBase):
+    """Plugin for a list of files"""
+
+    name = "ListOfFiles"
+
+    def get_stream_message_type(self, *_, **__) -> Type[DataBase]:
+        return ListOfFiles
+
+    def to_data_stream(
+        self, source_message: ListOfFiles, element_type: type
+    ) -> DataStream:
+        data_stream_list = []
+        for fname in source_message.files:
+            data_stream_list.append(
+                self._create_data_stream_from_file(
+                    fname=fname, element_type=element_type
+                )
+            )
+
+        return DataStream.chain(data_stream_list).flatten()
+
+    def get_field_number(self) -> int:
+        return 3
+
+
+class DirectoryDataStreamSourcePlugin(FilePluginBase):
+    """Plugin for a directory holding files"""
+
+    name = "Directory"
+
+    def get_stream_message_type(self, *_, **__) -> Type[DataBase]:
+        return Directory
+
+    def to_data_stream(
+        self, source_message: Directory, element_type: type
+    ) -> DataStream:
+        dirname = source_message.dirname
+        full_dirname = self._get_resolved_source_path(dirname)
+        extension = source_message.extension or "json"
+        if not dirname or not os.path.isdir(full_dirname):
+            raise CaikitRuntimeException(
+                grpc.StatusCode.INVALID_ARGUMENT,
+                f"Invalid {extension} directory source file: {full_dirname}",
+            )
+        files_with_ext = list(glob(os.path.join(full_dirname, "*." + extension)))
+        to_element_type = self._to_element_partial(element_type)
+        # make sure at least 1 file with the given extension exists
+        if len(files_with_ext) == 0:
+            raise CaikitRuntimeException(
+                grpc.StatusCode.INVALID_ARGUMENT,
+                f"directory {dirname} contains no source files with extension {extension}",
+            )
+        if extension == "json":
+            return DataStream.from_json_collection(full_dirname, extension).map(
+                to_element_type
+            )
+        if extension == "csv":
+            return DataStream.from_csv_collection(full_dirname).map(to_element_type)
+        if extension == "jsonl":
+            return DataStream.from_jsonl_collection(full_dirname).map(to_element_type)
+        raise CaikitRuntimeException(
+            grpc.StatusCode.INVALID_ARGUMENT,
+            f"Extension not supported! {extension}",
+        )
+
+    def get_field_number(self) -> int:
+        return 4
+
+
+class JsonDataStreamSourcePlugin(DataStreamSourcePlugin):
+    """This plugin is for inline data, elements are provided in a list.
+
+    This plugin has instantiation logic: it needs the stream's element type so that it can
+    generate a data model for List[element_type]"""
+
+    name = "JsonData"
+
+    # class-level cache required to avoid creating duplicate data model classes
+    stream_source_type_cache: Dict[Type[DataBase], Type[DataBase]] = {}
+
+    def get_stream_message_type(self, element_type: type) -> Type[DataBase]:
+        stream_message_type = self.__class__.stream_source_type_cache.get(element_type)
+        if stream_message_type:
+            return stream_message_type
+
+        package = get_runtime_service_package()
+        cls_name = _make_data_stream_source_type_name(element_type)
+        JsonData = make_dataobject(
+            package=package,
+            proto_name=f"{cls_name}JsonData",
+            name="JsonData",
+            attrs={"__qualname__": f"{cls_name}.JsonData"},
+            annotations={"data": List[element_type]},
+        )
+        self.__class__.stream_source_type_cache[element_type] = JsonData
+        return JsonData
+
+    def to_data_stream(self, source_message: Type[DataBase], *_, **__) -> DataStream:
+        """source_message should be of type self.get_stream_message_type
+        So it _should_ contain an attribute named `data`, which is a list"""
+        return DataStream.from_iterable(source_message.data)
+
+    def get_field_number(self) -> int:
+        return 1
+
+
+class S3FilesDataStreamSourcePlugin(DataStreamSourcePlugin):
+    """Unimplemented!"""
+
+    name = "S3Files"
+
+    def get_stream_message_type(self, *_, **__) -> Type[DataBase]:
+        return S3Files
+
+    def to_data_stream(self, *_, **__) -> DataStream:
+        error(
+            "<RUN80419785E>",
+            NotImplementedError(
+                "S3Files are not implemented as stream sources in this runtime."
+            ),
+        )
+
+    def get_field_number(self) -> int:
+        return 5
+
+
+## DataStreamPluginFactory #####################################################
+
+
+class DataStreamPluginFactory(ImportableFactory):
+    """The DataStreamPluginFactory is responsible for holding a registry of
+    plugin instances that will be used to create and manage data stream sources
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._plugins = None
+
+    def get_plugins(
+        self, plugins_config: Optional[aconfig.Config] = None
+    ) -> List[DataStreamSourcePlugin]:
+        """Builds the set of plugins to use for a data stream source of type element_type"""
+        if self._plugins is None:
+            self._plugins = []
+            if plugins_config is None:
+                plugins_config = get_config().data_streams.source_plugins
+            for name, cfg in plugins_config.items():
+                self._plugins.append(self.construct(cfg, name))
+
+            # Make sure field numbers are unique
+            field_numbers = [plugin.get_field_number() for plugin in self._plugins]
+            duplicate_field_number_names = [
+                plugin.name
+                for plugin in self._plugins
+                if field_numbers.count(plugin.get_field_number()) > 1
+            ]
+            error.value_check(
+                "<RUN69189361E>",
+                not duplicate_field_number_names,
+                "Duplicate plugin field numbers found for plugins: {}",
+                duplicate_field_number_names,
+            )
+        return self._plugins
+
+
+# Single default instance
+PluginFactory = DataStreamPluginFactory("DataStreamSource")
+PluginFactory.register(JsonDataStreamSourcePlugin)
+PluginFactory.register(FileDataStreamSourcePlugin)
+PluginFactory.register(ListOfFilesDataStreamSourcePlugin)
+PluginFactory.register(DirectoryDataStreamSourcePlugin)
+PluginFactory.register(S3FilesDataStreamSourcePlugin)
+
+
+## DataStreamSourceBase ########################################################
+
+
+class DataStreamSourceBase(DataStream):
+    """This base class acts as a sentinel so that dynamically generated data
+    stream source classes can be identified programmatically.
+    """
+
+    def __init__(self):
+        super().__init__(self._generator)
+
+    def _generator(self):
+        stream = self.to_data_stream()
+        return stream.generator_func(*stream.generator_args, **stream.generator_kwargs)
+
+    def __getstate__(self) -> bytes:
+        """A DataStreamSource is pickled by serializing its source
+        representation. This is particularly useful when sharing data streams
+        across subprocesses to run training in an isolated process.
+        """
+        return self.to_binary_buffer()
+
+    def __setstate__(self, pickle_bytes: bytes):
+        """Unpickling a DataStreamSource basically involves unpacking the
+        serialized source representation. The catch is that the oneof is
+        represented strangely in __dict__, so we need to explicitly set all
+        oneof members.
+        """
+        new_inst = self.__class__.from_binary_buffer(pickle_bytes)
+        setattr(self, new_inst.which_oneof("data_stream"), new_inst.data_stream)
+        self.generator_func = self._generator
+        self.generator_args = tuple()
+        self.generator_kwargs = {}
+
+    @cached_property
+    def name_to_plugin_map(self):
+        return {
+            plugin.get_field_name(self.ELEMENT_TYPE): plugin for plugin in self.PLUGINS
+        }
+
+    # pylint: disable=too-many-return-statements
+    def to_data_stream(self) -> DataStream:
+        """Convert to the target data stream type based on the source type"""
+
+        # Determine which of the value types is set
+        set_field = None
+        for field_name in self.get_proto_class().DESCRIPTOR.fields_by_name:
+            if getattr(self, field_name) is not None:
+                error.value_check(
+                    "<RUN80421785E>",
+                    set_field is None,
+                    "Found DataStreamSource with multiple sources set: {} and {}",
+                    set_field,
+                    field_name,
+                )
+                error.value_check(
+                    "<RUN80420785E>",
+                    field_name in self.name_to_plugin_map,
+                    "no data stream plugin found for field: {}",
+                    field_name,
+                )
+                set_field = field_name
+
+        # If no field is set, return an empty DataStream
+        if set_field is None:
+            log.debug3("Returning empty data stream")
+            return DataStream.from_iterable([])
+
+        # Get the correct plugin, and pass it the source field + the element
+        # type to serialize to
+        plugin = self.name_to_plugin_map[set_field]
+        return plugin.to_data_stream(getattr(self, set_field), self.ELEMENT_TYPE)
+
+
+## make_data_stream_source #####################################################
+
+
+def make_data_stream_source(
+    data_element_type: type,
+    plugin_factory: DataStreamPluginFactory = PluginFactory,
+    plugins_config: Optional[aconfig.Config] = None,
+) -> Type[DataBase]:
     """Dynamically create a data stream source message type that supports
     pulling an iterable of the given type from all valid data stream sources
     """
     log.debug2("Looking for DataStreamSource[%s]", data_element_type)
     if data_element_type not in _DATA_STREAM_SOURCE_TYPES:
         cls_name = _make_data_stream_source_type_name(data_element_type)
+        package = get_runtime_service_package()
+
         log.debug("Creating DataStreamSource[%s] -> %s", data_element_type, cls_name)
 
-        # Set up the "sub class." In python, this is the same as creating a
-        # standalone class with a __qualname__ that nests it under a parent
-        # class. We do this outside the declaration of the parent class so that
-        # this class can be referenced within the Union annotation for the
-        # outer class itself. This class needs to be created dynamically because
-        # it encapsulates type information about the elements of the data stream.
-        package = get_runtime_service_package()
-        JsonData = make_dataobject(
-            package=package,
-            proto_name=f"{cls_name}JsonData",
-            name="JsonData",
-            attrs={"__qualname__": f"{cls_name}.JsonData"},
-            annotations={"data": List[data_element_type]},
+        # Get the required plugins
+        plugins = plugin_factory.get_plugins(plugins_config)
+
+        # Make sure there are no field name duplicates
+        plug_to_name = {
+            plugin: plugin.get_field_name(data_element_type) for plugin in plugins
+        }
+        all_field_names = list(plug_to_name.values())
+        duplicates = {
+            plugin.name: field_name
+            for plugin, field_name in plug_to_name.items()
+            if all_field_names.count(field_name) > 1
+        }
+        error.value_check(
+            "<RUN66854455E>",
+            not duplicates,
+            "Duplicate plugin field names found for type {}: {}",
+            data_element_type,
+            duplicates,
         )
 
-        # Create the outer class that encapsulates the Union (oneof) or the
+        # Create the outer class that encapsulates the Union (oneof) of the
         # various types of input sources
+
+        # Determine the type stream message type for each source. This can
+        # potentially be expensive, so we do it once
+        stream_message_types = {
+            plugin.name: plugin.get_stream_message_type(data_element_type)
+            for plugin in plugins
+        }
+
+        # Build the type annotation for the data model
+        # This describes a large oneof containing all the info from each data
+        # stream source plugin
+        annotation_list = [
+            Annotated[
+                stream_message_types[plugin.name],
+                OneofField(plugin.get_field_name(data_element_type)),
+                FieldNumber(plugin.get_field_number()),
+            ]
+            for plugin in plugins
+        ]
+        data_stream_type_union = Union[tuple(annotation_list)]
+
+        # Create an attribute dictionary that will expose each of the source
+        # types on this datastream class itself. E.g. if I have the `JsonData`
+        # plugin enabled, this enables:
+        # >>> make_data_stream_source(some_type).JsonData
+        # to access the `JsonData` source message directly.
+        type_attrs = {
+            msg_type.__name__: msg_type for msg_type in stream_message_types.values()
+        }
+
         data_object = make_dataobject(
             package=package,
             name=cls_name,
             bases=(DataStreamSourceBase,),
-            attrs={
-                "ELEMENT_TYPE": data_element_type,
-                JsonData.__name__: JsonData,
-                File.__name__: File,
-                ListOfFiles.__name__: ListOfFiles,
-                Directory.__name__: Directory,
-                S3Files.__name__: S3Files,
-            },
-            annotations={
-                "data_stream": Union[
-                    Annotated[JsonData, OneofField(JsonData.__name__.lower())],
-                    Annotated[File, OneofField(File.__name__.lower())],
-                    Annotated[ListOfFiles, OneofField(ListOfFiles.__name__.lower())],
-                    Annotated[Directory, OneofField(Directory.__name__.lower())],
-                    Annotated[S3Files, OneofField(S3Files.__name__.lower())],
-                ],
-            },
+            attrs={"ELEMENT_TYPE": data_element_type, "PLUGINS": plugins, **type_attrs},
+            annotations={"data_stream": data_stream_type_union},
         )
 
         # Add this data stream source to the common data model and the module
diff --git a/docs/adrs/022-pluggable-runtime-io.md b/docs/adrs/022-pluggable-runtime-io.md
new file mode 100644
index 00000000..421821cb
--- /dev/null
+++ b/docs/adrs/022-pluggable-runtime-io.md
@@ -0,0 +1,65 @@
+# ADR N022: Pluggable Runtime IO
+
+Some users of `caikit` will require application-specific logic to fetch data from secure locations. These implementations will not be sufficiently generic to be maintained in the core of `caikit`, but for these security-minded applications, they are a must-have in order to satisfy security posture that fits the existing application.
+
+## Decision
+
+This proposal is to make the Input/Output mechanisms for `caikit.runtime` pluggable. Currently, this means updating `DataStreamSource` and the current `output_path` mechanism to use `oneof` semantics that can be extended by new implementations and bound at boot time with config options.
+
+**Example Implementation**
+```py
+from caikit.core.data_model import DataBase, DataObjectBase, DataStream, dataobject
+from caikit.runtime.service_generation.data_stream_source import DataStreamSourcePlugin, PluginFactory
+from typing import List, Type
+
+@dataobject
+class NumericStreamInput(DataObjectBase):
+    values: List[float]
+
+@dataobject
+class EncodedBytesStreamInput(DataObjectBase):
+    value: bytes
+    encoding: str
+
+class MyDataStreamPlugin(DataStreamSourcePlugin):
+
+    name = "MINE"
+
+    def get_stream_message_type(self, element_type: type) -> Type[DataBase]:
+        if element_type in [int, float]:
+            return NumericStreamInput
+        return EncodedBytesStreamInput
+
+    def get_field_number(self) -> int:
+        return 99 # Make this unique from all other plugins
+
+    def to_data_stream(self, source_message: Type[DataBase], element_type: type) -> DataStream:
+        if element_type in [int, float]:
+            return DataStream.from_iterable(source_message.values).map(element_type)
+        return DataStream.from_iterable(source_message.value.decode(source_message.encoding))
+
+
+# Register the plugin with the factory so it can be configured
+PluginFactory.register(MyDataStreamPlugin)
+```
+
+**Example Config**
+```yaml
+data_streams:
+    source_plugins:
+        inline:
+            type: JsonData
+        custom:
+            type: MINE
+```
+
+## Status
+
+choose one: Accepted
+
+## Consequences
+
+* New `config` section `data_streams.source_plugins` that will control which data stream source types are supported in the running `caikit` environment
+* Equivalent section for `output_path` (naming TBD)
+* An API breaking change to change `output_path` to something like `output_target` to be more generic and support a `oneof` for different output target representations
+* New base classes and factories to register user-defined implementations of the plugins


Test patch:
diff --git a/tests/data_model_helpers.py b/tests/data_model_helpers.py
index 42b9e2a6..92a8e659 100644
--- a/tests/data_model_helpers.py
+++ b/tests/data_model_helpers.py
@@ -124,8 +124,10 @@ def temp_dpool(inherit_global: bool = False, skip_inherit: Optional[List[str]] =
     # Nothing to do for protobuf 3.X
     except ImportError:
         pass
-    yield dpool
-    descriptor_pool._DEFAULT = global_dpool
+    try:
+        yield dpool
+    finally:
+        descriptor_pool._DEFAULT = global_dpool
 
 
 def justify_script_string(script_str):
diff --git a/tests/runtime/service_generation/test_data_stream_plugins.py b/tests/runtime/service_generation/test_data_stream_plugins.py
new file mode 100644
index 00000000..e108e7d6
--- /dev/null
+++ b/tests/runtime/service_generation/test_data_stream_plugins.py
@@ -0,0 +1,238 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Tests for the plugin mechanism for data stream sources
+"""
+# Standard
+from unittest import mock
+
+# Third Party
+import pytest
+
+# First Party
+import aconfig
+
+# Local
+from caikit.core.data_model.streams.data_stream import DataStream
+from caikit.interfaces.common.data_model.stream_sources import (
+    Directory,
+    File,
+    ListOfFiles,
+)
+from caikit.runtime.service_generation.data_stream_source import (
+    DataStreamPluginFactory,
+    DataStreamSourcePlugin,
+    DirectoryDataStreamSourcePlugin,
+    FileDataStreamSourcePlugin,
+    JsonDataStreamSourcePlugin,
+    ListOfFilesDataStreamSourcePlugin,
+    S3FilesDataStreamSourcePlugin,
+    make_data_stream_source,
+)
+from tests.data_model_helpers import reset_global_protobuf_registry, temp_dpool
+from tests.runtime.service_generation.test_data_stream_source import (
+    validate_data_stream,
+)
+import sample_lib
+
+
+def test_file_plugin(sample_json_file):
+    # Sample json file has SampleTrainingType data
+    element_type = sample_lib.data_model.SampleTrainingType
+    source_message = File(filename=sample_json_file)
+    stream = FileDataStreamSourcePlugin({}, "").to_data_stream(
+        source_message, element_type
+    )
+    validate_data_stream(stream, 2, element_type)
+
+
+def test_directory_plugin(sample_json_collection):
+    # Sample json collection has SampleTrainingType data
+    element_type = sample_lib.data_model.SampleTrainingType
+    source_message = Directory(dirname=sample_json_collection, extension="json")
+    stream = DirectoryDataStreamSourcePlugin({}, "").to_data_stream(
+        source_message, element_type
+    )
+    validate_data_stream(stream, 3, element_type)
+
+
+def test_list_of_files_plugin(sample_json_file, sample_csv_file, sample_jsonl_file):
+    # Samples all have SampleTrainingType data
+    element_type = sample_lib.data_model.SampleTrainingType
+    source_message = ListOfFiles(
+        files=[sample_json_file, sample_jsonl_file, sample_csv_file]
+    )
+    stream = ListOfFilesDataStreamSourcePlugin({}, "").to_data_stream(
+        source_message, element_type
+    )
+    validate_data_stream(stream, 9, element_type)
+
+
+def test_json_data_plugin():
+    element_type = sample_lib.data_model.SampleTrainingType
+    foo = sample_lib.data_model.SampleTrainingType(number=1, label="foo")
+    data = [foo] * 5
+
+    # This plugin is actually initialized so that it can build the source message type
+    plugin = JsonDataStreamSourcePlugin({}, "")
+    source_message_type = plugin.get_stream_message_type(element_type=element_type)
+    source_messsage = source_message_type(data=data)
+    stream = plugin.to_data_stream(
+        source_message=source_messsage, element_type=element_type
+    )
+
+    validate_data_stream(stream, 5, element_type)
+
+
+def test_json_data_plugin_can_be_reinitialized():
+    element_type = sample_lib.data_model.SampleTrainingType
+    strm_type = JsonDataStreamSourcePlugin({}, "inst1").get_stream_message_type(
+        element_type=element_type
+    )
+    strm_type2 = JsonDataStreamSourcePlugin({}, "inst2").get_stream_message_type(
+        element_type=element_type
+    )
+
+    assert strm_type is strm_type2
+
+
+## PluginFactory ###############################################################
+
+
+@pytest.fixture
+def reset_stream_source_types():
+    with reset_global_protobuf_registry():
+        with mock.patch(
+            "caikit.runtime.service_generation.data_stream_source._DATA_STREAM_SOURCE_TYPES",
+            {},
+        ):
+            with temp_dpool(True):
+                yield
+
+
+@pytest.fixture
+def plugin_factory():
+    fct = DataStreamPluginFactory("TestFactory")
+    fct.register(JsonDataStreamSourcePlugin)
+    fct.register(FileDataStreamSourcePlugin)
+    fct.register(ListOfFilesDataStreamSourcePlugin)
+    fct.register(DirectoryDataStreamSourcePlugin)
+    fct.register(S3FilesDataStreamSourcePlugin)
+    yield fct
+
+
+class SampleTrainingPlugin(DataStreamSourcePlugin):
+    """Dummy plugin that can be registered with a factory"""
+
+    name = "TEST"
+
+    def get_stream_message_type(self, *_, **__):
+        return sample_lib.data_model.SampleTrainingType
+
+    def to_data_stream(self, source_message, *_, **__):
+        count = self._config.get("count", 5)
+        DataStream.from_iterable(count * [source_message])
+
+    def get_field_number(self) -> int:
+        return self._config.get("field_number", 42)
+
+    def get_field_name(self, element_type):
+        return self._config.get("field_name", super().get_field_name(element_type))
+
+
+def test_plugin_factory_is_configurable(plugin_factory):
+    """Make sure new plugins can be registered and configured"""
+    cfg = aconfig.Config(
+        {"foo": {"type": SampleTrainingPlugin.name}},
+        override_env_vars=False,
+    )
+    # Before registration, it's a ValueError
+    with pytest.raises(ValueError):
+        plugin_factory.construct(cfg.foo, "foo")
+
+    # After registration, it's valid
+    plugin_factory.register(SampleTrainingPlugin)
+    plug = plugin_factory.construct(cfg.foo, "foo")
+    assert isinstance(plug, SampleTrainingPlugin)
+
+    # Make sure there's an instance in get_plugins
+    all_plugins = plugin_factory.get_plugins(cfg)
+    assert any(isinstance(plug, SampleTrainingPlugin) for plug in all_plugins)
+
+
+def test_plugin_factory_reuses_plugins(plugin_factory):
+    """Make sure that the list of plugins gets reused"""
+    plugins = plugin_factory.get_plugins()
+    assert plugin_factory.get_plugins() is plugins
+
+
+def test_no_duplicate_field_numbers(plugin_factory):
+    """Make sure that duplicate field numbers cause a ValueError"""
+    plugin_factory.register(SampleTrainingPlugin)
+    cfg = aconfig.Config(
+        {
+            "foo": {"type": SampleTrainingPlugin.name},
+            "bar": {"type": SampleTrainingPlugin.name},
+        },
+        override_env_vars=False,
+    )
+    with pytest.raises(ValueError):
+        plugin_factory.get_plugins(cfg)
+
+
+def test_no_duplicate_field_names(plugin_factory, reset_stream_source_types):
+    """Make sure that new data stream sources cannot be created with duplicate
+    field names
+    """
+    plugin_factory.register(SampleTrainingPlugin)
+    cfg = aconfig.Config(
+        {
+            "foo": {"type": SampleTrainingPlugin.name, "config": {"field_number": 42}},
+            "bar": {"type": SampleTrainingPlugin.name, "config": {"field_number": 43}},
+        },
+        override_env_vars=False,
+    )
+    with pytest.raises(ValueError):
+        make_data_stream_source(int, plugin_factory, cfg)
+
+
+def test_multiple_instances(plugin_factory, reset_stream_source_types):
+    """Make sure that multiple instances of a type with different configs can be
+    instantiated as long as field numbers and names are unique
+    """
+    plugin_factory.register(SampleTrainingPlugin)
+    cfg = aconfig.Config(
+        {
+            "foo": {
+                "type": SampleTrainingPlugin.name,
+                "config": {"field_number": 42, "field_name": "foo"},
+            },
+            "bar": {
+                "type": SampleTrainingPlugin.name,
+                "config": {"field_number": 43, "field_name": "bar"},
+            },
+        },
+        override_env_vars=False,
+    )
+    plugin_factory.get_plugins(cfg)
+
+    class LocalFoobar:
+        pass
+
+    strm_src_cls = make_data_stream_source(LocalFoobar, plugin_factory, cfg)
+    assert strm_src_cls.ELEMENT_TYPE is LocalFoobar
+    assert set(strm_src_cls.get_proto_class().DESCRIPTOR.fields_by_name.keys()) == {
+        "foo",
+        "bar",
+    }
diff --git a/tests/runtime/service_generation/test_data_stream_source.py b/tests/runtime/service_generation/test_data_stream_source.py
index 15bae671..762d259f 100644
--- a/tests/runtime/service_generation/test_data_stream_source.py
+++ b/tests/runtime/service_generation/test_data_stream_source.py
@@ -26,6 +26,15 @@ import caikit
 ## Helper functions
 
 
+@pytest.fixture(autouse=True)
+def train_service(sample_train_service):
+    """This autoused fixture ensures that the training APIs will be created
+    when individual tests are run. Each test, however, does not
+    use the fixture explicitly
+    """
+    pass
+
+
 def validate_data_stream(data_stream, length, data_item_type, data_item_length=None):
     assert isinstance(data_stream, DataStream)
     assert len(data_stream) == length
@@ -65,16 +74,6 @@ def test_multiple_make_data_stream_source():
     assert stream_type.from_proto(proto_repr).to_proto() == proto_repr
 
 
-def test_data_model_element_type():
-    stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
-    assert isinstance(stream_type._to_element_type({"number": 1}), SampleTrainingType)
-
-
-def test_primitive_element_type():
-    stream_type = caikit.interfaces.common.data_model.DataStreamSourceInt
-    assert isinstance(stream_type._to_element_type(1), int)
-
-
 def test_make_data_stream_source_types():
     assert issubclass(make_data_stream_source(int), DataStreamSourceBase)
     assert issubclass(make_data_stream_source(float), DataStreamSourceBase)