Send & Get Datasets (#2960)

* dataset inherits abstractobject

* added abstract dataset test

* formatting

* fix

* removed owner annotation

* formatting

* Added simplify& detail

* Edit simplify& detail

* Added PointerDataset type

* Added send()

* Added send(), create_pointer()

* Added get(), pointer.data(), pointer.target()

* Added wrap(), edited base.search(), edit dataset.__init__()

* Formatting

* Added test file, __repr__(), removed get() override

* Added tags, description to create_pointer()

* formatting

* changed repr of dataset ptr

* replaced == with torch.equal()

* adapt dataloader

* formatting

* added serede test!

* added serede test!

* test fixes

* test fixes

* serde test fix

* formatting

* hook_args fix

* Update test_dataset_pointer.py

* pointer location fix

* formatting

* formatting

* formatting

* formatting

* fedrated check pass

* formatting

Co-authored-by: Muhammed Abogazia <abogaziah@users.noreply.github.com>
Co-authored-by: Théo Ryffel <theo.leffyr@gmail.com>

Author: abogaziah

syft/frameworks/torch/fl/dataset.py

@@ -1,13 +1,16 @@
 import math
 import logging
-
+from syft.generic.object import AbstractObject
+from syft.workers.base import BaseWorker
+from syft.generic.pointers.pointer_dataset import PointerDataset
 import torch
 from torch.utils.data import Dataset
+import syft
 
 logger = logging.getLogger(__name__)
 
 
-class BaseDataset:
+class BaseDataset(AbstractObject):
     """
     This is a base class to be used for manipulating a dataset. This is composed
     of a .data attribute for inputs and a .targets one for labels. It is to
@@ -22,8 +25,10 @@ class BaseDataset:
 
     """
 
-    def __init__(self, data, targets, transform=None):
-
+    def __init__(self, data, targets, transform=None, owner=None, **kwargs):
+        if owner is None:
+            owner = syft.framework.hook.local_worker
+        super().__init__(owner=owner, **kwargs)
         self.data = data
         self.targets = targets
         self.transform_ = transform
@@ -68,21 +73,9 @@ def transform(self, transform):
 
             raise TypeError("Transforms can be applied only on torch tensors")
 
-    def send(self, worker):
-        """
-        Args:
-
-            worker[worker class]: worker to which the data must be sent
-
-        Returns:
-
-            self: Return the object instance with data sent to corresponding worker
-
-        """
-
-        self.data.send_(worker)
-        self.targets.send_(worker)
-        return self
+    def send(self, location: BaseWorker):
+        ptr = self.owner.send(self, workers=location)
+        return ptr
 
     def get(self):
         """
@@ -93,6 +86,12 @@ def get(self):
         self.targets.get_()
         return self
 
+    def get_data(self):
+        return self.data
+
+    def get_targets(self):
+        return self.targets
+
     def fix_prec(self, *args, **kwargs):
         """
             Converts data of BaseDataset into fixed precision
@@ -121,13 +120,81 @@ def share(self, *args, **kwargs):
         self.targets.share_(*args, **kwargs)
         return self
 
+    def create_pointer(
+        self, owner, garbage_collect_data, location=None, id_at_location=None, **kwargs
+    ):
+        """creats a pointer to the self dataset"""
+        if owner is None:
+            owner = self.owner
+
+        if location is None:
+            location = self.owner
+
+        owner = self.owner.get_worker(owner)
+        location = self.owner.get_worker(location)
+
+        return PointerDataset(
+            owner=owner,
+            location=location,
+            id_at_location=id_at_location or self.id,
+            garbage_collect_data=garbage_collect_data,
+            tags=self.tags,
+            description=self.description,
+        )
+
+    def __repr__(self):
+
+        fmt_str = "BaseDataset\n"
+        fmt_str += f"\tData: {self.data}\n"
+        fmt_str += f"\ttargets: {self.targets}"
+
+        if self.tags is not None and len(self.tags):
+            fmt_str += "\n\tTags: "
+            for tag in self.tags:
+                fmt_str += str(tag) + " "
+
+        if self.description is not None:
+            fmt_str += "\n\tDescription: " + str(self.description).split("\n")[0] + "..."
+
+        return fmt_str
+
     @property
     def location(self):
         """
             Get location of the data
         """
         return self.data.location
 
+    @staticmethod
+    def simplify(worker, dataset: "BaseDataset") -> tuple:
+        chain = None
+        if hasattr(dataset, "child"):
+            chain = syft.serde.msgpack.serde._simplify(worker, dataset.child)
+        return (
+            syft.serde.msgpack.serde._simplify(worker, dataset.data),
+            syft.serde.msgpack.serde._simplify(worker, dataset.targets),
+            dataset.id,
+            syft.serde.msgpack.serde._simplify(worker, dataset.tags),
+            syft.serde.msgpack.serde._simplify(worker, dataset.description),
+            chain,
+        )
+
+    @staticmethod
+    def detail(worker, dataset_tuple: tuple) -> "BaseDataset":
+        data, targets, id, tags, description, chain = dataset_tuple
+        dataset = BaseDataset(
+            syft.serde.msgpack.serde._detail(worker, data),
+            syft.serde.msgpack.serde._detail(worker, targets),
+            owner=worker,
+            id=id,
+            tags=syft.serde.msgpack.serde._detail(worker, tags),
+            description=syft.serde.msgpack.serde._detail(worker, description),
+        )
+        if chain is not None:
+            chain = syft.serde.msgpack.serde._detail(worker, chain)
+            dataset.child = chain
+        return dataset
+
 
 def dataset_federate(dataset, workers):
     """
@@ -172,11 +239,11 @@ def __init__(self, datasets):
             self.datasets[worker_id] = dataset
 
         # Check that data and targets for a worker are consistent
-        for worker_id in self.workers:
+        """for worker_id in self.workers:
             dataset = self.datasets[worker_id]
-            assert len(dataset.data) == len(
-                dataset.targets
-            ), "On each worker, the input and target must have the same number of rows."
+            assert (
+                dataset.data.shape == dataset.targets.shape
+            ), "On each worker, the input and target must have the same number of rows.""" ""
 
     @property
     def workers(self):

syft/generic/frameworks/hook/hook_args.py

@@ -658,6 +658,7 @@ def register_response(
 
     try:
         assert attr not in ambiguous_functions
+        assert attr not in ambiguous_methods
 
         # Load the utility function to register the response and transform tensors with pointers
         register_response_function = register_response_functions[attr_id]

syft/generic/pointers/pointer_dataset.py

@@ -0,0 +1,79 @@
+from typing import List
+from typing import Union
+
+import syft as sy
+from syft.generic.pointers.object_pointer import ObjectPointer
+from syft.workers.abstract import AbstractWorker
+
+
+class PointerDataset(ObjectPointer):
+    def __init__(
+        self,
+        location: "AbstractWorker" = None,
+        id_at_location: Union[str, int] = None,
+        owner: "AbstractWorker" = None,
+        garbage_collect_data: bool = True,
+        id: Union[str, int] = None,
+        tags: List[str] = None,
+        description: str = None,
+    ):
+        if owner is None:
+            owner = sy.framework.hook.local_worker
+        super().__init__(
+            location=location,
+            id_at_location=id_at_location,
+            owner=owner,
+            garbage_collect_data=garbage_collect_data,
+            id=id,
+            tags=tags,
+            description=description,
+        )
+
+    @property
+    def data(self):
+        command = ("get_data", self.id_at_location, [], {})
+        ptr = self.owner.send_command(message=command, recipient=self.location).wrap()
+        return ptr
+
+    @property
+    def targets(self):
+        command = ("get_targets", self.id_at_location, [], {})
+        ptr = self.owner.send_command(message=command, recipient=self.location).wrap()
+        return ptr
+
+    def wrap(self):
+        return self
+
+    def __repr__(self):
+        type_name = type(self).__name__
+        out = f"[" f"{type_name} | " f"owner: {str(self.owner.id)}, id:{self.id}"
+
+        if self.point_to_attr is not None:
+            out += "::" + str(self.point_to_attr).replace(".", "::")
+
+        big_str = False
+
+        if self.tags is not None and len(self.tags):
+            big_str = True
+            out += "\n\tTags: "
+            for tag in self.tags:
+                out += str(tag) + " "
+
+        if big_str and hasattr(self, "shape"):
+            out += "\n\tShape: " + str(self.shape)
+
+        if self.description is not None:
+            big_str = True
+            out += "\n\tDescription: " + str(self.description).split("\n")[0] + "..."
+
+        return out
+
+    def __len__(self):
+        command = ("__len__", self.id_at_location, [], {})
+        len = self.owner.send_command(message=command, recipient=self.location)
+        return len
+
+    def __getitem__(self, index):
+        command = ("__getitem__", self.id_at_location, [index], {})
+        data_elem, target_elem = self.owner.send_command(message=command, recipient=self.location)
+        return data_elem.wrap(), target_elem.wrap()

syft/serde/msgpack/serde.py

@@ -72,6 +72,7 @@
 from syft.serde.msgpack.native_serde import MAP_NATIVE_SIMPLIFIERS_AND_DETAILERS
 from syft.workers.abstract import AbstractWorker
 from syft.workers.base import BaseWorker
+from syft.frameworks.torch.fl import BaseDataset
 
 from syft.exceptions import GetNotPermittedError
 from syft.exceptions import ResponseSignatureError
@@ -133,6 +134,7 @@
     PlanCommandMessage,
     GradFunc,
     String,
+    BaseDataset,
     ExecuteWorkerFunctionMessage,
 ]
 

test/serde/msgpack/test_msgpack_serde_full.py

@@ -71,6 +71,7 @@
 samples[syft.frameworks.torch.tensors.interpreters.autograd.AutogradTensor] = make_autogradtensor
 samples[syft.frameworks.torch.tensors.interpreters.private.PrivateTensor] = make_privatetensor
 samples[syft.frameworks.torch.tensors.interpreters.placeholder.PlaceHolder] = make_placeholder
+samples[syft.frameworks.torch.fl.dataset.BaseDataset] = make_basedataset
 
 samples[syft.messaging.message.CommandMessage] = make_command_message
 samples[syft.messaging.message.ObjectMessage] = make_objectmessage

test/serde/serde_helpers.py

@@ -668,6 +668,45 @@ def compare(detailed, original):
     ]
 
 
+# syft.frameworks.torch.fl.dataset
+def make_basedataset(**kwargs):
+    workers = kwargs["workers"]
+    alice, bob, james = workers["alice"], workers["bob"], workers["james"]
+    dataset = syft.BaseDataset(torch.tensor([1, 2, 3, 4]), torch.tensor([5, 6, 7, 8]))
+    dataset.tag("#tag1").describe("desc")
+
+    def compare(detailed, original):
+        assert type(detailed) == syft.BaseDataset
+        assert (detailed.data == original.data).all()
+        assert (detailed.targets == original.targets).all()
+        assert detailed.id == original.id
+        assert detailed.tags == original.tags
+        assert detailed.description == original.description
+        return True
+
+    return [
+        {
+            "value": dataset,
+            "simplified": (
+                CODE[syft.frameworks.torch.fl.dataset.BaseDataset],
+                (
+                    msgpack.serde._simplify(syft.hook.local_worker, dataset.data),
+                    msgpack.serde._simplify(syft.hook.local_worker, dataset.targets),
+                    dataset.id,
+                    msgpack.serde._simplify(
+                        syft.hook.local_worker, dataset.tags
+                    ),  # (set of str) tags
+                    msgpack.serde._simplify(
+                        syft.hook.local_worker, dataset.description
+                    ),  # (str) description
+                    msgpack.serde._simplify(syft.hook.local_worker, dataset.child),
+                ),
+            ),
+            "cmp_detailed": compare,
+        }
+    ]
+
+
 # syft.execution.plan.Plan
 def make_plan(**kwargs):
     # Function to plan

test/torch/federated/test_dataset.py

@@ -15,17 +15,11 @@ def test_base_dataset(workers):
     assert len(dataset) == 4
     assert dataset[2] == (3, 3)
 
-    dataset.send(bob)
+    dataset = dataset.send(bob)
     assert dataset.data.location.id == "bob"
     assert dataset.targets.location.id == "bob"
     assert dataset.location.id == "bob"
 
-    dataset.get()
-    with pytest.raises(AttributeError):
-        assert dataset.data.location.id == 0
-    with pytest.raises(AttributeError):
-        assert dataset.targets.location.id == 0
-
 
 def test_base_dataset_transform():
 
@@ -61,11 +55,12 @@ def test_federated_dataset(workers):
     assert fed_dataset.workers == ["bob", "alice"]
     assert len(fed_dataset) == 6
 
-    fed_dataset["alice"].get()
-    assert (fed_dataset["alice"].data == alice_base_dataset.data).all()
-    assert fed_dataset["alice"][2] == (5, 5)
-    assert len(fed_dataset["alice"]) == 4
-    assert len(fed_dataset) == 6
+    alice_remote_data = fed_dataset["alice"].get()
+    del fed_dataset.datasets["alice"]
+    assert (alice_remote_data.data == alice_base_dataset.data).all()
+    assert alice_remote_data[2] == (5, 5)
+    assert len(alice_remote_data) == 4
+    assert len(fed_dataset) == 2
 
     assert isinstance(fed_dataset.__str__(), str)
 
@@ -114,3 +109,12 @@ def test_federated_dataset_search(workers):
         counter += 1
 
     assert counter == len(train_loader), f"{counter} == {len(fed_dataset)}"
+
+
+def test_abstract_dataset():
+    inputs = th.tensor([1, 2, 3, 4.0])
+    targets = th.tensor([1, 2, 3, 4.0])
+    dataset = BaseDataset(inputs, targets, id=1)
+
+    assert dataset.id == 1
+    assert dataset.description == None