adding temporary workaround for getSize issue and fixed passing base_…

…data_id
modin-project · Mar 9, 2023 · 6a63384 · 6a63384
1 parent 02f405d
commit 6a63384
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/unidist/core/backends/mpi/core/common.py b/unidist/core/backends/mpi/core/common.py
@@ -118,8 +118,7 @@ def __del__(self):
     def __getstate__(self):
         """Remove a reference to garbage collector for correct `pickle` serialization."""
         attributes = self.__dict__.copy()
-        if hasattr(self, "_gc"):
-            del attributes["_gc"]
+        del attributes["_gc"]
         return attributes
 
     def base_data_id(self):

diff --git a/unidist/core/backends/mpi/core/controller/api.py b/unidist/core/backends/mpi/core/controller/api.py
@@ -238,7 +238,8 @@ def put(data):
         dest_rank = RoundRobin.get_instance().schedule_rank()
         object_store.put_data_owner(data_id, dest_rank)
         object_store.put_data_size(data_id, data)
-        push_data_directly_to_worker(dest_rank, data_id, data)
+        base_data_id = data_id.base_data_id()
+        push_data_directly_to_worker(dest_rank, base_data_id, data)
 
     logger.debug("PUT {} id".format(data_id._id))
 

diff --git a/unidist/core/backends/mpi/core/controller/object_store.py b/unidist/core/backends/mpi/core/controller/object_store.py
@@ -75,7 +75,15 @@ def put_data_size(self, data_id, data):
         data : object
             Data to be put.
         """
-        self._data_sizes[data_id] = sys.getsizeof(data)
+        # if conditions are a temprary fix as pandas sizeof has an issue
+        # https://github.com/pandas-dev/pandas/issues/51858
+        if "DataFrame" in str(type(data)):
+            size = data.memory_usage().sum()
+        elif "Series" in str(type(data)):
+            size = data.memory_usage()
+        else:
+            size = sys.getsizeof(data)
+        self._data_sizes[data_id] = size
 
     def put_data_owner(self, data_id, rank):
         """