Bump fsspec and simpler wds compr

setup.py

@@ -127,7 +127,7 @@
     "multiprocess<0.70.20",  # to align with dill<0.3.9 (see above)
     # to save datasets locally or on any filesystem
     # minimum 2023.1.0 to support protocol=kwargs in fsspec's `open`, `get_fs_token_paths`, etc.: see https://github.com/fsspec/filesystem_spec/pull/1143
-    "fsspec[http]>=2023.1.0,<=2026.4.0",
+    "fsspec[http]>=2023.1.0,<=2026.6.0",
     # To get datasets from the Datasets Hub on huggingface.co
     "huggingface-hub>=0.25.0,<2.0",
     # Utilities from PyPA to e.g., compare versions

src/datasets/filesystems/__init__.py

@@ -1,4 +1,5 @@
 import importlib
+import io
 import shutil
 import warnings
 from typing import List
@@ -18,11 +19,15 @@
     compression.ZstdFileSystem,
 ]
 
+EXTENSION_TO_COMPRESSION_FS_FILE_CLS: dict[str, type[io.BytesIO]] = {}
 # Register custom filesystems
 for fs_class in COMPRESSION_FILESYSTEMS:
     if fs_class.protocol in fsspec.registry and fsspec.registry[fs_class.protocol] is not fs_class:
         warnings.warn(f"A filesystem protocol was already set for {fs_class.protocol} and will be overwritten.")
     fsspec.register_implementation(fs_class.protocol, fs_class, clobber=True)
+    for extension in fs_class.extensions:
+        if fs_class.compression in fsspec.compression.compr:
+            EXTENSION_TO_COMPRESSION_FS_FILE_CLS[extension] = fsspec.compression.compr[fs_class.compression]
 
 
 def is_remote_filesystem(fs: fsspec.AbstractFileSystem) -> bool:

src/datasets/packaged_modules/webdataset/webdataset.py

@@ -4,14 +4,14 @@
 from itertools import islice
 from typing import Any, Callable
 
-import fsspec
 import numpy as np
 import pyarrow as pa
 
 import datasets
 from datasets.builder import Key
 from datasets.features.features import cast_to_python_objects
-from datasets.utils.file_utils import SINGLE_FILE_COMPRESSION_EXTENSION_TO_PROTOCOL, xbasename
+from datasets.filesystems import EXTENSION_TO_COMPRESSION_FS_FILE_CLS
+from datasets.utils.file_utils import xbasename
 
 
 logger = datasets.utils.logging.get_logger(__name__)
@@ -29,8 +29,6 @@ class WebDataset(datasets.GeneratorBasedBuilder):
     @classmethod
     def _get_pipeline_from_tar(cls, tar_path, tar_iterator):
         current_example = {}
-        fs: fsspec.AbstractFileSystem = fsspec.filesystem("memory")
-        streaming_download_manager = datasets.StreamingDownloadManager()
         for filename, f in tar_iterator:
             example_key, field_name = base_plus_ext(filename)
             if example_key is None:
@@ -43,15 +41,14 @@ def _get_pipeline_from_tar(cls, tar_path, tar_iterator):
                 current_example = {}
             current_example["__key__"] = example_key
             current_example["__url__"] = tar_path
-            current_example[field_name] = f.read()
-            if field_name.split(".")[-1].lower() in SINGLE_FILE_COMPRESSION_EXTENSION_TO_PROTOCOL:
-                fs.write_bytes(filename, current_example[field_name])
-                extracted_file_path = streaming_download_manager.extract(f"memory://{filename}")
-                with fsspec.open(extracted_file_path) as f:
-                    current_example[field_name] = f.read()
-                fs.delete(filename)
-                data_extension = xbasename(extracted_file_path).split(".")[-1].lower()
+            last_extension = "." + filename.split(".")[-1].lower()
+            if last_extension in EXTENSION_TO_COMPRESSION_FS_FILE_CLS:
+                extracted_filename = ".".join(filename.split(".")[:-1])
+                with EXTENSION_TO_COMPRESSION_FS_FILE_CLS[last_extension](f, mode="rb") as extracted_f:
+                    current_example[field_name] = extracted_f.read()
+                data_extension = xbasename(extracted_filename).split(".")[-1].lower()
             else:
+                current_example[field_name] = f.read()
                 data_extension = field_name.split(".")[-1].lower()
             if data_extension in cls.DECODERS:
                 current_example[field_name] = cls.DECODERS[data_extension](current_example[field_name])

tests/test_fingerprint.py

@@ -365,13 +365,13 @@ def f(x):
     def test_hash_torch_compiled_module(self):
         m = TorchModule()
         next(iter(m.parameters())).data.fill_(1.0)
+        mc = torch.compile(m)
         hash1 = Hasher.hash(m)
-        m = torch.compile(m)
-        hash2 = Hasher.hash(m)
+        hash2 = Hasher.hash(mc)
         m = TorchModule()
         next(iter(m.parameters())).data.fill_(2.0)
-        m = torch.compile(m)
-        hash3 = Hasher.hash(m)
+        mc = torch.compile(m)
+        hash3 = Hasher.hash(mc)
         self.assertEqual(hash1, hash2)
         self.assertNotEqual(hash1, hash3)
         self.assertNotEqual(hash2, hash3)