Repository: caikit/caikit
Pull request #553 (2023-10-30 23:43:41)

:zap::bug: Add proper caching of streams in sources
<!--  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**:

We came across two issues while testing custom plugins for DataStreamSources:
1. The `.to_data_stream` method is called on each read through the stream, which is really slow when your plugin implementation is reading from a remote resource, and unnecessary when your plugin already cached the remote data locally for stream re-entrancy.
2.  Calling `len` on unpickled DataStreamSources fails, because the base `DataStream` class sets a `_length` attribute that is not re-created on the DataStreamSourceBase's `__set_state__`

This PR caches both the `_length` property and the result of calling `to_data_stream` to fix both

**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/data_model/streams/data_stream.py b/caikit/core/data_model/streams/data_stream.py
index deee46cb..baf9f48f 100644
--- a/caikit/core/data_model/streams/data_stream.py
+++ b/caikit/core/data_model/streams/data_stream.py
@@ -17,6 +17,7 @@
 """
 # Standard
 from collections.abc import Iterable
+from functools import cached_property
 from glob import glob
 from io import UnsupportedOperation
 from typing import Dict, Generic, List, Tuple, TypeVar, Union
@@ -103,7 +104,6 @@ class DataStream(Generic[T]):
 
         self.generator_func = generator_func
         self.generator_args, self.generator_kwargs = args, kwargs
-        self._length = None
 
     @classmethod
     def from_iterable(cls, data: typing.Iterable[T]) -> "DataStream[T]":
@@ -995,14 +995,20 @@ class DataStream(Generic[T]):
         return generator
 
     def __len__(self):
+        """See property method self._length"""
+        return self._length
+
+    @cached_property
+    def _length(self):
         """Return the number of data items contained in this data stream.  This requires that the
-        data stream be iterated over, which may be time consuming.  This value is then stored
+        data stream be iterated over, which may be time-consuming.  This value is then stored
         internally so that subsequent calls do not iterate over the data stream again.
-        """
-        if self._length is None:
-            self._length = sum(1 for data_item in self)
 
-        return self._length
+        This is implemented as a cached_property so that subclasses of DataStream which implement
+        their own __getstate__ and __setstate__ do not have to account for the existence of
+        self._length
+        """
+        return sum(1 for _ in self)
 
     def __or__(self, module):
         """Feed this data stream into the `.stream` method of a module.  This is syntactic sugar
diff --git a/caikit/runtime/service_generation/data_stream_source.py b/caikit/runtime/service_generation/data_stream_source.py
index fcbd216e..d9ec73c5 100644
--- a/caikit/runtime/service_generation/data_stream_source.py
+++ b/caikit/runtime/service_generation/data_stream_source.py
@@ -422,8 +422,9 @@ class DataStreamSourceBase(DataStream):
         super().__init__(self._generator)
 
     def _generator(self):
-        stream = self.to_data_stream()
-        return stream.generator_func(*stream.generator_args, **stream.generator_kwargs)
+        return self._stream.generator_func(
+            *self._stream.generator_args, **self._stream.generator_kwargs
+        )
 
     def __getstate__(self) -> bytes:
         """A DataStreamSource is pickled by serializing its source
@@ -450,6 +451,13 @@ class DataStreamSourceBase(DataStream):
             plugin.get_field_name(self.ELEMENT_TYPE): plugin for plugin in self.PLUGINS
         }
 
+    @cached_property
+    def _stream(self):
+        """The internal _stream is cached here so that the result of calling to_data_stream can be
+        re-read, rather than requiring to_data_stream to be invoked on every read through the
+        stream"""
+        return self.to_data_stream()
+
     # pylint: disable=too-many-return-statements
     def to_data_stream(self) -> DataStream:
         """Convert to the target data stream type based on the source type"""


Test patch:
diff --git a/tests/core/data_model/streams/test_data_stream.py b/tests/core/data_model/streams/test_data_stream.py
index d3756e7d..1eed24a3 100644
--- a/tests/core/data_model/streams/test_data_stream.py
+++ b/tests/core/data_model/streams/test_data_stream.py
@@ -167,6 +167,26 @@ def test_data_stream_from_jsonl_is_pickleable(tmp_path):
     post_pickle_vals = list(pickled_stream)
 
     assert pre_pickle_vals == post_pickle_vals
+    # Interesting: Technically this is a stream of length 1 where the one element is [1,2,3,4,5,6]
+    validate_data_stream(pickled_stream, 1, list)
+
+
+def test_data_stream_from_json_is_pickleable(tmp_path):
+    tmpdir = str(tmp_path)
+
+    data = [1, 2, 3, 4, 5, 6]
+    filepath = os.path.join(tmpdir, "foo.json")
+    with open(filepath, "w") as f:
+        json.dump(data, f)
+
+    stream = DataStream.from_json_array(filepath)
+
+    pre_pickle_vals = list(stream)
+    pickled_stream = pickle.loads(pickle.dumps(stream))
+    post_pickle_vals = list(pickled_stream)
+
+    assert pre_pickle_vals == post_pickle_vals
+    validate_data_stream(pickled_stream, 6, int)
 
 
 def test_bad_json_stream(tmp_path):
diff --git a/tests/runtime/service_generation/test_data_stream_source.py b/tests/runtime/service_generation/test_data_stream_source.py
index 3564bc0c..62c91b01 100644
--- a/tests/runtime/service_generation/test_data_stream_source.py
+++ b/tests/runtime/service_generation/test_data_stream_source.py
@@ -104,6 +104,8 @@ def test_pickle_round_trip_primitive():
     round_trip = pickle.loads(pickle.dumps(inst))
     assert round_trip.to_dict() == inst.to_dict()
 
+    validate_data_stream(round_trip, 2, float)
+
 
 def test_pickle_round_trip_data_model():
     """Make sure that a source wrapping a data model object can be round-tripped
@@ -122,6 +124,17 @@ def test_pickle_round_trip_data_model():
     round_trip = pickle.loads(pickle.dumps(inst))
     assert round_trip.to_dict() == inst.to_dict()
 
+    validate_data_stream(round_trip, 2, Foo)
+
+
+def test_pickle_round_trip_file(sample_json_file):
+    stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
+    data_stream = stream_type(file=FileReference(filename=sample_json_file))
+
+    round_trip = pickle.loads(pickle.dumps(data_stream))
+
+    validate_data_stream(round_trip, 2, SampleTrainingType)
+
 
 def test_data_stream_source_as_data_stream():
     """Make sure that a DataStreamSource works exactly like a DataStream"""
@@ -544,3 +557,20 @@ def test_s3_not_implemented():
     ) as e:
         for val in ds:
             _ = val
+
+
+def test_datastream_sources_not_repeatedly_read(sample_json_file):
+    """This test ensures that `to_data_stream` is only called once on a single instance of a
+    DataStreamSource. This allows source plugin authors to control how data is cached for the
+    life of the stream"""
+    stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
+    data_stream = stream_type(file=FileReference(filename=sample_json_file))
+
+    # Read the datastream as normal
+    validate_data_stream(data_stream, 2, SampleTrainingType)
+
+    # Delete the set source field so that .to_data_stream does not know how to find a source
+    data_stream.file = None
+
+    # Check that we can still read through the stream
+    validate_data_stream(data_stream, 2, SampleTrainingType)