Repository: caikit/caikit
Issue: Closures in DataStreams prevent pickling

## Describe the bug

The DataStream's `from_*` methods all create closures around local functions to provide the `generator_func` argument to the initializer. This prevents `DataStream`s created in this manner from being pickled. Pickling is important for at least one specific use case for us- using `ray` to execute things remotely. Not being able to pickle streams effectively blocks us from being able to write a shared train backend with ray, since most training things involve streams of training data.

## Platform

Please provide details about the environment you are using, including the following:

- Interpreter version: all
- Library version: 0.5

## Sample Code

```
stream = DataStream.from_iterable([1,2,3])
pickle.dumps(stream)
```

## Expected behavior

it should pickle

## Observed behavior

it does not

## Additional context

The `DataStream` class was designed to hold args and kwargs for the generator function as private member variable anyway, so closures aren't required to implement these `from_*` classmethods

Changed files:
tests/core/data_model/streams/test_data_stream.py, caikit/runtime/service_generation/data_stream_source.py, caikit/core/data_model/streams/data_stream.py

diff --git a/caikit/core/data_model/streams/data_stream.py b/caikit/core/data_model/streams/data_stream.py
index 0b9b7af..7b06888 100644
--- a/caikit/core/data_model/streams/data_stream.py
+++ b/caikit/core/data_model/streams/data_stream.py
@@ -125,9 +125,11 @@ class DataStream(Generic[T]):
             3
         """
         error.type_check("<COR88684982E>", Iterable, data=data)
+        return cls(cls._from_iterable_generator, data)
 
-        # lambda closure over `data` that returns an iterator over `data`
-        return cls(lambda: iter(data))
+    @classmethod
+    def _from_iterable_generator(cls, data: typing.Iterable[T]) -> typing.Iterator[T]:
+        return iter(data)
 
     @classmethod
     def from_jsonl(cls, filename: str) -> "DataStream[Dict]":
@@ -168,28 +170,28 @@ class DataStream(Generic[T]):
         """
         error.file_check("<COR32600575E>", filename)
 
-        def generator_func():
-            # open the file (closure around `filename`)
-            with open(filename, mode="rb") as json_fh:
-                log.debug2("Loading JSON array file: {}".format(filename))
-                lines = json_fh.readlines()
-
-                try:
-                    for line in lines:
-                        if line.strip():  # ignore empty lines
-                            yield json.loads(line)
-                except json.JSONDecodeError as e:
-                    error(
-                        "<COR55596551E>",
-                        ValueError(f"Invalid JSON object in `{line}`, error: {e.msg}"),
-                    )
-                except TypeError:
-                    error(
-                        "<COR35596551E>",
-                        ValueError("Invalid JSON object in `{}`".format(line)),
-                    )
+        return cls(cls._from_jsonl_generator, filename)
 
-        return cls(generator_func)
+    @classmethod
+    def _from_jsonl_generator(cls, filename):
+        with open(filename, mode="rb") as json_fh:
+            log.debug2("Loading JSON array file:  %s", filename)
+            lines = json_fh.readlines()
+
+            try:
+                for line in lines:
+                    if line.strip():  # ignore empty lines
+                        yield json.loads(line)
+            except json.JSONDecodeError as e:
+                error(
+                    "<COR55596551E>",
+                    ValueError(f"Invalid JSON object in `{line}`, error: {e.msg}"),
+                )
+            except TypeError:
+                error(
+                    "<COR35596551E>",
+                    ValueError("Invalid JSON object in `{}`".format(line)),
+                )
 
     @classmethod
     def from_json_array(cls, filename: str) -> "DataStream[Dict]":
@@ -228,24 +230,25 @@ class DataStream(Generic[T]):
         """
         error.file_check("<COR39609575E>", filename)
 
-        def generator_func():
-            # open the file (closure around `filename`)
-            with open(filename, mode="rb") as json_fh:
-                log.debug2("Loading JSON array file: {}".format(filename))
-
-                # for each {} object of the array
-                try:
-                    for item_idx, obj in enumerate(ijson.items(json_fh, "item")):
-                        log.debug2("Loading object index %d", item_idx)
-                        yield obj
+        return cls(cls._from_json_array_generator, filename)
 
-                except ijson.JSONError:
-                    error(
-                        "<COR85596551E>",
-                        ValueError("Invalid JSON object in `{}`".format(filename)),
-                    )
-
-        return cls(generator_func)
+    @classmethod
+    def _from_json_array_generator(cls, filename):
+        # open the file
+        with open(filename, mode="rb") as json_fh:
+            log.debug2("Loading JSON array file: %s", filename)
+
+            # for each {} object of the array
+            try:
+                for item_idx, obj in enumerate(ijson.items(json_fh, "item")):
+                    log.debug2("Loading object index %d", item_idx)
+                    yield obj
+
+            except ijson.JSONError:
+                error(
+                    "<COR85596551E>",
+                    ValueError("Invalid JSON object in `{}`".format(filename)),
+                )
 
     @classmethod
     def from_csv(cls, filename: str, *args, skip=0, **kwargs) -> "DataStream[List]":
@@ -288,20 +291,20 @@ class DataStream(Generic[T]):
                 ),
             )
 
-        # generator function that yields lists each containing the columns of the csv file
-        def generator_func(*csv_args, **csv_kwargs):
-            # open the csv file (closure around `filename`)
-            with open(filename, mode="r", encoding="utf8") as fh:
-                # skip lines if requested
-                for _ in range(skip):
-                    # pylint: disable=stop-iteration-return
-                    next(fh)
-
-                # for each line of the csv file, yield a list
-                for line in csv.reader(fh, *csv_args, **csv_kwargs):
-                    yield line
+        return cls(cls._from_csv_generator, filename, skip, *args, **kwargs)
 
-        return cls(generator_func, *args, **kwargs)
+    @classmethod
+    def _from_csv_generator(cls, filename, skip, *csv_args, **csv_kwargs):
+        # open the csv file (closure around `filename`)
+        with open(filename, mode="r", encoding="utf8") as fh:
+            # skip lines if requested
+            for _ in range(skip):
+                # pylint: disable=stop-iteration-return
+                next(fh)
+
+            # for each line of the csv file, yield a list
+            for line in csv.reader(fh, *csv_args, **csv_kwargs):
+                yield line
 
     @classmethod
     def from_header_csv(cls, filename: str, *args, **kwargs) -> "DataStream[Dict]":
@@ -343,16 +346,16 @@ class DataStream(Generic[T]):
                 ),
             )
 
-        # generator function that yields dicts
-        def generator_func(*csv_args, **csv_kwargs):
-            # open the csv file (closure around `filename`)
-            with open(filename, mode="r", encoding="utf8") as fh:
+        return cls(cls._from_header_csv_generator, filename, *args, **kwargs)
 
-                # for each line of the csv file, yield a dict
-                for line in csv.DictReader(fh, *csv_args, **csv_kwargs):
-                    yield line
+    @classmethod
+    def _from_header_csv_generator(cls, filename, *csv_args, **csv_kwargs):
+        # open the csv file (closure around `filename`)
+        with open(filename, mode="r", encoding="utf8") as fh:
 
-        return cls(generator_func, *args, **kwargs)
+            # for each line of the csv file, yield a dict
+            for line in csv.DictReader(fh, *csv_args, **csv_kwargs):
+                yield line
 
     @classmethod
     def from_txt(cls, filename: str) -> "DataStream[str]":
@@ -388,15 +391,16 @@ class DataStream(Generic[T]):
         """
         error.file_check("<COR79693043E>", filename)
 
-        def generator_func():
-            # open the file (closure around `filename`)
-            with open(filename, mode="r", encoding="utf8") as fh:
-                # for each line of the file
-                for line in fh:
-                    # strip new lines and carriage returns and yield the line
-                    yield line.rstrip("\\n\\r")
+        return cls(cls._from_txt_generator, filename)
 
-        return cls(generator_func)
+    @classmethod
+    def _from_txt_generator(cls, filename):
+        # open the file (closure around `filename`)
+        with open(filename, mode="r", encoding="utf8") as fh:
+            # for each line of the file
+            for line in fh:
+                # strip new lines and carriage returns and yield the line
+                yield line.rstrip("\\n\\r")
 
     @classmethod
     def from_file(cls, filename: str) -> "DataStream[Union[Dict, Tuple, str]]":
@@ -420,18 +424,14 @@ class DataStream(Generic[T]):
         _, file_ext = os.path.splitext(filename)
         # choose the right from_* fn
         if file_ext.lower() == ".json":
-            log.debug2(
-                "Detected .json extension, loading {} as a json file".format(filename)
-            )
+            log.debug2("Detected .json extension, loading %s as a json file", filename)
             return DataStream.from_json_array(filename)
 
         if file_ext.lower() == ".csv":
-            log.debug2(
-                "Detected .csv extension, loading {} as a csv file".format(filename)
-            )
+            log.debug2("Detected .csv extension, loading %s as a csv file", filename)
             return DataStream.from_csv(filename)
 
-        log.debug2("Loading {} as a raw text file".format(filename))
+        log.debug2("Loading %s as a raw text file", filename)
         # TODO: test this at some point (this path is unused currently)
         return DataStream.from_txt(filename)
 
@@ -465,12 +465,13 @@ class DataStream(Generic[T]):
         # verify that `dirname` exists
         cls._verify_dir(dirname)
 
-        def generator_func():
-            # glob `*.txt` files in `dirname` (closure around `dirname`)
-            for filename in glob(os.path.join(dirname, "*." + extension)):
-                yield file_opener(filename)
+        return cls(cls._from_collection_generator, dirname, extension, file_opener)
 
-        return cls(generator_func)
+    @classmethod
+    def _from_collection_generator(cls, dirname, extension, file_opener):
+        # glob `*.txt` files in `dirname` (closure around `dirname`)
+        for filename in glob(os.path.join(dirname, "*." + extension)):
+            yield file_opener(filename)
 
     @classmethod
     def from_txt_collection(cls, dirname: str, extension="txt") -> "DataStream[str]":
@@ -540,17 +541,18 @@ class DataStream(Generic[T]):
         # verify that `dirname` exists
         cls._verify_dir(dirname)
 
-        def generator_func():
-            # list of data_streams created from different files
-            data_stream_list = []
-            # glob `*.txt` files in `dirname` (closure around `dirname`)
-            for filename in glob(os.path.join(dirname, "*.csv")):
-                data_stream_list.append(cls.from_header_csv(filename=filename))
-            # yield the combined data item once flattened
-            for data_item in DataStream.chain(data_stream_list).flatten():
-                yield data_item
+        return cls(cls._from_csv_collection_generator, dirname)
 
-        return cls(generator_func)
+    @classmethod
+    def _from_csv_collection_generator(cls, dirname):
+        # list of data_streams created from different files
+        data_stream_list = []
+        # glob `*.txt` files in `dirname` (closure around `dirname`)
+        for filename in glob(os.path.join(dirname, "*.csv")):
+            data_stream_list.append(cls.from_header_csv(filename=filename))
+        # yield the combined data item once flattened
+        for data_item in DataStream.chain(data_stream_list).flatten():
+            yield data_item
 
     @classmethod
     def from_jsonl_collection(cls, dirname: str) -> "DataStream[Dict]":
@@ -569,17 +571,18 @@ class DataStream(Generic[T]):
         # verify that `dirname` exists
         cls._verify_dir(dirname)
 
-        def generator_func():
-            # list of data_streams created from different files
-            data_stream_list = []
-            # glob `*.txt` files in `dirname` (closure around `dirname`)
-            for filename in glob(os.path.join(dirname, "*.jsonl")):
-                data_stream_list.append(cls.from_jsonl(filename=filename))
-            # yield the combined data item once flattened
-            for data_item in DataStream.chain(data_stream_list).flatten():
-                yield data_item
-
-        return cls(generator_func)
+        return cls(cls._from_jsonl_collection_generator, dirname)
+
+    @classmethod
+    def _from_jsonl_collection_generator(cls, dirname):
+        # list of data_streams created from different files
+        data_stream_list = []
+        # glob `*.txt` files in `dirname` (closure around `dirname`)
+        for filename in glob(os.path.join(dirname, "*.jsonl")):
+            data_stream_list.append(cls.from_jsonl(filename=filename))
+        # yield the combined data item once flattened
+        for data_item in DataStream.chain(data_stream_list).flatten():
+            yield data_item
 
     def train_test_split(
         self, test_split=0.25, seed=None
diff --git a/caikit/runtime/service_generation/data_stream_source.py b/caikit/runtime/service_generation/data_stream_source.py
index 38c7fe1..eac3c8b 100644
--- a/caikit/runtime/service_generation/data_stream_source.py
+++ b/caikit/runtime/service_generation/data_stream_source.py
@@ -63,8 +63,11 @@ class DataStreamSourceBase(DataStream):
     """
 
     def __init__(self):
-        """Validate oneof semantics"""
-        super().__init__(lambda: self.to_data_stream().generator_func())
+        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
diff --git a/tests/core/data_model/streams/test_data_stream.py b/tests/core/data_model/streams/test_data_stream.py
index 746ffa1..8c57dcc 100644
--- a/tests/core/data_model/streams/test_data_stream.py
+++ b/tests/core/data_model/streams/test_data_stream.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 
 # Standard
+import json
 import os
+import pickle
 
 # Local
 from caikit.core import data_model as core_dm
@@ -53,6 +55,23 @@ def build_test_augmentor(produces_none):
     return TestStreamAugmentor()
 
 
+def test_data_stream_from_jsonl_is_pickleable(tmp_path):
+    tmpdir = str(tmp_path)
+
+    data = [1, 2, 3, 4, 5, 6]
+    filepath = os.path.join(tmpdir, "foo.jsonl")
+    with open(filepath, "w") as f:
+        json.dump(data, f)
+
+    stream = core_dm.DataStream.from_jsonl(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
+
+
 class TestDataStream(TestCaseBase):
     def setUp(self):
         self.list_data_stream = core_dm.DataStream.from_iterable(