From 370f90ffd1416e28956894753d058db53031f66a Mon Sep 17 00:00:00 2001
From: jafermarq <javier@flower.dev>
Date: Wed, 17 Jan 2024 22:58:11 +0000
Subject: [PATCH] updates

---
 examples/advanced-tensorflow/README.md        | 7 ++++---
 examples/advanced-tensorflow/client.py        | 8 +++-----
 examples/advanced-tensorflow/requirements.txt | 1 +
 examples/advanced-tensorflow/run.sh           | 4 ++--
 4 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/examples/advanced-tensorflow/README.md b/examples/advanced-tensorflow/README.md
index 0a8fc0d7884c..b21c0d2545ca 100644
--- a/examples/advanced-tensorflow/README.md
+++ b/examples/advanced-tensorflow/README.md
@@ -1,6 +1,6 @@
 # Advanced Flower Example (TensorFlow/Keras)
 
-This example demonstrates an advanced federated learning setup using Flower with TensorFlow/Keras. It differs from the quickstart example in the following ways:
+This example demonstrates an advanced federated learning setup using Flower with TensorFlow/Keras. This example uses [Flower Datasets](https://flower.dev/docs/datasets/) and it differs from the quickstart example in the following ways:
 
 - 10 clients (instead of just 2)
 - Each client holds a local dataset of 1/10 of the train datasets and 80% is training examples and 20% as test examples (note that by default only a small subset of this data is used when running the `run.sh` script)
@@ -57,10 +57,11 @@ pip install -r requirements.txt
 
 ## Run Federated Learning with TensorFlow/Keras and Flower
 
-The included `run.sh` will call a script to generate certificates (which will be used by server and clients), start the Flower server (using `server.py`), sleep for 2 seconds to ensure the the server is up, and then start 10 Flower clients (using `client.py`). You can simply start everything in a terminal as follows:
+The included `run.sh` will call a script to generate certificates (which will be used by server and clients), start the Flower server (using `server.py`), sleep for 10 seconds to ensure the the server is up, and then start 10 Flower clients (using `client.py`). You can simply start everything in a terminal as follows:
 
 ```shell
-poetry run ./run.sh
+# Once you have activated your environment
+./run.sh
 ```
 
 The `run.sh` script starts processes in the background so that you don't have to open eleven terminal windows. If you experiment with the code example and something goes wrong, simply using `CTRL + C` on Linux (or `CMD + C` on macOS) wouldn't normally kill all these processes, which is why the script ends with `trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT` and `wait`. This simply allows you to stop the experiment using `CTRL + C` (or `CMD + C`). If you change the script and anything goes wrong you can still use `killall python` (or `killall python3`) to kill all background processes (or a more specific command if you have other Python processes running that you don't want to kill).
diff --git a/examples/advanced-tensorflow/client.py b/examples/advanced-tensorflow/client.py
index 2185727e7258..033f20b1b027 100644
--- a/examples/advanced-tensorflow/client.py
+++ b/examples/advanced-tensorflow/client.py
@@ -76,7 +76,7 @@ def main() -> None:
     # Parse command line argument `partition`
     parser = argparse.ArgumentParser(description="Flower")
     parser.add_argument(
-        "--partition",
+        "--client-id",
         type=int,
         default=0,
         choices=range(0, 10),
@@ -86,9 +86,7 @@ def main() -> None:
     )
     parser.add_argument(
         "--toy",
-        type=bool,
-        default=False,
-        required=False,
+        action='store_true',
         help="Set to true to quicky run the client using only 10 datasamples. "
         "Useful for testing purposes. Default: False",
     )
@@ -101,7 +99,7 @@ def main() -> None:
     model.compile("adam", "sparse_categorical_crossentropy", metrics=["accuracy"])
 
     # Load a subset of CIFAR-10 to simulate the local data partition
-    x_train, y_train, x_test, y_test = load_partition(args.partition)
+    x_train, y_train, x_test, y_test = load_partition(args.client_id)
 
     if args.toy:
         x_train, y_train = x_train[:10], y_train[:10]
diff --git a/examples/advanced-tensorflow/requirements.txt b/examples/advanced-tensorflow/requirements.txt
index 7a70c46a8128..0cb5fe8c07af 100644
--- a/examples/advanced-tensorflow/requirements.txt
+++ b/examples/advanced-tensorflow/requirements.txt
@@ -1,3 +1,4 @@
 flwr>=1.0, <2.0
+flwr-datasets = { extras = ["vision"], version = ">=0.0.2,<1.0.0" }
 tensorflow-cpu>=2.9.1, != 2.11.1 ; platform_machine == "x86_64"
 tensorflow-macos>=2.9.1, != 2.11.1 ; sys_platform == "darwin" and platform_machine == "arm64"
diff --git a/examples/advanced-tensorflow/run.sh b/examples/advanced-tensorflow/run.sh
index a19142fec5b4..9537bbe7cbb0 100755
--- a/examples/advanced-tensorflow/run.sh
+++ b/examples/advanced-tensorflow/run.sh
@@ -5,11 +5,11 @@
 echo "Starting server"
 
 python server.py &
-sleep 3  # Sleep for 3s to give the server enough time to start
+sleep 10  # Sleep for 10s to give the server enough time to start
 
 for i in $(seq 0 9); do
     echo "Starting client $i"
-    python client.py --partition=${i} --toy True &
+    python client.py --client-id=${i} --toy &
 done
 
 # This will allow you to use CTRL+C to stop all background processes