Forward fix / Update dill_available API for torchdata

Summary: Changes from the PyTorch repo (D53082622) broke torchdata. I updated the dill_available API for torchdata to keep everything in sync.

Differential Revision: D53086369

test/test_serialization.py

@@ -17,11 +17,11 @@
 import torchdata.datapipes.iter as iterdp
 import torchdata.datapipes.map as mapdp
 from _utils._common_utils_for_test import create_temp_dir, create_temp_files
-from torch.utils.data.datapipes.utils.common import DILL_AVAILABLE
+from torch.utils._import_utils import dill_available
 from torchdata.datapipes.iter import IterableWrapper
 from torchdata.datapipes.map import SequenceWrapper
 
-if DILL_AVAILABLE:
+if dill_available():
     import dill
 
     dill.extend(use_dill=False)
@@ -87,7 +87,7 @@ def _filter_by_module_availability(datapipes):
         filter_set.update([iterdp.IoPathFileLister, iterdp.IoPathFileOpener, iterdp.IoPathSaver])
     if rarfile is None:
         filter_set.update([iterdp.RarArchiveLoader])
-    if torcharrow is None or not DILL_AVAILABLE:
+    if torcharrow is None or not dill_available():
         filter_set.update([iterdp.DataFrameMaker, iterdp.ParquetDataFrameLoader])
     return [dp for dp in datapipes if dp[0] not in filter_set]
 
@@ -374,7 +374,7 @@ def test_serializable_with_dill(self) -> None:
         # Skipping value comparison for these DataPipes
         dp_skip_comparison = {iterdp.OnDiskCacheHolder, iterdp.ParagraphAggregator}
         for dpipe, dp_args, dp_kwargs in unpicklable_datapipes:
-            if DILL_AVAILABLE:
+            if dill_available():
                 try:
                     if dpipe in dp_skip_comparison:  # Make sure they are picklable/loadable (no value comparison)
                         datapipe = dpipe(input_dp, *dp_args, **dp_kwargs)  # type: ignore[call-arg]

torchdata/datapipes/iter/util/cacheholder.py

@@ -21,13 +21,14 @@
 except ImportError:
     portalocker = None
 
-from torch.utils.data.datapipes.utils.common import _check_unpickable_fn, DILL_AVAILABLE
+from torch.utils.data.datapipes.utils.common import _check_unpickable_fn
+from torch.utils._import_utils import dill_available
 
 from torch.utils.data.graph import traverse_dps
 from torchdata.datapipes import functional_datapipe
 from torchdata.datapipes.iter import IterableWrapper, IterDataPipe
 
-if DILL_AVAILABLE:
+if dill_available():
     import dill
 
     dill.extend(use_dill=False)

torchdata/datapipes/iter/util/converter.py

@@ -9,9 +9,10 @@
 from typing import Callable, Dict, Optional
 
 from torch.utils.data import IterDataPipe, MapDataPipe
-from torch.utils.data.datapipes.utils.common import _check_unpickable_fn, DILL_AVAILABLE
+from torch.utils.data.datapipes.utils.common import _check_unpickable_fn
+from torch.utils._import_utils import dill_available
 
-if DILL_AVAILABLE:
+if dill_available():
     import dill
 
     dill.extend(use_dill=False)
@@ -108,7 +109,7 @@ def __len__(self):
         return len(self._map)  # type: ignore[arg-type]
 
     def __getstate__(self):
-        if DILL_AVAILABLE:
+        if dill_available():
             dill_key_value_fn = dill.dumps(self.key_value_fn)
         else:
             dill_key_value_fn = self.key_value_fn
@@ -120,7 +121,7 @@ def __getstate__(self):
 
     def __setstate__(self, state):
         (self.datapipe, dill_key_value_fn, self._map) = state
-        if DILL_AVAILABLE:
+        if dill_available():
             self.key_value_fn = dill.loads(dill_key_value_fn)  # type: ignore[assignment]
         else:
             self.key_value_fn = dill_key_value_fn  # type: ignore[assignment]

torchdata/datapipes/iter/util/dataframemaker.py

@@ -7,7 +7,7 @@
 from functools import partial
 from typing import List, Optional, TypeVar
 
-from torch.utils.data.datapipes.utils.common import DILL_AVAILABLE
+from torch.utils._import_utils import dill_available
 
 from torchdata.datapipes import functional_datapipe
 from torchdata.datapipes.iter import IterDataPipe
@@ -19,7 +19,7 @@
     torcharrow = None
     parquet = None
 
-if DILL_AVAILABLE:
+if dill_available():
     import dill
 
     dill.extend(use_dill=False)
@@ -150,7 +150,7 @@ def __iter__(self):
                 yield torcharrow.from_arrow(row_group, dtype=self.dtype)
 
     def __getstate__(self):
-        if DILL_AVAILABLE:
+        if dill_available():
             dill_dtype = dill.dumps(self.dtype)
         else:
             dill_dtype = self.dtype
@@ -161,7 +161,7 @@ def __getstate__(self):
 
     def __setstate__(self, state):
         (self.source_dp, dill_dtype, self.columns, self.device, self.use_threads) = state
-        if DILL_AVAILABLE:
+        if dill_available():
             self.dtype = dill.loads(dill_dtype)  # type: ignore[assignment]
         else:
             self.dtype = dill_dtype  # type: ignore[assignment]