Dataloader improvements

deeplake/core/dataset/indra_dataset_view.py

@@ -60,6 +60,40 @@ def __init__(
         except:
             pass
 
+    def __getstate__(self) -> Dict[str, Any]:
+        keys = [
+            "path",
+            "_read_only",
+            "group_index",
+            "storage",
+            "_token",
+            "verbose",
+            "enabled_tensors",
+            "index"
+        ]
+
+        state = {k: getattr(self, k) for k in keys}
+        return state
+
+    def __setstate__(self, state):
+        from indra import api  # type: ignore
+
+        d: Dict[str, Any] = {}
+        self.storage = state["storage"]
+        d["indra_ds"] = api.load_from_storage(self.storage.core)
+        d["group_index"] = state["group_index"]
+        d["enabled_tensors"] = state["enabled_tensors"]
+        d["verbose"] = state["verbose"]
+        d["_token"] = state["_token"]
+        self.__dict__.update(d)
+        self._view_base = None
+        self._view_entry = None
+        self._read_only = state["_read_only"]
+        self._locked_out = False
+        self._query_string = None
+        index = state["index"]
+        self.indra_ds = self[list(index.values[0].value)].indra_ds
+
     @property
     def meta(self):
         return DatasetMeta()
@@ -97,6 +131,10 @@ def commit_id(self) -> str:
     def libdeeplake_dataset(self):
         return self.indra_ds
 
+    @libdeeplake_dataset.setter
+    def libdeeplake_dataset(self, new_indra_ds):
+        self.indra_ds = new_indra_ds
+
     def merge(self, *args, **kwargs):
         raise InvalidOperationError(
             "merge", "merge method cannot be called on a Dataset view."
@@ -188,23 +226,31 @@ def __getitem__(
                     )
                     for x in item
                 ]
-                return IndraDatasetView(
+                ret = IndraDatasetView(
                     indra_ds=self.indra_ds,
                     enabled_tensors=enabled_tensors,
                 )
+                if hasattr(self, "_tql_query"):
+                    ret._tql_query = self._tql_query
+                return ret
             elif isinstance(item, tuple) and len(item) and isinstance(item[0], str):
                 ret = self
                 for x in item:
                     ret = self[x]
                 return ret
             else:
-                return IndraDatasetView(
+                ret = IndraDatasetView(
                     indra_ds=self.indra_ds[item],
                 )
+                if hasattr(self, "_tql_query"):
+                    ret._tql_query = self._tql_query
+                return ret
         else:
             raise InvalidKeyTypeError(item)
+
         raise AttributeError("Dataset has no attribute - {item}")
 
+
     def __getattr__(self, key):
         try:
             ret = self.__getitem__(key)

deeplake/core/io.py

@@ -293,9 +293,11 @@ def __init__(
         self.tensors = tensors
         self.pad_tensors = pad_tensors
         self.decode_method = decode_method
-        jpeg_png_compressed_tensors, json_tensors, list_tensors = check_tensors(
-            self.dataset, tensors, verbose
-        )
+        (
+            jpeg_png_compressed_tensors,
+            json_tensors,
+            list_tensors,
+        ) = check_tensors(self.dataset, tensors, verbose)
         (
             raw_tensors,
             pil_compressed_tensors,

deeplake/enterprise/convert_to_libdeeplake.py

@@ -8,6 +8,7 @@
 from deeplake.core.storage.azure import AzureProvider
 from deeplake.util.remove_cache import get_base_storage
 from deeplake.util.exceptions import EmptyTokenException
+from deeplake.core.dataset.indra_dataset_view import IndraDatasetView
 
 from deeplake.util.dataset import try_flushing  # type: ignore
 import importlib
@@ -65,6 +66,7 @@ def _get_indra_ds_from_azure_provider(
     storage = IndraProvider(
         path,
         read_only=provider.read_only,
+        origin_path=provider.root,
         token=token,
         account_name=account_name,
         account_key=account_key,
@@ -169,7 +171,7 @@ def dataset_to_libdeeplake(hub2_dataset: Dataset):
     token = (
         hub2_dataset.client.get_token()
         if (hub2_dataset.token is None or hub2_dataset._token == "")
-        and hub2_dataset.client
+        and hasattr(hub2_dataset, "client") and hub2_dataset.client
         else hub2_dataset.token
     )
     if token is None or token == "":
@@ -248,5 +250,11 @@ def dataset_to_libdeeplake(hub2_dataset: Dataset):
     if slice_ != slice(None):
         if isinstance(slice_, tuple):
             slice_ = list(slice_)
-        libdeeplake_dataset = libdeeplake_dataset[slice_]
-    return libdeeplake_dataset
+        from deeplake.core.index import Index
+        try:
+            idx = Index(libdeeplake_dataset.indexes)
+        except:
+            idx = Index(slice(0, len(libdeeplake_dataset)))
+        if isinstance(slice_, slice) or (list(slice_) != list(idx.values[0].value)):
+            libdeeplake_dataset = libdeeplake_dataset[slice_]
+    return libdeeplake_dataset

deeplake/enterprise/dataloader.py

@@ -1,9 +1,10 @@
 from typing import Callable, Dict, List, Optional, Union
 import deeplake
-from deeplake.enterprise.convert_to_libdeeplake import dataset_to_libdeeplake
+
 from deeplake.enterprise.dummy_dataloader import DummyDataloader  # type: ignore
 from deeplake.util.scheduling import create_fetching_schedule, find_primary_tensor
 from deeplake.core.seed import DeeplakeRandom
+from deeplake.util.exceptions import EmptyTensorError, MacOSEnvironmentError
 from deeplake.enterprise.util import (
     handle_mode,
     raise_indra_installation_error,
@@ -22,6 +23,8 @@
 from deeplake.util.dataset import map_tensor_keys
 from functools import partial
 import importlib
+import os
+import sys
 
 try:
     from torch.utils.data.dataloader import DataLoader, _InfiniteConstantSampler
@@ -113,6 +116,7 @@ def __init__(
         _ignore_errors=False,
         _verbose=False,
         _offset=None,
+        _pin_memory=False,
         **kwargs,
     ):
         import_indra_loader()
@@ -139,6 +143,7 @@ def __init__(
         self._ignore_errors = _ignore_errors
         self._verbose = _verbose
         self._offset = _offset
+        self._pin_memory = _pin_memory
         for k, v in kwargs.items():
             setattr(self, k, v)
 
@@ -345,13 +350,16 @@ def shuffle(self, shuffle: bool = True, buffer_size: int = 2048):
         all_vars = self.__dict__.copy()
         all_vars["_shuffle"] = shuffle
         all_vars["_buffer_size"] = buffer_size
-        if shuffle:
-            schedule = create_fetching_schedule(
-                self._orig_dataset, self._primary_tensor_name
-            )
-            if schedule is not None:
-                ds = self._orig_dataset  # type: ignore
-                all_vars["_orig_dataset"] = ds[schedule]
+
+        # TODO check the view dataset shuffle
+        # if shuffle:
+        #     schedule = create_fetching_schedule(
+        #         self._orig_dataset, self._primary_tensor_name
+        #     )
+        #     if schedule is not None:
+        #         ds = self._orig_dataset  # type: ignore
+        #         all_vars["_orig_dataset"] = ds[schedule]
+
         all_vars["_dataloader"] = None
         return self.__class__(**all_vars)
 
@@ -483,6 +491,7 @@ def pytorch(
         return_index: bool = True,
         decode_method: Optional[Dict[str, str]] = None,
         persistent_workers: bool = False,
+        pin_memory: bool = False,
     ):
         """Creates a PyTorch Dataloader on top of the ``DeepLakeDataLoader`` from the Deep Lake dataset. During iteration, the data from all tensors will be streamed on-the-fly from the storage location.
         Understanding the parameters below is critical for achieving fast streaming for your use-case
@@ -498,6 +507,7 @@ def pytorch(
             distributed (bool): Used for DDP training. Distributes different sections of the dataset to different ranks. Defaults to ``False``.
             return_index (bool): Used to idnetify where loader needs to retur sample index or not. Defaults to ``True``.
             persistent_workers (bool): If ``True``, the data loader will not shutdown the worker processes after a dataset has been consumed once. Defaults to ``False``.
+            pin_memory (bool): If ``True``, the data loader will copy Tensors into device/CUDA pinned memory before returning them. Defaults to ``False``.
             decode_method (Dict[str, str], Optional): A dictionary of decode methods for each tensor. Defaults to ``None``.
 
 
@@ -554,6 +564,7 @@ def pytorch(
         all_vars["_mode"] = mode
         all_vars["_persistent_workers"] = persistent_workers
         all_vars["_dataloader"] = None
+        all_vars["_pin_memory"] = pin_memory
         if distributed:
             all_vars["_world_size"] = torch.distributed.get_world_size()
         return self.__class__(**all_vars)
@@ -740,8 +751,7 @@ def __create_dummy_dataloader(
 
     def __get_indra_dataloader(
         self,
-        dataset,
-        indra_dataset,
+        deeplake_dataset,
         tensors: Optional[List[str]] = None,
         raw_tensors: Optional[List[str]] = None,
         pil_compressed_tensors: Optional[List[str]] = None,
@@ -774,25 +784,25 @@ def __get_indra_dataloader(
             json_tensors=json_tensors or [],
             list_tensors=list_tensors or [],
         )
-
         loader_meta = LoaderMetaInfo(
             context=self.multiprocessing_context,
             distributed=self._distributed,
+            mode=self._mode,
             upcast=self._mode == "pytorch"
             and self.__is_upcast_needed(
-                dataset, tensors
+                deeplake_dataset, tensors
             ),  # upcast to handle unsupported dtypes,
             return_index=self._return_index,
             verbose=self._verbose,
             ignore_errors=self._ignore_errors,
             prefetch_factor=self._prefetch_factor,
             offset=self._offset,
-            primary_tensor=self._primary_tensor_name,
             worker_init_fn=self.worker_init_fn,
+            pin_memory=self.pin_memory,
         )
 
         return INDRA_LOADER(  # type: ignore [misc]
-            indra_dataset,
+            deeplake_dataset=deeplake_dataset,
             batch_size=self._batch_size,
             num_threads=num_threads,
             shuffle=self._shuffle,
@@ -833,9 +843,11 @@ def __iter__(self):
             dataset = self._orig_dataset
             tensors = self._tensors or map_tensor_keys(dataset, None)
 
-            jpeg_png_compressed_tensors, json_tensors, list_tensors = check_tensors(
-                dataset, tensors
-            )
+            (
+                jpeg_png_compressed_tensors,
+                json_tensors,
+                list_tensors,
+            ) = check_tensors(dataset, tensors)
             (
                 raw_tensors,
                 pil_compressed_tensors,
@@ -855,6 +867,7 @@ def __iter__(self):
             self._fill_sample_info_tensors(
                 dataset, sample_info_tensors, json_tensors, list_tensors
             )
+
             tensors.extend(sample_info_tensors)
             htype_dict, ndim_dict, tensor_info_dict = get_htype_ndim_tensor_info_dicts(
                 dataset, data_tensors, tensor_info_tensors
@@ -867,14 +880,8 @@ def __iter__(self):
                     pil_compressed_tensors=pil_compressed_tensors,
                 )
             else:
-                if not hasattr(self, "_indra_dataset"):
-                    indra_dataset = dataset_to_libdeeplake(dataset)
-                else:
-                    indra_dataset = self._indra_dataset
-
                 self._dataloader = self.__get_indra_dataloader(
                     dataset,
-                    indra_dataset,
                     tensors=tensors,
                     raw_tensors=raw_tensors,
                     pil_compressed_tensors=pil_compressed_tensors,
@@ -887,11 +894,23 @@ def __iter__(self):
 
         dataset_read(self._orig_dataset)
 
+        self._check_environment()
         if self._iterator is not None:
             self._iterator = iter(self._dataloader)
 
         return self
 
+    def _check_environment(self):
+        if sys.platform == "darwin":
+            import multiprocessing as mp
+
+            if mp.get_start_method() == "fork":
+                env_vars = os.environ
+                no_proxy = env_vars.get("NO_PROXY", "")
+                init_check = env_vars.get("OBJC_DISABLE_INITIALIZE_FORK_SAFETY", "")
+                if no_proxy != "*" or init_check != "YES":
+                    raise MacOSEnvironmentError
+
     def __setattr__(self, attr, val):
         if (
             attr == "_iterator"

deeplake/enterprise/libdeeplake_query.py

@@ -49,7 +49,13 @@ def query(dataset, query_string: str):
     dsv = ds.query(query_string)
     from deeplake.enterprise.convert_to_libdeeplake import INDRA_API
 
-    if not isinstance(dataset, IndraDatasetView) and INDRA_API.tql.parse(query_string).is_filter and len(dsv.indexes) < INDRA_DATASET_SAMPLES_THRESHOLD:  # type: ignore
+    try:
+        is_gt = len(dsv.indexes) < INDRA_DATASET_SAMPLES_THRESHOLD
+    except:
+        is_gt = False
+        pass
+
+    if not isinstance(dataset, IndraDatasetView) and INDRA_API.tql.parse(query_string).is_filter and is_gt:  # type: ignore
         indexes = list(dsv.indexes)
         return dataset.no_view_dataset[indexes]
     else:

deeplake/enterprise/test_pytorch.py

@@ -827,8 +827,6 @@ def test_batch_sampler_attribute(local_auth_ds):
 @pytest.mark.slow
 @pytest.mark.flaky
 def test_pil_decode_method(local_auth_ds):
-    from indra.pytorch.exceptions import CollateExceptionWrapper  # type: ignore
-
     with local_auth_ds as ds:
         ds.create_tensor("x", htype="image", sample_compression="jpeg")
         ds.x.extend(np.random.randint(0, 255, (10, 10, 10, 3), np.uint8))
@@ -840,7 +838,7 @@ def test_pil_decode_method(local_auth_ds):
         assert batch["x"].shape == (1, 10, 10, 3)
 
     ptds = ds.dataloader().pytorch(decode_method={"x": "pil"})
-    with pytest.raises(CollateExceptionWrapper):
+    with pytest.raises(AttributeError):
         for _ in ptds:
             continue
 

deeplake/integrations/pytorch/common.py

@@ -175,7 +175,13 @@ def validate_decode_method(
             pil_compressed_tensors.append(tensor_name)
         elif decode_method == "data":
             data_tensors.append(tensor_name)
-    return raw_tensors, pil_compressed_tensors, json_tensors, list_tensors, data_tensors
+    return (
+        raw_tensors,
+        pil_compressed_tensors,
+        json_tensors,
+        list_tensors,
+        data_tensors,
+    )
 
 
 def find_additional_tensors_and_info(dataset, data_tensors):

deeplake/util/exceptions.py

@@ -847,6 +847,17 @@ def __init__(self, org_id, creds_key):
         )
 
 
+class MacOSEnvironmentError(Exception):
+    def __init__(self):
+        message = (
+            "When using the multiprocessing mode 'fork' on MacOS, "
+            "you need to execute the following commands in the terminal:\n"
+            "export OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES\n"
+            "export NO_PROXY=*"
+        )
+        super().__init__(message)
+
+
 class UnableToReadFromUrlError(Exception):
     def __init__(self, url, status_code):
         super().__init__(f"Unable to read from url {url}. Status code: {status_code}")