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,21 +293,26 @@ 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,
+            medical_tensors,
+        ) = check_tensors(self.dataset, tensors, verbose)
         (
             raw_tensors,
             pil_compressed_tensors,
             json_tensors,
             list_tensors,
             data_tensors,
+            medical_tensors,
         ) = validate_decode_method(
             self.decode_method,
             tensors,
             jpeg_png_compressed_tensors,
             json_tensors,
             list_tensors,
+            medical_tensors,
         )
         sample_info_tensors, tensor_info_tensors = find_additional_tensors_and_info(
             dataset, data_tensors

deeplake/core/storage/gcs.py

@@ -528,5 +528,7 @@ def get_object_from_full_url(self, url: str):
 
     def get_creds(self):
         d = self.scoped_credentials.get_token_info()
-        d["expiration"] = self.expiration or ""
+        d["expiration"] = (
+            self.expiration if hasattr(self, "expiration") and self.expiration else ""
+        )
         return d

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,
@@ -168,7 +170,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 == "":
@@ -247,5 +249,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
@@ -112,6 +115,7 @@ def __init__(
         _ignore_errors=False,
         _verbose=False,
         _offset=None,
+        _pin_memory=False,
         **kwargs,
     ):
         import_indra_loader()
@@ -137,6 +141,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)
 
@@ -343,11 +348,6 @@ 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.dataset, self._primary_tensor_name)
-            if schedule is not None:
-                ds = self.dataset.no_view_dataset  # type: ignore
-                all_vars["dataset"] = ds[schedule]
         all_vars["_dataloader"] = None
         return self.__class__(**all_vars)
 
@@ -479,6 +479,7 @@ def pytorch(
         return_index: bool = True,
         decode_method: Optional[Dict[str, str]] = None,
         persistent_workers: bool = False,
+        pin_memory: bool = False,
     ):
         """Returns a :class:`DeepLakeDataLoader` object.
 
@@ -492,6 +493,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``.
 
 
@@ -548,6 +550,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)
@@ -734,13 +737,13 @@ 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,
         json_tensors: Optional[List[str]] = None,
         list_tensors: Optional[List[str]] = None,
+        medical_tensors: Optional[List[str]] = None,
         htype_dict: Optional[dict] = None,
         ndim_dict: Optional[dict] = None,
         tensor_info_dict: Optional[dict] = None,
@@ -767,26 +770,27 @@ def __get_indra_dataloader(
             pil_compressed_tensors=pil_compressed_tensors or [],
             json_tensors=json_tensors or [],
             list_tensors=list_tensors or [],
+            medical_tensors=medical_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,
@@ -801,30 +805,62 @@ def __get_indra_dataloader(
             info=info,
         )
 
+    def _fill_sample_info_tensors(
+        self,
+        dataset,
+        sample_info_tensors,
+        json_tensors,
+        list_tensors,
+    ):
+        for tensor_name in sample_info_tensors:
+            tensor = dataset._get_tensor_from_root(tensor_name)
+            if len(tensor) == 0:
+                raise EmptyTensorError(
+                    f" the dataset has an empty tensor {tensor_name}, pytorch dataloader can't be created."
+                    f" Please either populate the tensor or pass tensors argument to .pytorch that excludes this"
+                    f" tensor."
+                )
+            meta = tensor.meta
+            if meta.htype == "json":
+                json_tensors.append(tensor_name)
+            elif meta.htype == "list":
+                list_tensors.append(tensor_name)
+            elif meta.htype == "tag":
+                list_tensors.append(tensor_name)
+
     def __iter__(self):
         if self._dataloader is None:
             dataset = self.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,
+                medical_tensors,
+            ) = check_tensors(dataset, tensors)
             (
                 raw_tensors,
                 pil_compressed_tensors,
                 json_tensors,
                 list_tensors,
                 data_tensors,
+                medical_tensors,
             ) = validate_decode_method(
                 self._decode_method,
                 tensors,
                 jpeg_png_compressed_tensors,
                 json_tensors,
                 list_tensors,
+                medical_tensors,
             )
             sample_info_tensors, tensor_info_tensors = find_additional_tensors_and_info(
                 dataset, data_tensors
             )
+            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
@@ -837,31 +873,38 @@ 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,
                     json_tensors=json_tensors,
                     list_tensors=list_tensors,
+                    medical_tensors=medical_tensors,
                     htype_dict=htype_dict,
                     ndim_dict=ndim_dict,
                     tensor_info_dict=tensor_info_dict,
                 )
 
         dataset_read(self.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"