Merge branch 'main' into intro-run-mgmt

.editorconfig

@@ -16,6 +16,14 @@ profile = black
 indent_style = space
 indent_size = 2
 
+[*.md]
+indent_style = space
+indent_size = 2
+
 [*.yml]
 indent_style = space
 indent_size = 2
+
+[*.toml]
+indent_style = space
+indent_size = 4

benchmarks/flowertune-llm/README.md

@@ -1,4 +1,4 @@
-![](_static/flower_llm.png)
+[![FlowerTune LLM Leaderboard](_static/flower_llm.png)](https://flower.ai/benchmarks/llm-leaderboard)
 
 # FlowerTune LLM Leaderboard
 
@@ -27,15 +27,16 @@ flwr new --framework=FlowerTune
 The `flwr new` command will generate a directory with the following structure:
 
 ```bash
-<project-name>
-       ├── README.md           # <- Instructions
-       ├── pyproject.toml      # <- Environment dependencies and configs
-       └── <project_name>
-                  ├── client_app.py   # <- Flower ClientApp build
-                  ├── dataset.py      # <- Dataset and tokenizer build
-                  ├── models.py       # <- Model build
-                  ├── server_app.py   # <- Flower ServerApp build
-                  └── strategy.py     # <- Flower strategy build
+<project_name>
+├── README.md           # Instructions
+├── pyproject.toml      # Environment dependencies and configs
+└── <project_name>
+    ├── __init__.py
+    ├── client_app.py   # Flower ClientApp build
+    ├── dataset.py      # Dataset and tokenizer build
+    ├── models.py       # Model build
+    ├── server_app.py   # Flower ServerApp build
+    └── strategy.py     # Flower strategy build
 ```
 
 This can serve as the starting point for you to build up your own federated LLM fine-tuning methods.

benchmarks/flowertune-llm/evaluation/README.md

@@ -5,7 +5,7 @@ If you are participating [LLM Leaderboard](https://flower.ai/benchmarks/llm-lead
 
 ## How to run
 
-Navigate to the directory corresponding to your selected challenge (`general NLP`, `finance`, `medical`, or `code`) and follow the instructions there to execute the evaluation.
+Navigate to the directory corresponding to your selected challenge ([`general NLP`](https://github.com/adap/flower/tree/main/benchmarks/flowertune-llm/evaluation/general-nlp), [`finance`](https://github.com/adap/flower/tree/main/benchmarks/flowertune-llm/evaluation/finance), [`medical`](https://github.com/adap/flower/tree/main/benchmarks/flowertune-llm/evaluation/medical), or [`code`](https://github.com/adap/flower/tree/main/benchmarks/flowertune-llm/evaluation/code)) and follow the instructions there to execute the evaluation.
 
 > [!NOTE]  
 > If you wish to participate in the LLM Leaderboard, you must not modify the evaluation code and should use the exact command provided in the respective directory to run the evaluation.

datasets/flwr_datasets/partitioner/__init__.py

@@ -27,6 +27,7 @@
 from .partitioner import Partitioner
 from .pathological_partitioner import PathologicalPartitioner
 from .shard_partitioner import ShardPartitioner
+from .size_partitioner import SizePartitioner
 from .square_partitioner import SquarePartitioner
 
 __all__ = [
@@ -42,5 +43,6 @@
     "Partitioner",
     "PathologicalPartitioner",
     "ShardPartitioner",
+    "SizePartitioner",
     "SquarePartitioner",
 ]

datasets/flwr_datasets/partitioner/size_partitioner.py

@@ -0,0 +1,128 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SizePartitioner class."""
+
+
+import warnings
+from collections.abc import Sequence
+
+import datasets
+from flwr_datasets.partitioner.partitioner import Partitioner
+
+
+class SizePartitioner(Partitioner):
+    """Partitioner that creates each partition with the size specified by a user.
+
+    Parameters
+    ----------
+    partition_sizes : Sequence[int]
+        The size of each partition. partition_id 0 will have partition_sizes[0]
+        samples, partition_id 1 will have partition_sizes[1] samples, etc.
+
+    Examples
+    --------
+    >>> from flwr_datasets import FederatedDataset
+    >>> from flwr_datasets.partitioner import SizePartitioner
+    >>>
+    >>> partition_sizes = [15_000, 5_000, 30_000]
+    >>> partitioner = SizePartitioner(partition_sizes)
+    >>> fds = FederatedDataset(dataset="cifar10", partitioners={"train": partitioner})
+    """
+
+    def __init__(self, partition_sizes: Sequence[int]) -> None:
+        super().__init__()
+        self._pre_ds_validate_partition_sizes(partition_sizes)
+        self._partition_sizes = partition_sizes
+        self._partition_id_to_indices: dict[int, list[int]] = {}
+        self._partition_id_to_indices_determined = False
+
+    def load_partition(self, partition_id: int) -> datasets.Dataset:
+        """Load a single partition of the size of partition_sizes[partition_id].
+
+        For example if given partition_sizes=[20_000, 10_000, 30_000],
+        then partition_id=0 will return a partition of size 20_000,
+        partition_id=1 will return a partition of size 10_000, etc.
+
+        Parameters
+        ----------
+        partition_id : int
+            The index that corresponds to the requested partition.
+
+        Returns
+        -------
+        dataset_partition : Dataset
+            Single dataset partition.
+        """
+        self._determine_partition_id_to_indices_if_needed()
+        return self.dataset.select(self._partition_id_to_indices[partition_id])
+
+    @property
+    def num_partitions(self) -> int:
+        """Total number of partitions."""
+        self._determine_partition_id_to_indices_if_needed()
+        return len(self._partition_sizes)
+
+    @property
+    def partition_id_to_indices(self) -> dict[int, list[int]]:
+        """Partition id to indices (the result of partitioning)."""
+        self._determine_partition_id_to_indices_if_needed()
+        return self._partition_id_to_indices
+
+    def _determine_partition_id_to_indices_if_needed(
+        self,
+    ) -> None:
+        """Create an assignment of indices to the partition indices."""
+        if self._partition_id_to_indices_determined:
+            return
+        self._post_ds_validate_partition_sizes()
+        start = 0
+        end = 0
+        for partition_id, partition_size in enumerate(self._partition_sizes):
+            end += partition_size
+            indices = list(range(start, end))
+            self._partition_id_to_indices[partition_id] = indices
+            start = end
+        self._partition_id_to_indices_determined = True
+
+    def _pre_ds_validate_partition_sizes(self, partition_sizes: Sequence[int]) -> None:
+        """Check if the partition sizes are valid (no information about the dataset)."""
+        if not isinstance(partition_sizes, Sequence):
+            raise ValueError("Partition sizes must be a sequence.")
+        if len(partition_sizes) == 0:
+            raise ValueError("Partition sizes must not be empty.")
+        if not all(
+            isinstance(partition_size, int) for partition_size in partition_sizes
+        ):
+            raise ValueError("All partition sizes must be integers.")
+        if not all(partition_size > 0 for partition_size in partition_sizes):
+            raise ValueError("All partition sizes must be greater than zero.")
+
+    def _post_ds_validate_partition_sizes(self) -> None:
+        """Validate the partition sizes against the dataset size."""
+        desired_partition_sizes = sum(self._partition_sizes)
+        dataset_size = len(self.dataset)
+        if desired_partition_sizes > dataset_size:
+            raise ValueError(
+                f"The sum of partition sizes sum({self._partition_sizes})"
+                f"= {desired_partition_sizes} is greater than the size of"
+                f" the dataset {dataset_size}."
+            )
+        if desired_partition_sizes < dataset_size:
+            warnings.warn(
+                f"The sum of partition sizes is {desired_partition_sizes}, which is"
+                f"smaller than the size of the dataset: {dataset_size}. "
+                f"Ignore this warning if it is the desired behavior.",
+                stacklevel=1,
+            )