Use stream_config instead of kwargs

mosaicml · Jan 7, 2025 · 52de4dc · 52de4dc
1 parent 607bd9c
commit 52de4dc
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 10 deletions.
diff --git a/docs/source/dataset_configuration/mixing_data_sources.md b/docs/source/dataset_configuration/mixing_data_sources.md
@@ -29,6 +29,7 @@ dataset = StreamingDataset(
     remote='s3://some/path',
     local='/local/path',
     stream_name='my_stream',
+    stream_config={'arg1': 'val1'},
 )
 ```
 

diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py
@@ -311,7 +311,7 @@ class StreamingDataset(Array, IterableDataset):
             devices need to see the same partition of the dataset. Defaults to ``None``.
         stream_name (str): The name of the Stream to use which is registered in streams_registry.
             Defaults to ``stream``.
-        kwargs (any): Additional arguments to pass to the Stream constructor.
+        stream_config (dict[str, any]): Additional arguments to pass to the Stream constructor.
     """
 
     def __init__(self,
@@ -340,7 +340,7 @@ def __init__(self,
                  allow_unsafe_types: bool = False,
                  replication: Optional[int] = None,
                  stream_name: str = 'stream',
-                 **kwargs: Any) -> None:
+                 stream_config: Optional[dict[str, Any]] = None) -> None:
         # Global arguments (which do not live in Streams).
         self.predownload = predownload
         self.cache_limit = cache_limit
@@ -444,16 +444,16 @@ def __init__(self,
             for stream in streams:
                 stream.apply_default(default)
         else:
-            kwargs = {
+            stream_config = stream_config or {}
+            stream_config.update({
                 'remote': remote,
                 'local': local,
                 'split': split,
                 'download_retry': download_retry,
                 'download_timeout': download_timeout,
                 'validate_hash': validate_hash,
                 'keep_zip': keep_zip,
-                **kwargs,
-            }
+            })
 
             # Construct a Stream instance using registry-based construction
             default = construct_from_registry(
@@ -462,7 +462,7 @@ def __init__(self,
                 partial_function=False,
                 pre_validation_function=None,
                 post_validation_function=None,
-                kwargs=kwargs,
+                kwargs=stream_config,
             )
 
             streams = [default]

diff --git a/tests/test_streaming.py b/tests/test_streaming.py
@@ -1098,10 +1098,14 @@ def __init__(
 
     streams_registry.register('custom_stream', func=CustomStream)
 
-    dataset = StreamingDataset(local=local_dir,
-                               remote=remote_dir,
-                               stream_name='custom_stream',
-                               custom_arg=100)
+    dataset = StreamingDataset(
+        local=local_dir,
+        remote=remote_dir,
+        stream_name='custom_stream',
+        stream_config={
+            'custom_arg': 100,
+        },
+    )
 
     assert len(dataset.streams) == 1
     assert isinstance(dataset.streams[0], CustomStream)