Update embedded-devices example (#2384)

adap · Oct 6, 2023 · fa0fd8a · fa0fd8a
1 parent 7e14c96
commit fa0fd8a
Show file tree

Hide file tree

Showing 22 changed files with 587 additions and 730 deletions.
diff --git a/doc/source/ref-changelog.md b/doc/source/ref-changelog.md
@@ -16,6 +16,8 @@
 
   - FedProx ([#2210](https://github.com/adap/flower/pull/2210), [#2286](https://github.com/adap/flower/pull/2286))
 
+- **Update Flower Examples** ([#2384](https://github.com/adap/flower/pull/2384))
+
 - **General updates to baselines** ([#2301](https://github.com/adap/flower/pull/2301).[#2305](https://github.com/adap/flower/pull/2305), [#2307](https://github.com/adap/flower/pull/2307), [#2327](https://github.com/adap/flower/pull/2327))
 
 - **General updates to the simulation engine** ([#2331](https://github.com/adap/flower/pull/2331))

diff --git a/examples/embedded-devices/Dockerfile b/examples/embedded-devices/Dockerfile
@@ -1,28 +1,13 @@
-ARG BASE_IMAGE_TYPE=cpu
-# these images have been pushed to Dockerhub but you can find
-# each Dockerfile used in the `base_images` directory 
-FROM jafermarq/jetsonfederated_$BASE_IMAGE_TYPE:latest
+ARG BASE_IMAGE
 
-RUN apt-get install wget -y
+# Pull the base image from NVIDIA
+FROM $BASE_IMAGE
 
-# Download and extract CIFAR-10
-# To keep things simple, we keep this as part of the docker image.
-# If the dataset is already in your system you can mount it instead.
-ENV DATA_DIR=/app/data/cifar-10
-RUN mkdir -p $DATA_DIR
-WORKDIR $DATA_DIR
-RUN wget https://www.cs.toronto.edu/\~kriz/cifar-10-python.tar.gz 
-RUN tar -zxvf cifar-10-python.tar.gz
-
-WORKDIR /app
-# Scripts needed for Flower client
-ADD client.py /app
-ADD utils.py /app
-
-# update pip
+# Update pip
 RUN pip3 install --upgrade pip
 
-# making sure the latest version of flower is installed
-RUN pip3 install flwr>=1.0.0
+# Install flower
+RUN pip3 install flwr>=1.0
+RUN pip3 install tqdm==4.65.0
 
-ENTRYPOINT ["python3","-u","./client.py"]
+WORKDIR /client
diff --git a/examples/embedded-devices/README.md b/examples/embedded-devices/README.md
diff --git a/examples/embedded-devices/_static/diagram.png b/examples/embedded-devices/_static/diagram.png
diff --git a/examples/embedded-devices/_static/rpi_imager.png b/examples/embedded-devices/_static/rpi_imager.png
diff --git a/examples/embedded-devices/base_images/README.md b/examples/embedded-devices/base_images/README.md
diff --git a/examples/embedded-devices/base_images/cpu/Dockerfile b/examples/embedded-devices/base_images/cpu/Dockerfile
diff --git a/examples/embedded-devices/base_images/cpu/build.sh b/examples/embedded-devices/base_images/cpu/build.sh
diff --git a/examples/embedded-devices/base_images/gpu/Dockerfile b/examples/embedded-devices/base_images/gpu/Dockerfile
diff --git a/examples/embedded-devices/base_images/gpu/build.sh b/examples/embedded-devices/base_images/gpu/build.sh
diff --git a/examples/embedded-devices/build_image.sh b/examples/embedded-devices/build_image.sh
diff --git a/examples/embedded-devices/build_jetson_flower_client.sh b/examples/embedded-devices/build_jetson_flower_client.sh
@@ -0,0 +1,42 @@
+#!/bin/bash
+
+if [ -z "${CI}" ]; then
+    BUILDKIT=1
+else
+    BUILDKIT=0
+fi
+
+# This script build a docker image that's ready to run your flower client.
+# Depending on your choice of ML framework (TF or PyTorch), the appropiate
+# base image from NVIDIA will be pulled. This ensures you get the best
+# performance out of your Jetson device.
+
+BASE_PYTORCH=nvcr.io/nvidia/l4t-pytorch:r35.1.0-pth1.13-py3
+BASE_TF=nvcr.io/nvidia/l4t-tensorflow:r35.3.1-tf2.11-py3
+EXTRA=""
+
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    -p|--pytorch)
+      BASE_IMAGE=$BASE_PYTORCH
+      shift
+      ;;
+    -t|--tensorflow)
+      BASE_IMAGE=$BASE_TF
+      shift
+      ;;
+    -r|--no-cache)
+      EXTRA="--no-cache"
+      shift
+      ;;
+    -*|--*)
+      echo "Unknown option $1 (pass either --pytorch or --tensorflow)"
+      exit 1
+      ;;
+  esac
+done
+
+DOCKER_BUILDKIT=${BUILDKIT} docker build $EXTRA \
+                                        --build-arg BASE_IMAGE=$BASE_IMAGE \
+                                        . \
+                                        -t flower_client:latest
diff --git a/examples/embedded-devices/client.py b/examples/embedded-devices/client.py