Trigger the partitioning in the num_partitions

adap · Mar 11, 2024 · 02411f7 · 02411f7
1 parent 4211941
commit 02411f7
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 0 deletions.
diff --git a/datasets/flwr_datasets/partitioner/dirichlet_partitioner.py b/datasets/flwr_datasets/partitioner/dirichlet_partitioner.py
@@ -135,6 +135,8 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
     @property
     def num_partitions(self) -> int:
         """Total number of partitions."""
+        self._check_num_partitions_correctness_if_needed()
+        self._determine_node_id_to_indices_if_needed()
         return self._num_partitions
 
     def _initialize_alpha(

diff --git a/datasets/flwr_datasets/partitioner/inner_dirichlet_partitioner.py b/datasets/flwr_datasets/partitioner/inner_dirichlet_partitioner.py
@@ -122,6 +122,12 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
     @property
     def num_partitions(self) -> int:
         """Total number of partitions."""
+        self._check_num_partitions_correctness_if_needed()
+        self._check_partition_sizes_correctness_if_needed()
+        self._check_the_sum_of_partition_sizes()
+        self._determine_num_unique_classes_if_needed()
+        self._alpha = self._initialize_alpha_if_needed(self._initial_alpha)
+        self._determine_node_id_to_indices_if_needed()
         return self._num_partitions
 
     def _initialize_alpha_if_needed(

diff --git a/datasets/flwr_datasets/partitioner/natural_id_partitioner.py b/datasets/flwr_datasets/partitioner/natural_id_partitioner.py
@@ -68,6 +68,8 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
     @property
     def num_partitions(self) -> int:
         """Total number of partitions."""
+        if len(self._node_id_to_natural_id) == 0:
+            self._create_int_node_id_to_natural_id()
         return len(self._node_id_to_natural_id)
 
     @property

diff --git a/datasets/flwr_datasets/partitioner/shard_partitioner.py b/datasets/flwr_datasets/partitioner/shard_partitioner.py
@@ -182,6 +182,10 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
     @property
     def num_partitions(self) -> int:
         """Total number of partitions."""
+        self._check_num_partitions_correctness_if_needed()
+        self._check_possibility_of_partitions_creation()
+        self._sort_dataset_if_needed()
+        self._determine_node_id_to_indices_if_needed()
         return self._num_partitions
 
     def _determine_node_id_to_indices_if_needed(self) -> None:  # pylint: disable=R0914

diff --git a/datasets/flwr_datasets/partitioner/size_partitioner.py b/datasets/flwr_datasets/partitioner/size_partitioner.py
@@ -87,6 +87,7 @@ def load_partition(self, node_id: int) -> datasets.Dataset:
     @property
     def num_partitions(self) -> int:
         """Total number of partitions."""
+        self._determine_node_id_to_indices_if_needed()
         return self._num_partitions
 
     @property