Repository: caikit/caikit
Pull request #365 (2023-07-25 20:00:01)

✨ S3 streams
Implements #363 , https://github.com/caikit/caikit/issues/300

This PR adds an `S3Files` source for datastreams.
This is not implemented in the core of caikit yet, we don't have support for pulling files from S3. But it's left open to extension for other projects to implement as needed.

This also removes the dynamic generation of the `Files` `ListOfFiles` and `Directory` data models, in favor of static class definitions.


Patch:
diff --git a/caikit/interfaces/common/data_model/stream_sources.py b/caikit/interfaces/common/data_model/stream_sources.py
new file mode 100644
index 00000000..5eef5c42
--- /dev/null
+++ b/caikit/interfaces/common/data_model/stream_sources.py
@@ -0,0 +1,69 @@
+# 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.
+"""
+This file contains interfaces required to generate DataStreamSource[T] classes
+"""
+
+# Standard
+from typing import List
+
+# First Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
+
+# Local
+from caikit.core.data_model import PACKAGE_COMMON, DataObjectBase, dataobject
+
+
+@dataobject(PACKAGE_COMMON)
+class File(DataObjectBase):
+    filename: Annotated[str, FieldNumber(1)]
+
+
+@dataobject(PACKAGE_COMMON)
+class ListOfFiles(DataObjectBase):
+    files: Annotated[List[str], FieldNumber(1)]
+
+
+@dataobject(PACKAGE_COMMON)
+class Directory(DataObjectBase):
+    dirname: Annotated[str, FieldNumber(1)]
+    extension: Annotated[str, FieldNumber(2)]
+
+
+@dataobject(PACKAGE_COMMON)
+class S3Base(DataObjectBase):
+    # URI info
+    endpoint: Annotated[str, FieldNumber(2)]  # begins with `http://` or `https://`
+    region: Annotated[str, FieldNumber(3)]
+    bucket: Annotated[str, FieldNumber(4)]
+
+    # HMAC credentials
+    accessKey: Annotated[str, FieldNumber(5)]
+    secretKey: Annotated[str, FieldNumber(6)]
+
+    # IAM credentials
+    IAM_id: Annotated[str, FieldNumber(7)]
+    IAM_api_key: Annotated[str, FieldNumber(8)]
+
+
+@dataobject(PACKAGE_COMMON)
+class S3Files(S3Base):
+    # List of file paths relative to the bucket
+    files: Annotated[List[str], FieldNumber(1)]
+
+
+@dataobject(PACKAGE_COMMON)
+class S3Path(S3Base):
+    # Path relative to the bucket
+    path: Annotated[str, FieldNumber(1)]
diff --git a/caikit/runtime/service_generation/data_stream_source.py b/caikit/runtime/service_generation/data_stream_source.py
index 2d83df60..37ef65ec 100644
--- a/caikit/runtime/service_generation/data_stream_source.py
+++ b/caikit/runtime/service_generation/data_stream_source.py
@@ -29,6 +29,13 @@ import alog
 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.toolkit.errors import error_handler
+from caikit.interfaces.common.data_model.stream_sources import (
+    Directory,
+    File,
+    ListOfFiles,
+    S3Files,
+)
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 import caikit
 
@@ -41,6 +48,7 @@ import caikit.interfaces.common
 _DATA_STREAM_SOURCE_TYPES = {}
 
 log = alog.use_channel("DSTRM-SRC")
+error = error_handler.get(log)
 
 
 class DataStreamSourceBase(DataStream):
@@ -89,6 +97,15 @@ class DataStreamSourceBase(DataStream):
             log.debug3("Returning empty data stream")
             return DataStream.from_iterable([])
 
+        # 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."
+                ),
+            )
+
         # If jsondata, pull from the data elements directly
         if set_field == "jsondata":
             log.debug3("Pulling data stream from inline json")
@@ -194,11 +211,12 @@ def make_data_stream_source(data_element_type: Type) -> Type[DataBase]:
         cls_name = _make_data_stream_source_type_name(data_element_type)
         log.debug("Creating DataStreamSource[%s] -> %s", data_element_type, cls_name)
 
-        # Set up the "sub classes." In python, this is the same as creating a
+        # 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
-        # these classes can be referenced within the Union annotation for the
-        # outer class itself.
+        # 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 = "caikit_data_model.runtime"
         JsonData = make_dataobject(
             package=package,
@@ -207,27 +225,6 @@ def make_data_stream_source(data_element_type: Type) -> Type[DataBase]:
             attrs={"__qualname__": f"{cls_name}.JsonData"},
             annotations={"data": List[data_element_type]},
         )
-        File = make_dataobject(
-            package=package,
-            proto_name=f"{cls_name}File",
-            name="File",
-            attrs={"__qualname__": f"{cls_name}.File"},
-            annotations={"filename": str},
-        )
-        ListOfFiles = make_dataobject(
-            package=package,
-            proto_name=f"{cls_name}ListOfFiles",
-            name="ListOfFiles",
-            attrs={"__qualname__": f"{cls_name}.ListOfFiles"},
-            annotations={"files": List[str]},
-        )
-        Directory = make_dataobject(
-            package=package,
-            proto_name=f"{cls_name}Directory",
-            name="Directory",
-            attrs={"__qualname__": f"{cls_name}.Directory"},
-            annotations={"dirname": str, "extension": str},
-        )
 
         # Create the outer class that encapsulates the Union (oneof) or the
         # various types of input sources
@@ -241,6 +238,7 @@ def make_data_stream_source(data_element_type: Type) -> Type[DataBase]:
                 File.__name__: File,
                 ListOfFiles.__name__: ListOfFiles,
                 Directory.__name__: Directory,
+                S3Files.__name__: S3Files,
             },
             annotations={
                 "data_stream": Union[
@@ -248,6 +246,7 @@ def make_data_stream_source(data_element_type: Type) -> Type[DataBase]:
                     Annotated[File, OneofField(File.__name__.lower())],
                     Annotated[ListOfFiles, OneofField(ListOfFiles.__name__.lower())],
                     Annotated[Directory, OneofField(Directory.__name__.lower())],
+                    Annotated[S3Files, OneofField(S3Files.__name__.lower())],
                 ],
             },
         )


Test patch:
diff --git a/tests/runtime/service_generation/test_data_stream_source.py b/tests/runtime/service_generation/test_data_stream_source.py
index d7125bb2..182b03a8 100644
--- a/tests/runtime/service_generation/test_data_stream_source.py
+++ b/tests/runtime/service_generation/test_data_stream_source.py
@@ -12,6 +12,7 @@ import pytest
 # Local
 from caikit.core.data_model import DataObjectBase, dataobject
 from caikit.core.data_model.streams.data_stream import DataStream
+from caikit.interfaces.common.data_model.stream_sources import S3Files
 from caikit.runtime.service_generation.data_stream_source import (
     DataStreamSourceBase,
     _make_data_stream_source_type_name,
@@ -507,3 +508,22 @@ def test_make_data_stream_source_no_files_w_ext_dir(
     with pytest.raises(CaikitRuntimeException) as e:
         ds.to_data_stream()
     assert "contains no source files with extension" in e.value.message
+
+
+def test_s3_not_implemented(sample_train_service, sample_jsonl_dir):
+    stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
+    ds = stream_type(s3files=S3Files())
+    # Explicit .to_data_stream will fail
+    with pytest.raises(
+        NotImplementedError,
+        match="S3Files are not implemented as stream sources in this runtime.",
+    ) as e:
+        ds.to_data_stream()
+
+    # And so would iterating on the data stream source directly
+    with pytest.raises(
+        NotImplementedError,
+        match="S3Files are not implemented as stream sources in this runtime.",
+    ) as e:
+        for val in ds:
+            _ = val